Support fixed features in Service API

[AngelFP]

Addresses #746 (also in the wishlist #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.

❗ 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 ⚠️

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     

☔ View full report in Codecov by Sentry.
📢 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 #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:

• [Feature Request] Support fixing parameter values in get_next_trial() in ServiceAPI  #1951
• [Feature Request] Specifying fixed features to AxClient.get_next_trial (to conduct contextual BO using Service API) #746
• How to include extended features (e.g. state variables like temperature or pressure) in a restricted search space: "Fixed parameter ... with out of design value" #905
• Incorporate domain knowledge in categorial parameters #772
• [Question] Changing parameter constraints before generating a new trial #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.

[sgbaird]

@Cesar-Cardoso This is really the only thing that's changing. The rest is plumbing to be able to pass fixed_features down to _gen_new_generator_run(...).

[Cesar-Cardoso]

Hi folks! Apologies for the slow response from my part. Support for fixed features is being added now in #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

Ax/ax/service/tests/test_ax_client.py

Lines 2838 to 2863 in 18d0842

	def test_gen_fixed_features(self) -> None:
	ax_client = AxClient(random_seed=RANDOM_SEED)
	ax_client.create_experiment(
	parameters=[
	{"name": "x", "type": "range", "bounds": [-5.0, 10.0]},
	{"name": "y", "type": "range", "bounds": [0.0, 15.0]},
	],
	name="fixed_features",
	)
	with mock.patch.object(
	GenerationStrategy, "gen", wraps=ax_client.generation_strategy.gen
	) as mock_gen:
	with self.subTest("fixed_features is None"):
	params, idx = ax_client.get_next_trial()
	call_kwargs = mock_gen.call_args_list[0][1]
	ff = call_kwargs["fixed_features"]
	self.assertIsNone(ff)
	with self.subTest("fixed_features is set"):
	fixed_features = FixedFeatures(
	parameters={"x": 0.0, "y": 5.0}, trial_index=0
	)
	params, idx = ax_client.get_next_trial(fixed_features=fixed_features)
	call_kwargs = mock_gen.call_args_list[1][1]
	ff = call_kwargs["fixed_features"]
	self.assertEqual(ff.parameters, fixed_features.parameters)
	self.assertEqual(ff.trial_index, 0)