Skip to content

Add fit_tracking_metrics option to configure_generation_strategy#5238

Open
saitcakmak wants to merge 2 commits into
facebook:mainfrom
saitcakmak:fit-tracking-metrics-dispatch
Open

Add fit_tracking_metrics option to configure_generation_strategy#5238
saitcakmak wants to merge 2 commits into
facebook:mainfrom
saitcakmak:fit-tracking-metrics-dispatch

Conversation

@saitcakmak

Copy link
Copy Markdown
Contributor

Summary

Implements the feature requested in #4688: a fit_tracking_metrics option on the high-level Client.configure_generation_strategy API, so users can disable modeling of tracking metrics without manually editing the generation strategy.

client.configure_generation_strategy(fit_tracking_metrics=False)

When False, only metrics in the optimization config (objectives and outcome constraints) are modeled. Tracking metrics are still recorded but not modeled by the BoTorch model — useful when many tracking metrics are kept only for book-keeping.

Changes

  • ax/api/client.py — added fit_tracking_metrics: bool = True to configure_generation_strategy and _choose_generation_strategy, threading it into the dispatch struct. Docstrings flag that when False, model-dependent analyses (cross-validation, sensitivity analysis, etc.) will not be produced for tracking metrics, that it requires an optimization config, and that it has no effect for random_search.
  • ax/api/utils/structs.py — added the fit_tracking_metrics field to GenerationStrategyDispatchStruct.
  • ax/api/utils/generation_strategy_dispatch.py_get_mbm_node injects fit_tracking_metrics into the MBM node's generator_kwargs; only the MBM node receives it (Sobol/random-search nodes don't model outcomes).
  • Tests — added test_gs_fit_tracking_metrics (dispatch) and test_configure_generation_strategy_with_fit_tracking_metrics (client); updated existing tests that assert the full MBM generator_kwargs dict.

Implementation note

fit_tracking_metrics is a TorchAdapter.__init__ kwarg. The registry's kwarg-splitter (ax/adapter/registry.py) automatically routes any key matching an adapter constructor param into the adapter, so adding it to the node's generator_kwargs is sufficient — no changes to registry.py, torch.py, or base.py.

Test plan

  • Dispatch tests: 13 passed. Client tests pass. flake8 + ufmt clean.
  • End-to-end verified with a genuine tracking metric m2: fit_tracking_metrics=True → modeled outcomes ['m1', 'm2']; False['m1'] (tracking metric excluded from modeling).

Closes #4688

Allows disabling modeling of tracking metrics via the high-level Client
interface (client.configure_generation_strategy(fit_tracking_metrics=False))
instead of manually editing the generation strategy.

The option is threaded through GenerationStrategyDispatchStruct and injected
into the MBM node's generator_kwargs, where the registry routes it to the
TorchAdapter constructor. When False, only metrics in the optimization config
are modeled; tracking metrics are still recorded but not modeled, and
model-dependent analyses are not produced for them.

Closes facebook#4688
@meta-cla meta-cla Bot added the CLA Signed Do not delete this pull request or issue due to inactivity. label Jun 25, 2026
Move pyrefly-ignore comment to its own line (matching repo convention)
to bring the line under 88 chars. Pre-existing error from facebook#5236, unrelated
to this PR's feature but surfaced by pre-commit.
@meta-codesync

meta-codesync Bot commented Jun 25, 2026

Copy link
Copy Markdown

@saitcakmak has imported this pull request. If you are a Meta employee, you can view this in D109735998.

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.56%. Comparing base (025d7e5) to head (b857fff).

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #5238   +/-   ##
=======================================
  Coverage   96.56%   96.56%           
=======================================
  Files         619      619           
  Lines       70273    70296   +23     
=======================================
+ Hits        67861    67884   +23     
  Misses       2412     2412           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed Do not delete this pull request or issue due to inactivity.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[FEATURE REQUEST]: setting fit_tracking_metrics in configure_generation_strategy

2 participants