support loading start time from env (#33)

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

Author: kerthcet

alphatrion/trial/trial.py

@@ -1,4 +1,5 @@
 import contextvars
+import os
 import uuid
 from datetime import UTC, datetime
 
@@ -93,11 +94,26 @@ def __init__(self, exp_id: int, config: TrialConfig | None = None):
         self._step = 0
         self._context = Context(
             cancel_func=self._stop,
-            timeout=self._config.max_duration_seconds
-            if self._config.max_duration_seconds > 0
-            else None,
+            timeout=self._timeout(),
         )
 
+    def _timeout(self) -> int | None:
+        timeout = self._config.max_duration_seconds
+        if 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)
+        return timeout
+
     def stopped(self) -> bool:
         return self._context.cancelled()
 

alphatrion/utils/context.py

@@ -14,7 +14,9 @@ def __init__(self, cancel_func: Callable | None = None, timeout=None):
         self._timeout = timeout
 
     async def start(self):
-        if self._timeout:
+        # 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))
 
     async def _auto_cancel(self, timeout):

tests/unit/trial/test_trial.py

@@ -0,0 +1,50 @@
+import os
+import unittest
+from datetime import UTC, datetime, timedelta
+
+from alphatrion.trial.trial import Trial, TrialConfig
+
+
+class TestTrial(unittest.IsolatedAsyncioTestCase):
+    def test_timeout(self):
+        test_cases = [
+            {
+                "name": "No timeout",
+                "config": TrialConfig(),
+                "started_at": None,
+                "expected": None,
+            },
+            {
+                "name": "Positive timeout",
+                "config": TrialConfig(max_duration_seconds=10),
+                "started_at": None,
+                "expected": 10,
+            },
+            {
+                "name": "Zero timeout",
+                "config": TrialConfig(max_duration_seconds=0),
+                "started_at": None,
+                "expected": 0,
+            },
+            {
+                "name": "Negative timeout",
+                "config": TrialConfig(max_duration_seconds=-5),
+                "started_at": None,
+                "expected": None,
+            },
+            {
+                "name": "With started_at, positive timeout",
+                "config": TrialConfig(max_duration_seconds=5),
+                "started_at": (
+                    datetime.now(UTC) - timedelta(seconds=3)
+                ).isoformat(),
+                "expected": 2,
+            },
+        ]
+
+        for case in test_cases:
+            if case["started_at"]:
+                os.environ["ALPHATRION_TRIAL_START_TIME"] = case["started_at"]
+            with self.subTest(name=case["name"]):
+                trial = Trial(exp_id=1, config=case["config"])
+                self.assertEqual(trial._timeout(), case["expected"])