experimaestro 2.0.0a8__py3-none-any.whl → 2.0.0b4__py3-none-any.whl

experimaestro/tests/test_progress.py

@@ -90,54 +90,6 @@ def test_progress_basic():
                 assert info.progress == v
 
 
-def test_progress_multiple():
-    """Test that even with two schedulers, we get notified"""
-    max_wait = 5
-
-    with TemporaryExperiment(
-        "progress-progress-multiple-1", maxwait=max_wait, port=0
-    ) as xp1:
-        assert xp1.server is not None
-        assert xp1.server.port > 0
-
-        listener1 = ProgressListener()
-        xp1.scheduler.addlistener(listener1)
-
-        out = ProgressingTask.C().submit()
-        path = out.path  # type: Path
-        job = out.__xpm__.job
-
-        logger.info("Waiting for job to start (1)")
-        while job.state.notstarted():
-            time.sleep(1e-2)
-
-        with TemporaryExperiment(
-            "progress-progress-multiple-2",
-            workdir=xp1.workdir,
-            maxwait=max_wait,
-            port=0,
-        ) as xp2:
-            assert xp2.server is not None
-            assert xp2.server.port > 0
-            listener2 = ProgressListener()
-            xp2.scheduler.addlistener(listener2)
-
-            out = ProgressingTask.C().submit()
-            job = out.__xpm__.job  # type: CommandLineJob
-            logger.info("Waiting for job to start (2)")
-            while job.state.notstarted():
-                time.sleep(1e-2)
-
-            # Both schedulers should receive the job progress information
-            logger.info("Checking job progress")
-            progresses = [i / 10.0 for i in range(11)]
-            for v in progresses:
-                writeprogress(path, v)
-                if v < 1:
-                    assert listener1.progresses.get()[0].progress == v
-                    assert listener2.progresses.get()[0].progress == v
-
-
 NestedTasks = Tuple[str, Union[int, List["NestedTasks"]]]
 
 

experimaestro/tests/test_resumable_task.py

@@ -0,0 +1,480 @@
+"""Tests for ResumableTask with timeout retry logic"""
+
+import json
+from pathlib import Path
+from experimaestro import field, ResumableTask, Task, Param, GracefulTimeout
+from experimaestro.scheduler.workspace import RunMode
+from experimaestro.scheduler import JobState, JobFailureStatus
+from experimaestro.scheduler.jobs import JobStateError
+from experimaestro.scheduler.interfaces import JobState as JobStateClass
+from experimaestro.connectors import Process, ProcessState
+from experimaestro.connectors.local import LocalConnector
+from experimaestro.launchers.direct import DirectLauncher
+from experimaestro.commandline import CommandLineJob
+from .utils import TemporaryExperiment
+
+
+class MockTimeoutCommandLineJob(CommandLineJob):
+    """CommandLineJob that simulates timeouts based on attempt count"""
+
+    def __init__(self, *args, timeout_count=0, checkpoint_file=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.timeout_count = timeout_count
+        self.test_checkpoint_file = checkpoint_file
+
+    async def aio_run(self):
+        """Override to simulate timeout behavior"""
+        # Execute the actual task in a thread
+        from experimaestro.utils.asyncio import asyncThreadcheck
+
+        await asyncThreadcheck("execute", self.config.execute)
+
+        # Return a mock process that simulates timeout based on attempt count
+        return MockTimeoutProcess(self.test_checkpoint_file, self.timeout_count)
+
+    async def aio_process(self):
+        """No existing process"""
+        return None
+
+
+class MockTimeoutProcess(Process):
+    """Process that returns TIMEOUT for first N attempts"""
+
+    def __init__(self, checkpoint_file: Path, timeout_count: int):
+        self.checkpoint_file = checkpoint_file
+        self.timeout_count = timeout_count
+
+    def wait(self) -> int:
+        # Always return success (code 0) - timeout detection is in get_job_state
+        return 0
+
+    async def aio_state(self, timeout: float | None = None) -> ProcessState:
+        return ProcessState.DONE
+
+    def get_job_state(self, code: int) -> "JobState":
+        """Return TIMEOUT for first timeout_count attempts"""
+        # Read attempt count from checkpoint
+        attempt = 1
+        if self.checkpoint_file.exists():
+            attempt = int(self.checkpoint_file.read_text())
+
+        # Return TIMEOUT for first timeout_count attempts
+        if attempt <= self.timeout_count:
+            return JobStateError(JobFailureStatus.TIMEOUT)
+
+        return JobState.DONE
+
+
+class MockTimeoutLauncher(DirectLauncher):
+    """Launcher that creates jobs simulating timeouts"""
+
+    def __init__(self, timeout_count: int, checkpoint_file: Path):
+        super().__init__(LocalConnector())
+        self.timeout_count = timeout_count
+        self.checkpoint_file = checkpoint_file
+
+
+# Monkey-patch the task type to use our mock job
+# This is done by overriding the task factory
+def create_mock_timeout_task(timeout_count: int, checkpoint_file: Path):
+    """Create a task type that uses MockTimeoutCommandLineJob"""
+
+    def mock_job_factory(commandline):
+        class MockCommandLineTask:
+            def __init__(self, commandline):
+                self.commandline = commandline
+
+            def __call__(
+                self,
+                pyobject,
+                *,
+                launcher=None,
+                workspace=None,
+                run_mode=None,
+                max_retries=None,
+            ):
+                return MockTimeoutCommandLineJob(
+                    self.commandline,
+                    pyobject,
+                    launcher=launcher,
+                    workspace=workspace,
+                    run_mode=run_mode,
+                    max_retries=max_retries,
+                    timeout_count=timeout_count,
+                    checkpoint_file=checkpoint_file,
+                )
+
+        return MockCommandLineTask(commandline)
+
+    return mock_job_factory
+
+
+class CountingResumableTask(ResumableTask):
+    """Resumable task that counts execution attempts"""
+
+    checkpoint: Param[Path]
+
+    def execute(self):
+        # Count attempts in checkpoint file
+        attempt = 1
+        if self.checkpoint.exists():
+            attempt = int(self.checkpoint.read_text()) + 1
+
+        self.checkpoint.write_text(str(attempt))
+
+
+class SimpleResumableTask(ResumableTask):
+    """Simple resumable task for testing"""
+
+    def execute(self):
+        # This would normally contain checkpoint logic
+        pass
+
+
+class SimpleNonResumableTask(Task):
+    """Simple non-resumable task for testing"""
+
+    def execute(self):
+        pass
+
+
+def test_resumable_task_has_resumable_flag():
+    """Test that ResumableTask instances are correctly identified"""
+    with TemporaryExperiment("resumable_flag", maxwait=0):
+        launcher = DirectLauncher(LocalConnector())
+
+        # Submit resumable task
+        resumable = SimpleResumableTask.C().submit(
+            launcher=launcher, run_mode=RunMode.DRY_RUN
+        )
+        assert resumable.__xpm__.job.resumable is True
+
+        # Submit non-resumable task
+        non_resumable = SimpleNonResumableTask.C().submit(
+            launcher=launcher, run_mode=RunMode.DRY_RUN
+        )
+        assert non_resumable.__xpm__.job.resumable is False
+
+
+def test_max_retries_default():
+    """Test that default max_retries is 3"""
+    with TemporaryExperiment("max_retries_default", maxwait=0):
+        launcher = DirectLauncher(LocalConnector())
+
+        task = SimpleResumableTask.C().submit(
+            launcher=launcher, run_mode=RunMode.DRY_RUN
+        )
+
+        # Default should be 3 (from workspace settings)
+        assert task.__xpm__.job.max_retries == 3
+        assert task.__xpm__.job.retry_count == 0
+
+
+def test_max_retries_custom():
+    """Test that custom max_retries parameter is respected"""
+    with TemporaryExperiment("max_retries_custom", maxwait=0):
+        launcher = DirectLauncher(LocalConnector())
+
+        task = SimpleResumableTask.C().submit(
+            launcher=launcher, run_mode=RunMode.DRY_RUN, max_retries=5
+        )
+
+        assert task.__xpm__.job.max_retries == 5
+        assert task.__xpm__.job.retry_count == 0
+
+
+def test_max_retries_zero():
+    """Test that max_retries=0 is allowed (no retries)"""
+    with TemporaryExperiment("max_retries_zero", maxwait=0):
+        launcher = DirectLauncher(LocalConnector())
+
+        task = SimpleResumableTask.C().submit(
+            launcher=launcher, run_mode=RunMode.DRY_RUN, max_retries=0
+        )
+
+        assert task.__xpm__.job.max_retries == 0
+        assert task.__xpm__.job.retry_count == 0
+
+
+def test_resumable_task_succeeds_after_timeouts():
+    """Test that a resumable task retries and succeeds after timeouts"""
+    with TemporaryExperiment("resumable_timeout_success", maxwait=20) as xp:
+        checkpoint_file = xp.workspace.path / "checkpoint.txt"
+        launcher = DirectLauncher(LocalConnector())
+
+        # Create task config
+        task_config = CountingResumableTask.C(checkpoint=checkpoint_file)
+
+        # Create mock job directly
+        job = MockTimeoutCommandLineJob(
+            None,  # commandline not used
+            task_config,
+            workspace=xp.workspace,
+            launcher=launcher,
+            max_retries=5,
+            timeout_count=2,
+            checkpoint_file=checkpoint_file,
+        )
+
+        # Set the job on the config
+        task_config.__xpm__.job = job
+
+        # Submit to scheduler
+        from experimaestro.scheduler import experiment
+
+        experiment.CURRENT.submit(job)
+
+        # Should succeed after 3 attempts (2 timeouts + 1 success)
+        state = job.wait()
+        assert state == JobState.DONE
+        assert job.retry_count == 2
+        # Checkpoint should show 3 executions
+        assert int(checkpoint_file.read_text()) == 3
+
+
+def test_resumable_task_fails_after_max_retries():
+    """Test that a resumable task fails after exceeding max_retries"""
+    from experimaestro.scheduler import FailedExperiment
+    import pytest
+
+    with pytest.raises(FailedExperiment):
+        with TemporaryExperiment("resumable_timeout_fail", maxwait=20) as xp:
+            checkpoint_file = xp.workspace.path / "checkpoint.txt"
+            launcher = DirectLauncher(LocalConnector())
+
+            # Create task config
+            task_config = CountingResumableTask.C(checkpoint=checkpoint_file)
+
+            # Create mock job directly
+            job = MockTimeoutCommandLineJob(
+                None,  # commandline not used
+                task_config,
+                workspace=xp.workspace,
+                launcher=launcher,
+                max_retries=3,
+                timeout_count=10,
+                checkpoint_file=checkpoint_file,
+            )
+
+            # Set the job on the config
+            task_config.__xpm__.job = job
+
+            # Submit to scheduler
+            from experimaestro.scheduler import experiment
+
+            experiment.CURRENT.submit(job)
+
+            # Wait for job to complete (will fail)
+            state = job.wait()
+
+    # Verify job failed correctly
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.TIMEOUT
+    # retry_count should be max_retries + 1 (initial attempt + 3 retries = 4 total)
+    assert job.retry_count == 4
+    # Checkpoint should show 4 executions
+    assert int(checkpoint_file.read_text()) == 4
+
+
+# =============================================================================
+# Tests for JobState.from_path and .failed file format
+# =============================================================================
+
+
+def test_job_state_from_path_json_timeout(tmp_path):
+    """Test JobState.from_path reads JSON format with timeout reason"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text(
+        json.dumps({"code": 1, "reason": "timeout", "message": "Graceful"})
+    )
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.TIMEOUT
+
+
+def test_job_state_from_path_json_memory(tmp_path):
+    """Test JobState.from_path reads JSON format with memory reason"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text(json.dumps({"code": 1, "reason": "memory"}))
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.MEMORY
+
+
+def test_job_state_from_path_json_dependency(tmp_path):
+    """Test JobState.from_path reads JSON format with dependency reason"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text(json.dumps({"code": 1, "reason": "dependency"}))
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.DEPENDENCY
+
+
+def test_job_state_from_path_json_failed(tmp_path):
+    """Test JobState.from_path reads JSON format with failed reason"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text(json.dumps({"code": 1, "reason": "failed"}))
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.FAILED
+
+
+def test_job_state_from_path_json_unknown_reason(tmp_path):
+    """Test JobState.from_path handles unknown reason gracefully"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text(json.dumps({"code": 1, "reason": "unknown_reason"}))
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.FAILED  # Falls back to FAILED
+
+
+def test_job_state_from_path_legacy_integer_nonzero(tmp_path):
+    """Test JobState.from_path reads legacy integer format (non-zero = error)"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text("1")
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert isinstance(state, JobStateError)
+    assert state.failure_reason == JobFailureStatus.FAILED
+
+
+def test_job_state_from_path_legacy_integer_zero(tmp_path):
+    """Test JobState.from_path reads legacy integer format (zero = done)"""
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text("0")
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert state == JobState.DONE
+
+
+def test_job_state_from_path_done_file(tmp_path):
+    """Test JobState.from_path reads .done file"""
+    done_file = tmp_path / "test.done"
+    done_file.touch()
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert state == JobState.DONE
+
+
+def test_job_state_from_path_no_file(tmp_path):
+    """Test JobState.from_path returns None when no file exists"""
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert state is None
+
+
+def test_job_state_from_path_done_takes_precedence(tmp_path):
+    """Test JobState.from_path prefers .done over .failed"""
+    done_file = tmp_path / "test.done"
+    done_file.touch()
+    failed_file = tmp_path / "test.failed"
+    failed_file.write_text(json.dumps({"code": 1, "reason": "failed"}))
+
+    state = JobStateClass.from_path(tmp_path, "test")
+    assert state == JobState.DONE
+
+
+# =============================================================================
+# Tests for GracefulTimeout exception
+# =============================================================================
+
+
+class GracefulTimeoutTask(ResumableTask):
+    """Task that raises GracefulTimeout"""
+
+    checkpoint: Param[Path]
+    should_timeout: Param[bool] = field(ignore_default=True)
+
+    def execute(self):
+        # Count attempts in checkpoint file
+        attempt = 1
+        if self.checkpoint.exists():
+            attempt = int(self.checkpoint.read_text()) + 1
+        self.checkpoint.write_text(str(attempt))
+
+        # Raise GracefulTimeout on first attempt if should_timeout is True
+        if self.should_timeout and attempt == 1:
+            raise GracefulTimeout("Not enough time for another epoch")
+
+
+# =============================================================================
+# Tests for remaining_time() method with mock launcher
+# =============================================================================
+
+
+class MockLauncherWithRemainingTime(DirectLauncher):
+    """Mock launcher that provides remaining_time via launcher_info_code()"""
+
+    def __init__(self, remaining_time_value: float | None):
+        super().__init__(LocalConnector())
+        self.remaining_time_value = remaining_time_value
+
+    def launcher_info_code(self) -> str:
+        """Generate code to set up MockLauncherInformation with the specified value"""
+        if self.remaining_time_value is None:
+            return (
+                "    from experimaestro.tests.test_resumable_task import MockLauncherInformation\n"
+                "    from experimaestro import taskglobals\n"
+                "    taskglobals.Env.instance().launcher_info = MockLauncherInformation(None)\n"
+            )
+        return (
+            "    from experimaestro.tests.test_resumable_task import MockLauncherInformation\n"
+            "    from experimaestro import taskglobals\n"
+            f"    taskglobals.Env.instance().launcher_info = MockLauncherInformation({self.remaining_time_value})\n"
+        )
+
+
+class MockLauncherInformation:
+    """Mock launcher info for testing remaining_time()"""
+
+    def __init__(self, remaining: float | None):
+        self._remaining = remaining
+
+    def remaining_time(self) -> float | None:
+        return self._remaining
+
+
+class RemainingTimeTask(ResumableTask):
+    """Task that records the remaining_time() value"""
+
+    output_file: Param[Path]
+
+    def execute(self):
+        remaining = self.remaining_time()
+        self.output_file.write_text(str(remaining) if remaining is not None else "None")
+
+
+def test_remaining_time_with_mock_launcher():
+    """Test remaining_time() works with a mock launcher that provides launcher_info_code()"""
+    with TemporaryExperiment("remaining_time", maxwait=10) as xp:
+        output_file = xp.workspace.path / "remaining.txt"
+        launcher = MockLauncherWithRemainingTime(remaining_time_value=1234.5)
+
+        task = RemainingTimeTask.C(output_file=output_file).submit(launcher=launcher)
+
+        state = task.__xpm__.job.wait()
+        assert state == JobState.DONE
+
+        # Verify the task received the remaining time value
+        assert output_file.exists()
+        assert output_file.read_text() == "1234.5"
+
+
+def test_remaining_time_none_with_mock_launcher():
+    """Test remaining_time() returns None when launcher has no time limit"""
+    with TemporaryExperiment("remaining_time_none", maxwait=10) as xp:
+        output_file = xp.workspace.path / "remaining.txt"
+        launcher = MockLauncherWithRemainingTime(remaining_time_value=None)
+
+        task = RemainingTimeTask.C(output_file=output_file).submit(launcher=launcher)
+
+        state = task.__xpm__.job.wait()
+        assert state == JobState.DONE
+
+        # Verify the task received None
+        assert output_file.exists()
+        assert output_file.read_text() == "None"

experimaestro/tests/test_serializers.py

@@ -2,11 +2,12 @@ from typing import Optional
 from experimaestro import (
     Config,
     Param,
+    Task,
     state_dict,
     from_state_dict,
 )
 from experimaestro.core.context import SerializationContext
-from experimaestro.core.objects import ConfigMixin
+from experimaestro.core.objects import ConfigMixin, ConfigInformation, TaskStub
 
 
 class Object1(Config):
@@ -52,3 +53,142 @@ def test_serializers_types():
     config = MultiParamObject.C(x={"a": None})
     config.__xpm__.seal(context)
     state_dict(context, config)
+
+
+# --- Tests for partial_loading feature ---
+
+
+class MyTask(Task):
+    """A task for testing partial_loading"""
+
+    value: Param[int]
+
+    def execute(self):
+        pass
+
+
+class ConfigWithTask(Config):
+    """A config that references a task"""
+
+    name: Param[str]
+
+
+class TaskDependentConfig(Config):
+    """A config that is only used by a task"""
+
+    data: Param[str]
+
+
+class TaskWithDependency(Task):
+    """A task that uses another config"""
+
+    dep: Param[TaskDependentConfig]
+
+    def execute(self):
+        pass
+
+
+def test_partial_loading_skips_task_reference():
+    """Test that partial_loading skips loading task references"""
+    context = SerializationContext(save_directory=None)
+
+    # Create a config with a task
+    task = MyTask.C(value=42)
+    config = ConfigWithTask.C(name="test")
+    config.__xpm__.task = task
+
+    # Seal both
+    task.__xpm__.seal(context)
+    config.__xpm__.seal(context)
+
+    # Serialize
+    data = state_dict(context, [task, config])
+
+    # Load without partial_loading - task should be loaded
+    [loaded_task, loaded_config] = from_state_dict(data, partial_loading=False)
+    assert isinstance(loaded_config, ConfigWithTask)
+    assert isinstance(loaded_config.__xpm__.task, MyTask)
+
+    # Load with partial_loading - task should be a stub
+    [loaded_task2, loaded_config2] = from_state_dict(data, partial_loading=True)
+    assert isinstance(loaded_config2, ConfigWithTask)
+    assert isinstance(loaded_config2.__xpm__.task, TaskStub)
+    # Check that typename contains the task class name
+    assert "MyTask" in loaded_config2.__xpm__.task.typename
+
+
+def test_partial_loading_skips_task_dependencies():
+    """Test that partial_loading skips configs only used by tasks"""
+    context = SerializationContext(save_directory=None)
+
+    # Create a config that is only used by a task
+    dep = TaskDependentConfig.C(data="some data")
+    task = TaskWithDependency.C(dep=dep)
+    config = ConfigWithTask.C(name="main")
+    config.__xpm__.task = task
+
+    # Seal all
+    dep.__xpm__.seal(context)
+    task.__xpm__.seal(context)
+    config.__xpm__.seal(context)
+
+    # Serialize
+    data = state_dict(context, [dep, task, config])
+    definitions = data["objects"]
+
+    # Load with partial_loading - both task and its dependency should be stubs
+    objects = ConfigInformation.load_objects(
+        definitions, as_instance=False, partial_loading=True
+    )
+
+    # The main config should be loaded
+    main_obj = objects[definitions[-1]["id"]]
+    assert isinstance(main_obj, ConfigWithTask)
+
+    # The task should be a stub
+    task_obj = objects[definitions[1]["id"]]
+    assert isinstance(task_obj, TaskStub)
+
+    # The task dependency should also be a stub
+    dep_obj = objects[definitions[0]["id"]]
+    assert isinstance(dep_obj, TaskStub)
+
+
+def test_partial_loading_preserves_shared_configs():
+    """Test that configs used by both main object and task are not skipped"""
+    context = SerializationContext(save_directory=None)
+
+    # Create a shared config
+    shared = Object1.C()
+
+    # Create task and main config that both use the shared config
+    task = MyTask.C(value=1)
+    main = Object2.C(object=shared)
+    main.__xpm__.task = task
+
+    # Seal all
+    shared.__xpm__.seal(context)
+    task.__xpm__.seal(context)
+    main.__xpm__.seal(context)
+
+    # Serialize
+    data = state_dict(context, [shared, task, main])
+    definitions = data["objects"]
+
+    # Load with partial_loading
+    objects = ConfigInformation.load_objects(
+        definitions, as_instance=False, partial_loading=True
+    )
+
+    # The shared config should be loaded (not a stub) since main uses it
+    shared_obj = objects[definitions[0]["id"]]
+    assert isinstance(shared_obj, Object1)
+    assert not isinstance(shared_obj, TaskStub)
+
+    # The main config should be loaded
+    main_obj = objects[definitions[-1]["id"]]
+    assert isinstance(main_obj, Object2)
+
+    # The task should be a stub
+    task_obj = objects[definitions[1]["id"]]
+    assert isinstance(task_obj, TaskStub)