finish trials when exciting experiment (#74)

* finish trials when exciting experiment

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

* fix test error

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

* rename cancelled to completed

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

* rename cancelled to done

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

* release v0.0.4

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

* rename FINISHED to COMPLETED

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

---------

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

Author: kerthcet

README.md

@@ -34,7 +34,6 @@ pip install alphatrion
 ### Install from Source
 
 * Git clone the repository
-* Run `uv sync` to install dependencies from `pyproject.toml`.
 * Run `source start.sh` to activate the virtual environment.
 
 

alphatrion/experiment/base.py

@@ -54,7 +54,14 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.complete()
+
+    def complete(self):
+        for t in list(self._trials.values()):
+            t.complete()
         self._trials = dict()
+        # Set to None at the end of the experiment because
+        # it will be used in trial.complete().
         self._runtime.current_exp = None
 
     @classmethod

alphatrion/log/log.py

@@ -102,4 +102,4 @@ async def log_metrics(metrics: dict[str, float]):
         )
 
     if should_early_stop:
-        trial.cancel()
+        trial.complete()

alphatrion/metadata/sql_models.py

@@ -12,11 +12,11 @@
 class TrialStatus(enum.Enum):
     PENDING = "pending"
     RUNNING = "running"
-    FINISHED = "finished"
+    COMPLETED = "completed"
     FAILED = "failed"
 
 
-COMPLETED_STATUS = [TrialStatus.FINISHED, TrialStatus.FAILED]
+FINISHED_STATUS = [TrialStatus.COMPLETED, TrialStatus.FAILED]
 
 
 class Project(Base):
@@ -26,9 +26,11 @@ class Project(Base):
     name = Column(String, nullable=False)
     description = Column(String, nullable=True)
 
-    created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
     updated_at = Column(
-        DateTime(timezone=True), default=datetime.now(UTC), onupdate=datetime.now(UTC)
+        DateTime(timezone=True),
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
     )
     is_del = Column(Integer, default=0, comment="0 for not deleted, 1 for deleted")
 
@@ -43,9 +45,11 @@ class Experiment(Base):
     project_id = Column(UUID(as_uuid=True), nullable=False)
     meta = Column(JSON, nullable=True, comment="Additional metadata for the experiment")
 
-    created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
     updated_at = Column(
-        DateTime(timezone=True), default=datetime.now(UTC), onupdate=datetime.now(UTC)
+        DateTime(timezone=True),
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
     )
     is_del = Column(Integer, default=0, comment="0 for not deleted, 1 for deleted")
 
@@ -60,16 +64,19 @@ class Trial(Base):
     description = Column(String, nullable=True)
     meta = Column(JSON, nullable=True, comment="Additional metadata for the trial")
     params = Column(JSON, nullable=True, comment="Parameters for the experiment")
+    duration = Column(Float, nullable=True, comment="Duration of the trial in seconds")
     status = Column(
         Enum(TrialStatus),
         default=TrialStatus.PENDING,
         nullable=False,
         comment="Status of the trial",
     )
 
-    created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
     updated_at = Column(
-        DateTime(timezone=True), default=datetime.now(UTC), onupdate=datetime.now(UTC)
+        DateTime(timezone=True),
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
     )
 
 
@@ -80,9 +87,11 @@ class Run(Base):
     project_id = Column(UUID(as_uuid=True), nullable=False)
     trial_id = Column(UUID(as_uuid=True), nullable=False)
 
-    created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
     updated_at = Column(
-        DateTime(timezone=True), default=datetime.now(UTC), onupdate=datetime.now(UTC)
+        DateTime(timezone=True),
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
     )
 
 
@@ -96,9 +105,11 @@ class Model(Base):
     version = Column(String, nullable=False)
     meta = Column(JSON, nullable=True, comment="Additional metadata for the model")
 
-    created_at = Column(DateTime(timezone=True), default=datetime.now(UTC))
+    created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC))
     updated_at = Column(
-        DateTime(timezone=True), default=datetime.now(UTC), onupdate=datetime.now(UTC)
+        DateTime(timezone=True),
+        default=lambda: datetime.now(UTC),
+        onupdate=lambda: datetime.now(UTC),
     )
     is_del = Column(Integer, default=0, comment="0 for not deleted, 1 for deleted")
 

alphatrion/trial/trial.py

@@ -1,11 +1,10 @@
 import contextvars
-import os
 import uuid
 from datetime import UTC, datetime
 
 from pydantic import BaseModel, Field, model_validator
 
-from alphatrion.metadata.sql_models import COMPLETED_STATUS, TrialStatus
+from alphatrion.metadata.sql_models import FINISHED_STATUS, TrialStatus
 from alphatrion.run.run import Run
 from alphatrion.runtime.runtime import global_runtime
 from alphatrion.utils.context import Context
@@ -127,10 +126,6 @@ def __init__(self, exp_id: int, config: TrialConfig | None = None):
         self._config = config or TrialConfig()
         self._runtime = global_runtime()
         self._step = 0
-        self._context = Context(
-            cancel_func=self._stop,
-            timeout=self._timeout(),
-        )
         self._construct_meta()
         self._runs = dict()
         self._running_tasks = dict()
@@ -141,7 +136,7 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self.cancel()
+        self.complete()
         if self._token:
             current_trial_id.reset(self._token)
 
@@ -218,19 +213,18 @@ def should_early_stop(self, metric_key: str, metric_value: float) -> bool:
 
     def _timeout(self) -> int | None:
         timeout = self._config.max_runtime_seconds
-        if timeout < 0:
+        if timeout is None or timeout < 0:
             return None
 
-        # Adjust timeout based on the trial start time from environment variable,
-        # this is useful when running in cloud env when the trial process may be
-        # restarted.
-        start_time = os.environ.get("ALPHATRION_TRIAL_START_TIME", None)
-        if start_time is not None:
-            elapsed = (
-                datetime.now(UTC)
-                - datetime.fromisoformat(start_time).replace(tzinfo=UTC)
-            ).total_seconds()
-            timeout -= int(elapsed)
+        obj = self._get_obj()
+        if obj is None:
+            return timeout
+
+        elapsed = (
+            datetime.now(UTC) - obj.created_at.replace(tzinfo=UTC)
+        ).total_seconds()
+        timeout -= int(elapsed)
+
         return timeout
 
     # Make sure you have termination condition, either by timeout or by calling cancel()
@@ -240,7 +234,7 @@ def _timeout(self) -> int | None:
     async def wait(self):
         await self._context.wait()
 
-    def cancelled(self) -> bool:
+    def done(self) -> bool:
         return self._context.cancelled()
 
     # If the name is same in the same experiment, it will refer to the existing trial.
@@ -254,7 +248,7 @@ def _start(
         trial_obj = self._runtime._metadb.get_trial_by_name(
             trial_name=name, exp_id=self._exp_id
         )
-        # FIXME: what if the existing trial is finished, will lead to confusion?
+        # FIXME: what if the existing trial is completed, will lead to confusion?
         if trial_obj:
             self._id = trial_obj.uuid
         else:
@@ -268,25 +262,29 @@ def _start(
                 status=TrialStatus.RUNNING,
             )
 
+        self._context = Context(
+            cancel_func=self._stop,
+            timeout=self._timeout(),
+        )
+
         # We don't reset the trial id context var here, because
         # each trial runs in its own context.
         self._token = current_trial_id.set(self._id)
-        self._context.start()
 
-    # cancel function should be called manually as a pair of start
+    # complete function should be called manually as a pair of start
     # FIXME: watch for system signals to cancel the trial gracefully,
-    # or it could lead to trial not being marked as finished.
-    def cancel(self):
+    # or it could lead to trial not being marked as completed.
+    def complete(self):
         self._context.cancel()
 
     def _stop(self):
         trial = self._runtime._metadb.get_trial(trial_id=self._id)
-        if trial is not None and trial.status not in COMPLETED_STATUS:
+        if trial is not None and trial.status not in FINISHED_STATUS:
             duration = (
                 datetime.now(UTC) - trial.created_at.replace(tzinfo=UTC)
             ).total_seconds()
             self._runtime._metadb.update_trial(
-                trial_id=self._id, status=TrialStatus.FINISHED, duration=duration
+                trial_id=self._id, status=TrialStatus.COMPLETED, duration=duration
             )
 
         self._runtime.current_exp.unregister_trial(self._id)
@@ -321,7 +319,7 @@ def start_run(self, call_func: callable) -> Run:
             task.add_done_callback(
                 lambda t: (
                     setattr(self, "_total_runs_counter", self._total_runs_counter + 1),
-                    self.cancel()
+                    self.complete()
                     if self._total_runs_counter >= self._config.max_runs_per_trial
                     else None,
                 )

alphatrion/utils/context.py

@@ -4,7 +4,7 @@
 
 # Inspired by golang context package
 class Context:
-    def __init__(self, cancel_func: Callable | None = None, timeout=None):
+    def __init__(self, cancel_func: Callable | None = None, timeout: int | None = None):
         """A context for managing cancellation and timeouts.
         :param cancel_func: A function to call when the context is cancelled.
         :param timeout: Timeout in seconds. If None, no timeout is set.
@@ -13,9 +13,6 @@ def __init__(self, cancel_func: Callable | None = None, timeout=None):
         self._cancel_func = cancel_func
         self._timeout = timeout
 
-    def start(self):
-        # If timeout is None, it means no timeout is set.
-        # If timeout is negative, it means already timed out.
         if self._timeout is not None:
             asyncio.create_task(self._auto_cancel(self._timeout))
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphatrion"
-version = "0.0.3"
+version = "0.0.4"
 description = "⚒️ AlphaTrion is an open-source framework to help build GenAI applications, including experiment tracking, adaptive model routing, prompt optimization and performance evaluation."
 license = {text = "Apache-2.0"}
 readme = "README.md"

tests/integration/test_log.py

@@ -52,7 +52,7 @@ async def test_log_artifact():
         versions = exp._runtime._artifact.list_versions(exp_obj.uuid)
         assert len(versions) == 0
 
-        trial.cancel()
+        trial.complete()
 
         got_exp = exp._runtime._metadb.get_exp(exp_id=exp._id)
         assert got_exp is not None
@@ -61,7 +61,7 @@ async def test_log_artifact():
         got_trial = exp._runtime._metadb.get_trial(trial_id=trial._id)
         assert got_trial is not None
         assert got_trial.name == "first-trial"
-        assert got_trial.status == TrialStatus.FINISHED
+        assert got_trial.status == TrialStatus.COMPLETED
 
 
 @pytest.mark.asyncio
@@ -84,11 +84,11 @@ async def test_log_params():
         assert new_trial.status == TrialStatus.RUNNING
         assert current_trial_id.get() == trial.id
 
-        trial.cancel()
+        trial.complete()
 
         trial = exp.start_trial(name="second-trial", params={"param1": 0.1})
         assert current_trial_id.get() == trial.id
-        trial.cancel()
+        trial.complete()
 
 
 @pytest.mark.asyncio
@@ -135,7 +135,7 @@ async def log_metric(metrics: dict):
         assert run_id_2 is not None
         assert run_id_2 != run_id_1
 
-        trial.cancel()
+        trial.complete()
 
 
 @pytest.mark.asyncio
@@ -201,7 +201,7 @@ async def log_metric(value: float):
             versions = exp._runtime._artifact.list_versions(exp.id)
             assert len(versions) == 3
 
-            trial.cancel()
+            trial.complete()
 
 
 @pytest.mark.asyncio
@@ -266,7 +266,7 @@ async def log_metric(value: float):
             versions = exp._runtime._artifact.list_versions(exp.id)
             assert len(versions) == 3
 
-            trial.cancel()
+            trial.complete()
 
 
 @pytest.mark.asyncio
@@ -356,7 +356,7 @@ async def fake_work(value: float):
                 max_runs_per_trial=5,
             ),
         ) as trial:
-            while not trial.cancelled():
+            while not trial.done():
                 run = trial.start_run(lambda: fake_work(1))
                 await run.wait()