Change func name (#38)

* change start_trial to run_trial

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

* Change func name

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

---------

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

Author: kerthcet

alphatrion/experiment/base.py

@@ -47,7 +47,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 
     @classmethod
     @abstractmethod
-    def run(
+    def start(
         cls, name: str, description: str | None = None, meta: dict | None = None
     ) -> "Experiment":
         """Return a new experiment."""

alphatrion/experiment/craft_exp.py

@@ -16,7 +16,7 @@ def __init__(self):
         super().__init__()
 
     @classmethod
-    def run(
+    def start(
         cls,
         name: str,
         id: uuid.UUID | None = None,
@@ -50,9 +50,10 @@ def start_trial(
     ) -> Trial:
         """
         start_trial starts a new trial in this experiment.
-        You need to call trial.stop() to stop the trial for proper cleanup,
-        unless it's a timeout trial. Or you can use 'async with exp.run_trial(...)'
-        as trial, which will automatically stop the trial at the end of the context.
+        You need to call trial.cancel() to stop the trial for proper cleanup,
+        unless it's a timeout trial.
+        Or you can use 'async with exp.start_trial(...) as trial', which will
+        automatically stop the trial at the end of the context.
 
         :params description: the description of the trial
         :params meta: the metadata of the trial

alphatrion/trial/trial.py

@@ -103,7 +103,7 @@ async def __aenter__(self):
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
-        self.stop()
+        self.cancel()
 
     def _construct_meta(self):
         self._meta = dict()
@@ -167,7 +167,7 @@ def _timeout(self) -> int | None:
     def stopped(self) -> bool:
         return self._context.cancelled()
 
-    async def wait_stopped(self):
+    async def wait(self):
         await self._context.wait_cancelled()
 
     def _start(
@@ -194,8 +194,8 @@ def _start(
     def id(self) -> uuid.UUID:
         return self._id
 
-    # stop function should be called manually as a pair of start
-    def stop(self):
+    # cancel function should be called manually as a pair of start
+    def cancel(self):
         self._context.cancel()
 
     def _stop(self):

tests/integration/test_log.py

@@ -13,7 +13,7 @@
 async def test_log_artifact():
     alpha.init(project_id="test_project", artifact_insecure=True)
 
-    async with alpha.CraftExperiment.run(
+    async with alpha.CraftExperiment.start(
         name="context_exp",
         description="Context manager test",
         meta={"key": "value"},
@@ -49,7 +49,7 @@ async def test_log_artifact():
         versions = exp._runtime._artifact.list_versions(exp_obj.uuid)
         assert len(versions) == 0
 
-        trial.stop()
+        trial.cancel()
 
         got_exp = exp._runtime._metadb.get_exp(exp_id=exp._id)
         assert got_exp is not None
@@ -65,7 +65,7 @@ async def test_log_artifact():
 async def test_log_params():
     alpha.init(project_id="test_project", artifact_insecure=True)
 
-    async with alpha.CraftExperiment.run(name="test_experiment") as exp:
+    async with alpha.CraftExperiment.start(name="test_experiment") as exp:
         trial = exp.start_trial(description="First trial", params={"param1": 0.1})
 
         new_trial = exp._runtime._metadb.get_trial(trial_id=trial.id)
@@ -81,18 +81,18 @@ async def test_log_params():
         assert new_trial.status == TrialStatus.RUNNING
         assert current_trial_id.get() == trial.id
 
-        trial.stop()
+        trial.cancel()
 
         trial = exp.start_trial(description="Second trial", params={"param1": 0.1})
         assert current_trial_id.get() == trial.id
-        trial.stop()
+        trial.cancel()
 
 
 @pytest.mark.asyncio
 async def test_log_metrics():
     alpha.init(project_id="test_project", artifact_insecure=True)
 
-    async with alpha.CraftExperiment.run(name="test_experiment") as exp:
+    async with alpha.CraftExperiment.start(name="test_experiment") as exp:
         trial = exp.start_trial(description="First trial", params={"param1": 0.1})
 
         new_trial = exp._runtime._metadb.get_trial(trial_id=trial._id)
@@ -121,14 +121,14 @@ async def test_log_metrics():
         assert metrics[2].value == 0.96
         assert metrics[2].step == 2
 
-        trial.stop()
+        trial.cancel()
 
 
 @pytest.mark.asyncio
 async def test_log_metrics_with_save_on_max():
     alpha.init(project_id="test_project", artifact_insecure=True)
 
-    async with alpha.CraftExperiment.run(
+    async with alpha.CraftExperiment.start(
         name="context_exp",
         description="Context manager test",
         meta={"key": "value"},
@@ -182,7 +182,7 @@ async def test_log_metrics_with_save_on_max():
 async def test_log_metrics_with_save_on_min():
     alpha.init(project_id="test_project", artifact_insecure=True)
 
-    async with alpha.CraftExperiment.run(
+    async with alpha.CraftExperiment.start(
         name="context_exp",
         description="Context manager test",
         meta={"key": "value"},

tests/unit/experiment/test_craft_exp.py

@@ -1,19 +1,20 @@
 import asyncio
 import random
+from datetime import datetime, timedelta
 
 import pytest
 
 from alphatrion.experiment.craft_exp import CraftExperiment
 from alphatrion.metadata.sql_models import TrialStatus
 from alphatrion.runtime.runtime import init
-from alphatrion.trial.trial import TrialConfig, current_trial_id
+from alphatrion.trial.trial import Trial, TrialConfig, current_trial_id
 
 
 @pytest.mark.asyncio
 async def test_craft_experiment():
     init(project_id="test_project", artifact_insecure=True)
 
-    async with CraftExperiment.run(
+    async with CraftExperiment.start(
         name="context_exp",
         description="Context manager test",
         meta={"key": "value"},
@@ -28,7 +29,7 @@ async def test_craft_experiment():
         assert trial_obj is not None
         assert trial_obj.description == "First trial"
 
-        trial.stop()
+        trial.cancel()
 
         trial2 = trial._get_obj()
         assert trial2.status == TrialStatus.FINISHED
@@ -39,7 +40,7 @@ async def test_create_experiment_with_trial():
     init(project_id="test_project", artifact_insecure=True)
 
     trial_id = None
-    async with CraftExperiment.run(name="context_exp") as exp:
+    async with CraftExperiment.start(name="context_exp") as exp:
         async with exp.start_trial(description="First trial") as trial:
             trial_obj = trial._get_obj()
             assert trial_obj is not None
@@ -50,19 +51,44 @@ async def test_create_experiment_with_trial():
         assert trial_obj.status == TrialStatus.FINISHED
 
 
+@pytest.mark.asyncio
+async def test_create_experiment_with_trial_wait():
+    init(project_id="test_project", artifact_insecure=True)
+
+    async def fake_work(trial: Trial):
+        await asyncio.sleep(3)
+        trial.cancel()
+
+    trial_id = None
+    async with CraftExperiment.start(name="context_exp") as exp:
+        async with exp.start_trial(description="First trial") as trial:
+            trial_id = current_trial_id.get()
+
+            start_time = datetime.now()
+
+            asyncio.create_task(fake_work(trial))
+            assert datetime.now() - start_time <= timedelta(seconds=1)
+
+            await trial.wait()
+            assert datetime.now() - start_time >= timedelta(seconds=3)
+
+        trial_obj = exp._runtime._metadb.get_trial(trial_id=trial_id)
+        assert trial_obj.status == TrialStatus.FINISHED
+
+
 @pytest.mark.asyncio
 async def test_craft_experiment_with_context():
     init(project_id="test_project", artifact_insecure=True)
 
-    async with CraftExperiment.run(
+    async with CraftExperiment.start(
         name="context_exp",
         description="Context manager test",
         meta={"key": "value"},
     ) as exp:
         trial = exp.start_trial(
             description="First trial", config=TrialConfig(max_duration_seconds=2)
         )
-        await trial.wait_stopped()
+        await trial.wait()
         assert trial.stopped()
 
         trial = trial._get_obj()
@@ -81,15 +107,15 @@ async def fake_work(exp: CraftExperiment):
         # double check current trial id.
         assert trial.id == current_trial_id.get()
 
-        await trial.wait_stopped()
+        await trial.wait()
         assert trial.stopped()
         # we don't reset the current trial id.
         assert trial.id == current_trial_id.get()
 
         trial = trial._get_obj()
         assert trial.status == TrialStatus.FINISHED
 
-    async with CraftExperiment.run(
+    async with CraftExperiment.start(
         name="context_exp",
         description="Context manager test",
         meta={"key": "value"},