Remove ExperimentConfig (#77)

Signed-off-by: kerthcet <[email protected]>

Author: kerthcet

alphatrion/__init__.py

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

alphatrion/experiment/base.py

@@ -2,25 +2,10 @@
 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):
     """
@@ -29,8 +14,7 @@ class Experiment(ABC):
 
     __slots__ = ("_runtime", "_id", "_trials")
 
-    def __init__(self, config: ExperimentConfig | None = None):
-        self._config = config or ExperimentConfig()
+    def __init__(self):
         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, ExperimentConfig
+from alphatrion.experiment.base import Experiment
 from alphatrion.trial.trial import Trial, TrialConfig
 
 
@@ -10,23 +10,22 @@ class CraftExperiment(Experiment):
     Opposite to other experiment classes, you need to call all these methods yourself.
     """
 
-    def __init__(self, config: ExperimentConfig | None = None):
-        super().__init__(config=config)
+    def __init__(self):
+        super().__init__()
 
     @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(config=config)
+        exp = CraftExperiment()
         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.
@@ -63,15 +62,6 @@ 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)

tests/unit/experiment/test_craft_exp.py

@@ -5,7 +5,7 @@
 
 import pytest
 
-from alphatrion.experiment.craft_exp import CraftExperiment, ExperimentConfig
+from alphatrion.experiment.craft_exp import CraftExperiment
 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
@@ -194,39 +194,3 @@ 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.COMPLETED
-
-
-@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.COMPLETED