experimaestro 1.11.1__py3-none-any.whl → 2.0.0b4__py3-none-any.whl

experimaestro/tests/tasks/test_dynamic.py

@@ -0,0 +1,231 @@
+# Test for future task outputs handling
+# https://github.com/experimaestro/experimaestro-python/issues/90
+
+from functools import partial
+import json
+import logging
+from pathlib import Path
+import sys
+import time
+from experimaestro import (
+    Config,
+    Param,
+    Task,
+    ResumableTask,
+    DependentMarker,
+    LightweightTask,
+    field,
+    PathGenerator,
+)
+from experimaestro.core.arguments import Meta
+from experimaestro.tests.utils import TemporaryDirectory, TemporaryExperiment
+
+
+class Model(Config):
+    pass
+
+
+class Checkpoint(Config):
+    step: Param[int]
+    model: Param[Model]
+
+
+class CheckpointLoader(LightweightTask):
+    checkpoint: Param[Checkpoint]
+
+    def execute(self):
+        pass
+
+
+class Evaluate(Task):
+    model: Param[Model]
+
+    def execute(self):
+        pass
+
+
+class Validation(Config):
+    model: Param[Model]
+
+    def checkpoint(self, dep: DependentMarker, *, step: int) -> Checkpoint:
+        return dep(Checkpoint.C(model=self.model, step=step))
+
+    def compute(self, step: int):
+        self.register_task_output(self.checkpoint, step=step)
+
+
+class Learn(ResumableTask):
+    model: Param[Model]
+    validation: Param[Validation]
+
+    # Control files for synchronization with tests
+    max_step_file: Meta[Path] = field(default_factory=PathGenerator("max_step"))
+    state_file: Meta[Path] = field(default_factory=PathGenerator("state.json"))
+
+    def execute(self):
+        start_step = 0
+
+        if self.state_file.exists():
+            with self.state_file.open("r") as f:
+                state = json.load(f)
+                start_step = state.get("last_step", 0)
+                logging.info("Resuming from step %d", start_step)
+
+        # Wait for max_step_file to know how far to go
+        while not self.max_step_file.is_file():
+            time.sleep(0.1)
+
+        with self.max_step_file.open("r") as f:
+            max_step = int(f.read().strip())
+        self.max_step_file.unlink()
+
+        # Use absolute value for step comparison
+        # Negative max_step means: produce up to |max_step| then crash (simulate interruption)
+        # Positive max_step means: produce up to max_step then complete normally
+        abs_max = abs(max_step)
+
+        for step in [15, 30, 45]:
+            if step <= start_step:
+                logging.info("Skipping already processed step %d", step)
+                continue
+
+            if step > abs_max:
+                # We're past the limit, stop here
+                break
+
+            self.validation.compute(step)
+
+            # Save state after each checkpoint
+            with self.state_file.open("w") as f:
+                json.dump({"last_step": step}, f)
+
+            # If max_step is negative (e.g. -15), simulate exit after producing |max_step|
+            if max_step < 0 and step >= abs_max:
+                logging.warning("Simulating interruption after step %d", step)
+                sys.exit(1)
+
+
+def evaluate(evaluations, checkpoint: Checkpoint):
+    logging.warning("Evaluating checkpoint %s", checkpoint)
+    task = Evaluate.C(model=checkpoint.model)
+    checkpoint_loader = CheckpointLoader.C(checkpoint=checkpoint)
+    evaluations.append(task.submit(init_tasks=[checkpoint_loader]))
+
+
+def test_task_dynamic_simple():
+    """Test that dynamic task outputs trigger callbacks
+
+    This test verifies that callbacks are guaranteed to complete before
+    the experiment context exits. The callback waits for jobs to complete
+    before submitting evaluations, which validates that the synchronization
+    logic correctly waits for all callbacks to finish.
+    """
+    import asyncio
+
+    evaluations = []
+    xp_ref = [None]  # To access xp from callback
+
+    def collect_checkpoint(checkpoint: Checkpoint):
+        """Callback that waits for jobs to complete before evaluating
+
+        This simulates a real-world scenario where the callback needs to wait
+        for the triggering task to complete before it can proceed (e.g., to
+        read outputs from the task's directory).
+        """
+        logging.info("Received checkpoint %s, waiting for jobs to complete", checkpoint)
+        xp = xp_ref[0]
+
+        # Wait for unfinished jobs to become 0 (all tasks completed)
+        async def wait_for_jobs_done():
+            async with xp.scheduler.exitCondition:
+                while xp.unfinishedJobs > 0:
+                    await xp.scheduler.exitCondition.wait()
+
+        asyncio.run_coroutine_threadsafe(
+            wait_for_jobs_done(), xp.scheduler.loop
+        ).result()
+
+        # Now submit evaluation
+        logging.info("Jobs done, submitting evaluation for checkpoint %s", checkpoint)
+        evaluate(evaluations, checkpoint)
+
+    with TemporaryDirectory() as workdir:
+        with TemporaryExperiment("dynamic", maxwait=10, workdir=workdir) as xp:
+            xp_ref[0] = xp
+            model = Model.C()
+            validation = Validation.C(model=model)
+            learn = Learn.C(model=model, validation=validation)
+            learn.watch_output(validation.checkpoint, collect_checkpoint)
+
+            learn.submit()
+
+            # Allow the task to run up to step 30
+            learn.max_step_file.parent.mkdir(parents=True, exist_ok=True)
+            with learn.max_step_file.open("w") as f:
+                f.write("30")
+
+            logging.info("Experiment will wait for completion...")
+
+        assert len(evaluations) == 2, f"Expected 2 evaluations, got {len(evaluations)}"
+
+
+def test_task_dynamic_replay():
+    """Test that dynamic outputs are replayed when a task is restarted
+
+    Scenario:
+    1. First run: task produces checkpoint for step 15, then exits (simulated timeout)
+    2. Second run: task should replay the step 15 checkpoint and produce new ones
+    """
+    with TemporaryDirectory() as workdir:
+        # First run: produce one checkpoint then exit
+        evaluations_run1 = []
+        try:
+            with TemporaryExperiment("dynamic_replay", maxwait=5, workdir=workdir):
+                model = Model.C()
+                validation = Validation.C(model=model)
+                learn = Learn.C(model=model, validation=validation)
+                learn.watch_output(
+                    validation.checkpoint, partial(evaluate, evaluations_run1)
+                )
+
+                learn.submit()
+
+                # Allow task to produce step 15 checkpoint, then simulate crash
+                # Negative value means: produce up to |value| then exit with error
+                learn.max_step_file.parent.mkdir(parents=True, exist_ok=True)
+                with learn.max_step_file.open("w") as f:
+                    f.write("-15")
+
+        except Exception as e:
+            # Expected: the task will fail when trying to go past max_step
+            logging.info("First run ended (expected): %s", e)
+
+        # First run should have produced at least one evaluation (for step 15)
+        assert (
+            len(evaluations_run1) == 1
+        ), f"Run 1: Expected 1 evaluation, got {len(evaluations_run1)}"
+
+        # Second run: restart and continue
+        evaluations_run2 = []
+        with TemporaryExperiment("dynamic_replay", maxwait=30, workdir=workdir):
+            model = Model.C()
+            validation = Validation.C(model=model)
+            learn = Learn.C(model=model, validation=validation)
+            learn.watch_output(
+                validation.checkpoint, partial(evaluate, evaluations_run2)
+            )
+
+            learn.submit()
+
+            # Allow task to run to completion (step 45)
+            learn.max_step_file.parent.mkdir(parents=True, exist_ok=True)
+            with learn.max_step_file.open("w") as f:
+                f.write("45")
+
+        # Second run should have:
+        # - Replayed the step 15 checkpoint (from first run)
+        # - Produced step 30 and 45 checkpoints
+        # Total: 3 evaluations (but step 15 was replayed, not re-produced)
+        assert (
+            len(evaluations_run2) == 3
+        ), f"Run 2: Expected 3 evaluations, got {len(evaluations_run2)}"

experimaestro/tests/test_checkers.py

@@ -10,7 +10,7 @@ def test_choices():
     class TestChoices(Config):
         a: Annotated[str, Choices(["a", "b"])]
 
-    TestChoices(a="a").__xpm__.validate()
+    TestChoices.C(a="a").__xpm__.validate()
 
     with pytest.raises((ValueError, KeyError)):
-        TestChoices(a="c").__xpm__.validate()
+        TestChoices.C(a="c").__xpm__.validate()