Support fixed features in Service API

[AngelFP]

Addresses https://github.com/facebook/Ax/issues/746 (also in the wishlist https://github.com/facebook/Ax/issues/566).

As the title implies, this PR adds the possibility of specifying some ObservationFeatures as fixed_features in AxClient.get_next_trial and AxClient.get_next_trials. From my understanding, this is currently only possible with the developer API.

This is an option we would like to expose in optimas, where we use the Service API for most of the BO generators.

[facebook-github-bot]

@Cesar-Cardoso has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[codecov-commenter]

Codecov Report

Attention: 1 lines in your changes are missing coverage. Please review.

Comparison is base (fd80cba) 94.52% compared to head (44679e3) 94.51%. Report is 17 commits behind head on main.

:exclamation: Current head 44679e3 differs from pull request most recent head 1677261. Consider uploading reports for the commit 1677261 to get more accurate results

Files	Patch %	Lines
ax/service/ax_client.py	75.00%	1 Missing :warning:

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2015      +/-   ##
==========================================
- Coverage   94.52%   94.51%   -0.01%     
==========================================
  Files         460      460              
  Lines       44301    44303       +2     
==========================================
+ Hits        41874    41875       +1     
- Misses       2427     2428       +1     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

[AngelFP]

Hi all,

has there been any update on this? We're currently having to use a custom subclass of AxClient in order to have this feature, so it would be great if it could be merged :D (it would also address https://github.com/facebook/Ax/issues/1951).

Feel free to let me know if there's anything I should change in how I implemented the feature.

[sgbaird]

Seconded on @AngelFP's comment! I've been itching to implement a fixed feature option within https://honegumi.readthedocs.io/ to better serve the researchers in my circle (earlier today I had a conversation about fixed features with a colleague today), but it's going to be much more straightforward to do so with this PR merged. Anything that requires help? Does this perhaps require a unit test to be written?

EDIT: If I'm not mistaken, it's already covered by a unit test, correct?

https://github.com/AngelFP/Ax/blob/feature/fixed_features_service/ax/service/tests/test_ax_client.py#L2868-L2884

@Cesar-Cardoso and co., it looks like all checks are passing, except for "Facebook Internal - Linter". Does this check need to be fixed? Non- meta employees can't view details on this.

I also wanted to point out that there are many issues scattered around related to this:

• https://github.com/facebook/Ax/issues/1951
• https://github.com/facebook/Ax/issues/746
• https://github.com/facebook/Ax/issues/905
• https://github.com/facebook/Ax/issues/772
• https://github.com/facebook/Ax/issues/1931

[sgbaird]

As an aside, now actually having reviewed the changed files, @AngelFP's implementation looks really straightforward. The key lines are:

def _gen_new_generator_run(
        self, n: int = 1, fixed_features: Optional[ObservationFeatures] = None
    ) -> GeneratorRun:
	...
    if fixed_features:
        fixed_feats.update_features(fixed_features)
   ...

update_features(...) is the the lowest level piece that's changing. The rest is just plumbing.

[Cesar-Cardoso]

Hi folks! Apologies for the slow response from my part. Support for fixed features is being added now in https://github.com/facebook/Ax/pull/2372. Only difference being the use of FixedFeatures instead of ObservationFeatures, as ObservationFeatures isn't meant to be used as input for the API. We also had some internal discussion about using TParameterization directly instead of FixedFeatures, but decided against it.

[AngelFP]

Hi @Cesar-Cardoso , great to see this done! Thanks for looking into it.

[sgbaird]

I'll try to put a MWE up soon, but a starting point might be https://github.com/facebook/Ax/blob/18d08424c4a11ca8ed9a50b760c87826e886be22/ax/service/tests/test_ax_client.py#L2838-L2863