Scheduler updates experiment status from generation strategy

[CristianLara]

Summary:

Summary

This change enables the Ax Scheduler (Orchestrator) to receive and propagate suggested experiment status from the generation strategy, then persist it to the database. This is a key integration point that connects the experiment lifecycle tracking system to the scheduler workflow.

What Changed

1. Updated GenerationStrategy.gen() signature (/data/sandcastle/boxes/fbsource/fbcode/ax/generation_strategy/generation_strategy.py):

   • Changed return type from list[list[GeneratorRun]] to tuple including status

2. Updated Orchestrator flow (/data/sandcastle/boxes/fbsource/fbcode/ax/service/orchestrator.py):

   • Modified _get_next_trials() to return suggested experiment status as third element in tuple
   • Updated _gen_new_trials_from_generation_strategy() to receive and return suggested status from generation_strategy.gen()
   • Updated generate_candidates() to receive suggested status and pass it to _save_or_update_trials_and_generation_strategy_if_possible()

3. Updated database persistence (/data/sandcastle/boxes/fbsource/fbcode/ax/service/utils/with_db_settings_base.py):

   • Added suggested_experiment_status parameter to _save_or_update_trials_and_generation_strategy_if_possible()
   • Calls _update_experiment_status_in_db_if_possible() when suggested status is provided

Why This Change

This integrates the experiment lifecycle tracking feature (added in D86801911, D88089767, D88091834, D88096914) into the Scheduler/Orchestrator workflow. When the generation strategy suggests an experiment status (e.g., INITIALIZATION for Sobol exploration, OPTIMIZATION for model-based optimization), the scheduler now:

1. Receives this suggestion from the generation strategy
2. Updates the experiment status accordingly
3. Persists the status to the database

This enables automatic experiment lifecycle management based on what optimization phase the generation strategy is in, providing users with clear visibility into experiment progression.

Differential Revision: D87589267

[meta-codesync]

@CristianLara has exported this pull request. If you are a Meta employee, you can view the originating Diff in D87589267.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 99.15254% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 96.80%. Comparing base (5cf011e) to head (0303af7).

Files with missing lines	Patch %	Lines
ax/storage/sqa_store/save.py	92.85%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #4891      +/-   ##
==========================================
+ Coverage   96.79%   96.80%   +0.01%     
==========================================
  Files         595      595              
  Lines       62841    62949     +108     
==========================================
+ Hits        60826    60937     +111     
+ Misses       2015     2012       -3     

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

[meta-codesync]

This pull request has been merged in f0bd257.