Deal with total_concurrency and n weirdness: introduce ExperimentDesign.concurrency_limit (#4732)

Lena Kashtelyan · meta-codesync[bot] · commit 68d37c111613 · 2025-12-31T18:31:05.000-08:00
Summary: Pull Request resolved: #4732 As titled, adding a simple `ExperimentDesign` object. Putting it into properties for serialization for now, so as to not do duplicate work ahead of the storage refactor implementation (and also in case we change things while working on this stack). Differential Revision: D89770462 Privacy Context Container: L1307644
diff --git a/ax/core/experiment.py b/ax/core/experiment.py
@@ -14,6 +14,7 @@
 import warnings
 from collections import defaultdict, OrderedDict
 from collections.abc import Hashable, Iterable, Mapping, Sequence
+from dataclasses import dataclass
 from datetime import datetime
 from functools import partial, reduce
 from typing import Any, cast, Union
@@ -84,6 +85,29 @@
 }
 
 
+@dataclass
+class ExperimentDesign:
+    """Struct that holds "experiment design" configuration: these are
+    experiment-level settings that pertain to "how the experiment will be
+    run or conducted", but are agnostic to the specific evaluation
+    backend, to which the trials will be deployed.
+
+    NOTE: In the future, we might treat concurrency limit as expressed
+    in terms of "full arm equivalents" as opposed to just "number of arms",
+    to cover for the multi-fidelity cases.
+
+    Args:
+        concurrency_limit: Maximum number of arms to run within one or
+            multiple trials, in parallel. In experiments that consist of
+            `Trial`-s, this is equivalent to the total number of trials
+            that should run in parallel. In experiments with `BatchTrial`-s,
+            this total number of arms can be spread across one or
+            multiple `BatchTrial`-s.
+    """
+
+    concurrency_limit: int | None = None
+
+
 class Experiment(Base):
     """Base class for defining an experiment."""
 
@@ -150,6 +174,12 @@ def __init__(
         self._time_created: datetime = datetime.now()
         self._trials: dict[int, BaseTrial] = {}
         self._properties: dict[str, Any] = properties or {}
+        self._design: ExperimentDesign = ExperimentDesign()
+        # Restore ExperimentDesign from properties if present (for deserialization).
+        # TODO[drfreund, mpolson64]: Replace with proper storage as part of the
+        # refactor.
+        if (design_dict := self._properties.pop("design", None)) is not None:
+            self._design.concurrency_limit = design_dict.get("concurrency_limit")
 
         # Initialize trial type to runner mapping
         self._default_trial_type = default_trial_type
@@ -233,6 +263,11 @@ def experiment_type(self, experiment_type: str | None) -> None:
         """Set the type of the experiment."""
         self._experiment_type = experiment_type
 
+    @property
+    def design(self) -> ExperimentDesign:
+        """The experiment design configuration."""
+        return self._design
+
     @property
     def search_space(self) -> SearchSpace:
         """The search space for this experiment.
diff --git a/ax/core/tests/test_experiment_design.py b/ax/core/tests/test_experiment_design.py
@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+from ax.core import Experiment
+from ax.core.experiment import ExperimentDesign
+from ax.utils.common.testutils import TestCase
+from ax.utils.testing.core_stubs import get_branin_search_space
+
+
+class ExperimentDesignTest(TestCase):
+    """Tests for ExperimentDesign and related logic."""
+
+    def test_experiment_design_defaults(self) -> None:
+        """Test that ExperimentDesign has expected defaults."""
+        design = ExperimentDesign()
+        self.assertIsNone(design.concurrency_limit)
+
+    def test_experiment_design_with_concurrency_limit(self) -> None:
+        """Test ExperimentDesign with concurrency_limit set."""
+        design = ExperimentDesign(concurrency_limit=10)
+        self.assertEqual(design.concurrency_limit, 10)
+
+        design_zero = ExperimentDesign(concurrency_limit=0)
+        self.assertEqual(design_zero.concurrency_limit, 0)
+
+    def test_experiment_design_property(self) -> None:
+        """Test that Experiment.design property returns ExperimentDesign instance."""
+        experiment = Experiment(
+            name="test",
+            search_space=get_branin_search_space(),
+        )
+        self.assertIsInstance(experiment.design, ExperimentDesign)
+        self.assertIsNone(experiment.design.concurrency_limit)
+
+    def test_experiment_design_modification(self) -> None:
+        """Test that ExperimentDesign can be modified after experiment creation."""
+        experiment = Experiment(
+            name="test",
+            search_space=get_branin_search_space(),
+        )
+        # Modify concurrency_limit
+        experiment.design.concurrency_limit = 5
+        self.assertEqual(experiment.design.concurrency_limit, 5)
+
+        # Set to None
+        experiment.design.concurrency_limit = None
+        self.assertIsNone(experiment.design.concurrency_limit)
+
+    def test_experiment_design_restore_from_properties(self) -> None:
+        """Test that ExperimentDesign is restored from properties during init."""
+        # Simulate deserializing an experiment with design stored in properties
+        experiment = Experiment(
+            name="test",
+            search_space=get_branin_search_space(),
+            properties={"design": {"concurrency_limit": 25}},
+        )
+        # The design should be restored and the "design" key should be removed
+        # from properties
+        self.assertEqual(experiment.design.concurrency_limit, 25)
+        self.assertNotIn("design", experiment._properties)
+
+    def test_experiment_design_restore_from_properties_with_none(self) -> None:
+        """Test that ExperimentDesign handles None concurrency_limit in properties."""
+        experiment = Experiment(
+            name="test",
+            search_space=get_branin_search_space(),
+            properties={"design": {"concurrency_limit": None}},
+        )
+        self.assertIsNone(experiment.design.concurrency_limit)
+        # The "design" key should be removed from properties once it's consumed
+        # to recreate `ExperimentDesign`.
+        self.assertNotIn("design", experiment._properties)
+
+    def test_experiment_design_restore_from_properties_empty_dict(self) -> None:
+        """Test that ExperimentDesign handles empty design dict in properties."""
+        experiment = Experiment(
+            name="test",
+            search_space=get_branin_search_space(),
+            properties={"design": {}},
+        )
+        self.assertIsNone(experiment.design.concurrency_limit)
+        # The "design" key should be removed from properties once it's consumed
+        # to recreate `ExperimentDesign`.
+        self.assertNotIn("design", experiment._properties)
+
+    def test_experiment_design_not_affecting_other_properties(self) -> None:
+        """Test that ExperimentDesign restoration doesn't affect other properties."""
+        experiment = Experiment(
+            name="test",
+            search_space=get_branin_search_space(),
+            properties={
+                "design": {"concurrency_limit": 15},
+                "custom_property": "custom_value",
+                "another_property": 42,
+            },
+        )
+        self.assertEqual(experiment.design.concurrency_limit, 15)
+        # The "design" key should be removed from properties once it's consumed
+        # to recreate `ExperimentDesign`.
+        self.assertNotIn("design", experiment._properties)
+        self.assertEqual(experiment._properties["custom_property"], "custom_value")
+        self.assertEqual(experiment._properties["another_property"], 42)
diff --git a/ax/storage/json_store/encoders.py b/ax/storage/json_store/encoders.py
@@ -73,6 +73,13 @@
 
 def experiment_to_dict(experiment: Experiment) -> dict[str, Any]:
     """Convert Ax experiment to a dictionary."""
+    # Serialize ExperimentDesign into properties
+    properties = {
+        **experiment._properties,
+        "design": {
+            "concurrency_limit": experiment.design.concurrency_limit,
+        },
+    }
     return {
         "__type": experiment.__class__.__name__,
         "name": experiment._name,
@@ -87,7 +94,7 @@ def experiment_to_dict(experiment: Experiment) -> dict[str, Any]:
         "trials": experiment.trials,
         "is_test": experiment.is_test,
         "data_by_trial": experiment._data_by_trial,
-        "properties": experiment._properties,
+        "properties": properties,
         "_trial_type_to_runner": experiment._trial_type_to_runner,
     }
 
diff --git a/ax/storage/json_store/tests/test_json_store.py b/ax/storage/json_store/tests/test_json_store.py
@@ -1656,6 +1656,60 @@ def test_experiment_with_pruning_target_json_roundtrip(self) -> None:
             ).pruning_target_parameterization,
         )
 
+    def test_experiment_design_json_roundtrip(self) -> None:
+        """Test that ExperimentDesign is preserved through JSON serialization."""
+        # Setup: create experiment and set concurrency_limit
+        experiment = get_branin_experiment()
+        experiment.design.concurrency_limit = 42
+
+        # Execute: save and load experiment through JSON
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json") as f:
+            save_experiment(
+                experiment,
+                f.name,
+                encoder_registry=CORE_ENCODER_REGISTRY,
+                class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+            )
+            loaded_experiment = load_experiment(
+                f.name,
+                decoder_registry=CORE_DECODER_REGISTRY,
+                class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+            )
+
+        # Cleanup
+        os.remove(f.name)
+
+        # Assert: confirm ExperimentDesign is preserved
+        self.assertEqual(experiment, loaded_experiment)
+        self.assertEqual(loaded_experiment.design.concurrency_limit, 42)
+
+    def test_experiment_design_none_concurrency_json_roundtrip(self) -> None:
+        """Test that ExperimentDesign with None concurrency_limit is preserved."""
+        # Setup: create experiment with default (None) concurrency_limit
+        experiment = get_branin_experiment()
+        self.assertIsNone(experiment.design.concurrency_limit)
+
+        # Execute: save and load experiment through JSON
+        with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".json") as f:
+            save_experiment(
+                experiment,
+                f.name,
+                encoder_registry=CORE_ENCODER_REGISTRY,
+                class_encoder_registry=CORE_CLASS_ENCODER_REGISTRY,
+            )
+            loaded_experiment = load_experiment(
+                f.name,
+                decoder_registry=CORE_DECODER_REGISTRY,
+                class_decoder_registry=CORE_CLASS_DECODER_REGISTRY,
+            )
+
+        # Cleanup
+        os.remove(f.name)
+
+        # Assert: confirm ExperimentDesign is preserved with None
+        self.assertEqual(experiment, loaded_experiment)
+        self.assertIsNone(loaded_experiment.design.concurrency_limit)
+
     def test_multi_objective_from_json_warning(self) -> None:
         objectives = [get_objective()]
 
diff --git a/ax/storage/sqa_store/encoder.py b/ax/storage/sqa_store/encoder.py
@@ -227,6 +227,10 @@ def experiment_to_sqa(self, experiment: Experiment) -> SQAExperiment:
         elif experiment.runner:
             runners.append(self.runner_to_sqa(none_throws(experiment.runner)))
         properties = experiment._properties.copy()
+        # Serialize ExperimentDesign into properties
+        properties["design"] = {
+            "concurrency_limit": experiment.design.concurrency_limit,
+        }
         if (
             oc := experiment.optimization_config
         ) is not None and oc.pruning_target_parameterization is not None:
diff --git a/ax/storage/sqa_store/tests/test_sqa_store.py b/ax/storage/sqa_store/tests/test_sqa_store.py
@@ -339,6 +339,35 @@ def test_experiment_save_load(self) -> None:
             loaded_experiment = load_experiment(exp.name)
             self.assertEqual(loaded_experiment, exp)
 
+    def test_experiment_design_sqa_roundtrip(self) -> None:
+        """Test that ExperimentDesign is preserved through SQA serialization."""
+        # Create experiment and set concurrency_limit
+        experiment = get_experiment_with_batch_trial()
+        experiment.design.concurrency_limit = 42
+
+        # Save and load experiment through SQA
+        save_experiment(experiment)
+        loaded_experiment = load_experiment(experiment.name)
+
+        # Verify ExperimentDesign is preserved
+        self.assertEqual(loaded_experiment, experiment)
+        self.assertEqual(loaded_experiment.design.concurrency_limit, 42)
+
+    def test_experiment_design_none_concurrency_sqa_roundtrip(self) -> None:
+        """Test that ExperimentDesign with None concurrency_limit is preserved."""
+        # Create experiment with default (None) concurrency_limit
+        experiment = get_experiment_with_batch_trial()
+        experiment.name = "experiment_design_none_concurrency_test"
+        self.assertIsNone(experiment.design.concurrency_limit)
+
+        # Save and load experiment through SQA
+        save_experiment(experiment)
+        loaded_experiment = load_experiment(experiment.name)
+
+        # Verify ExperimentDesign with None is preserved
+        self.assertEqual(loaded_experiment, experiment)
+        self.assertIsNone(loaded_experiment.design.concurrency_limit)
+
     def test_saving_and_loading_experiment_with_aux_exp(self) -> None:
         aux_experiment = Experiment(
             name="test_aux_exp_in_SQAStoreTest",
