Expose ExperimentConfig (#71)

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

Makefile

@@ -1,12 +1,12 @@
 POETRY := poetry
 
 .PHONY: build
-build:
+build: lint
 	$(POETRY) build
 
 .PHONY: publish
 publish: build
-	$(POETRY) publish
+	$(POETRY) publish --username=__token__ --password=$(INFTYAI_PYPI_TOKEN)
 
 .PHONY: up
 up:

alphatrion/__init__.py

@@ -1,4 +1,4 @@
-from alphatrion.experiment.craft_exp import CraftExperiment
+from alphatrion.experiment.craft_exp import CraftExperiment, ExperimentConfig
 from alphatrion.log.log import log_artifact, log_metrics, log_params
 from alphatrion.runtime.runtime import init
 from alphatrion.tracing.tracing import task, workflow
@@ -9,6 +9,7 @@
     "log_params",
     "log_metrics",
     "CraftExperiment",
+    "ExperimentConfig",
     "init",
     "TrialConfig",
     "CheckpointConfig",

alphatrion/experiment/base.py

@@ -2,10 +2,25 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 
+from pydantic import BaseModel, Field
+
 from alphatrion.runtime.runtime import global_runtime
 from alphatrion.trial import trial
 
 
+class ExperimentConfig(BaseModel):
+    """
+    Configuration for Experiment.
+    """
+
+    max_runtime_seconds: int = Field(
+        default=-1,
+        description="Maximum runtime seconds for the experiment. \
+        It will overwrite the trial timeout if both are set. \
+        Default is -1 (no limit).",
+    )
+
+
 @dataclass
 class Experiment(ABC):
     """
@@ -14,7 +29,8 @@ class Experiment(ABC):
 
     __slots__ = ("_runtime", "_id", "_trials")
 
-    def __init__(self):
+    def __init__(self, config: ExperimentConfig | None = None):
+        self._config = config or ExperimentConfig()
         self._runtime = global_runtime()
         # All trials in this experiment, key is trial_id, value is Trial instance.
         self._trials = dict()

alphatrion/experiment/craft_exp.py

@@ -1,4 +1,4 @@
-from alphatrion.experiment.base import Experiment
+from alphatrion.experiment.base import Experiment, ExperimentConfig
 from alphatrion.trial.trial import Trial, TrialConfig
 
 
@@ -10,22 +10,23 @@ class CraftExperiment(Experiment):
     Opposite to other experiment classes, you need to call all these methods yourself.
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, config: ExperimentConfig | None = None):
+        super().__init__(config=config)
 
     @classmethod
     def setup(
         cls,
         name: str,
         description: str | None = None,
         meta: dict | None = None,
+        config: ExperimentConfig | None = None,
     ) -> "CraftExperiment":
         """
         Setup the experiment. If the name already exists in the same project,
         it will refer to the existing experiment instead of creating a new one.
         """
 
-        exp = CraftExperiment()
+        exp = CraftExperiment(config=config)
         exp_obj = exp._get_by_name(name=name, project_id=exp._runtime._project_id)
 
         # If experiment with the same name exists in the project, use it.
@@ -62,6 +63,15 @@ def start_trial(
         :return: the Trial instance
         """
 
+        config = config or TrialConfig()
+
+        if (
+            self._config is not None
+            and self._config.max_runtime_seconds > 0
+            and config.max_runtime_seconds < 0
+        ):
+            config.max_runtime_seconds = self._config.max_runtime_seconds
+
         trial = Trial(exp_id=self._id, config=config)
         trial._start(name=name, description=description, meta=meta, params=params)
         self.register_trial(id=trial.id, instance=trial)

alphatrion/trial/trial.py

@@ -47,9 +47,10 @@ class CheckpointConfig(BaseModel):
 class TrialConfig(BaseModel):
     """Configuration for an experiment."""
 
-    max_duration_seconds: int = Field(
+    max_runtime_seconds: int = Field(
         default=-1,
-        description="Maximum duration in seconds for the experiment. \
+        description="Maximum runtime seconds for the trial. \
+        Trial timeout will override experiment timeout if both are set. \
         Default is -1 (no limit).",
     )
     early_stopping_runs: int = Field(
@@ -216,7 +217,7 @@ def should_early_stop(self, metric_key: str, metric_value: float) -> bool:
         return self._early_stopping_counter >= self._config.early_stopping_runs
 
     def _timeout(self) -> int | None:
-        timeout = self._config.max_duration_seconds
+        timeout = self._config.max_runtime_seconds
         if timeout < 0:
             return None
 

tests/integration/test_log.py

@@ -325,7 +325,7 @@ async def fake_sleep(value: float):
             config=alpha.TrialConfig(
                 monitor_metric="accuracy",
                 early_stopping_runs=3,
-                max_duration_seconds=3,
+                max_runtime_seconds=3,
             ),
         ) as trial:
             start_time = datetime.now()

tests/unit/experiment/test_craft_exp.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from alphatrion.experiment.craft_exp import CraftExperiment
+from alphatrion.experiment.craft_exp import CraftExperiment, ExperimentConfig
 from alphatrion.metadata.sql_models import TrialStatus
 from alphatrion.runtime.runtime import global_runtime, init
 from alphatrion.trial.trial import Trial, TrialConfig, current_trial_id
@@ -35,6 +35,7 @@ async def test_craft_experiment():
         trial_obj = trial._get_obj()
         assert trial_obj.status == TrialStatus.FINISHED
 
+
 @pytest.mark.asyncio
 async def test_craft_experiment_with_no_context():
     init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
@@ -51,6 +52,7 @@ async def fake_work(trial: Trial):
         trial_obj = trial._get_obj()
         assert trial_obj.status == TrialStatus.FINISHED
 
+
 @pytest.mark.asyncio
 async def test_create_experiment_with_trial():
     init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
@@ -129,7 +131,7 @@ async def test_craft_experiment_with_context():
         meta={"key": "value"},
     ) as exp:
         trial = exp.start_trial(
-            name="first-trial", config=TrialConfig(max_duration_seconds=2)
+            name="first-trial", config=TrialConfig(max_runtime_seconds=2)
         )
         await trial.wait()
         assert trial.cancelled()
@@ -147,7 +149,7 @@ async def fake_work():
 
         duration = random.randint(1, 5)
         trial = exp.start_trial(
-            name="first-trial", config=TrialConfig(max_duration_seconds=duration)
+            name="first-trial", config=TrialConfig(max_runtime_seconds=duration)
         )
         # double check current trial id.
         assert trial.id == current_trial_id.get()
@@ -171,3 +173,39 @@ async def fake_work():
             fake_work(),
         )
         print("All trials finished.")
+
+
+@pytest.mark.asyncio
+async def test_craft_experiment_with_timeout():
+    init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
+
+    exp = CraftExperiment.setup(
+        name="timeout_exp",
+        config=ExperimentConfig(max_runtime_seconds=3),
+    )
+
+    async with exp.start_trial(name="first-trial") as trial:
+        await trial.wait()
+
+        trial_obj = trial._get_obj()
+        assert trial_obj.status == TrialStatus.FINISHED
+
+
+@pytest.mark.asyncio
+async def test_craft_experiment_with_timeout_overwrite():
+    init(project_id=uuid.uuid4(), artifact_insecure=True, init_tables=True)
+
+    exp = CraftExperiment.setup(
+        name="timeout_exp",
+        config=ExperimentConfig(max_runtime_seconds=3),
+    )
+
+    start_time = datetime.now()
+    async with exp.start_trial(
+        name="first-trial", config=TrialConfig(max_runtime_seconds=1)
+    ) as trial:
+        await trial.wait()
+        assert datetime.now() - start_time < timedelta(seconds=3)
+
+        trial_obj = trial._get_obj()
+        assert trial_obj.status == TrialStatus.FINISHED

tests/unit/trial/test_trial.py

@@ -18,25 +18,25 @@ def test_timeout(self):
             },
             {
                 "name": "Positive timeout",
-                "config": TrialConfig(max_duration_seconds=10),
+                "config": TrialConfig(max_runtime_seconds=10),
                 "started_at": None,
                 "expected": 10,
             },
             {
                 "name": "Zero timeout",
-                "config": TrialConfig(max_duration_seconds=0),
+                "config": TrialConfig(max_runtime_seconds=0),
                 "started_at": None,
                 "expected": 0,
             },
             {
                 "name": "Negative timeout",
-                "config": TrialConfig(max_duration_seconds=-5),
+                "config": TrialConfig(max_runtime_seconds=-5),
                 "started_at": None,
                 "expected": None,
             },
             {
                 "name": "With started_at, positive timeout",
-                "config": TrialConfig(max_duration_seconds=5),
+                "config": TrialConfig(max_runtime_seconds=5),
                 "started_at": (datetime.now(UTC) - timedelta(seconds=3)).isoformat(),
                 "expected": 2,
             },