Prevent training service tests from leaking running jobs

Author: bghira

simpletuner/simpletuner_sdk/server/services/local_gpu_allocator.py

@@ -53,7 +53,7 @@ def __init__(self):
         self._job_repo = None
         self._reconciled = False
 
-    async def reconcile_on_startup(self, max_entries: int = 20) -> Dict[str, int]:
+    async def reconcile_on_startup(self, max_entries: Optional[int] = None) -> Dict[str, int]:
         """Reconcile LOCAL jobs with actual process status.
 
         Called on startup to check running LOCAL jobs and handle orphaned ones:
@@ -62,7 +62,8 @@ async def reconcile_on_startup(self, max_entries: int = 20) -> Dict[str, int]:
         - Jobs without PIDs (legacy) are marked as failed
 
         Args:
-            max_entries: Maximum number of running entries to check (default 20).
+            max_entries: Optional maximum number of running entries to check.
+                When omitted, all running local jobs are reconciled.
 
         Returns:
             Dictionary with stats: {"orphaned": N, "adopted": N, "no_pid": N}
@@ -73,17 +74,14 @@ async def reconcile_on_startup(self, max_entries: int = 20) -> Dict[str, int]:
             return stats
 
         job_repo = self._get_job_repo()
-        # Only get LOCAL running jobs - this is bounded by the number of
-        # GPUs on this machine, so should be small (typically 0-8)
         running_jobs = await job_repo.get_running_local_jobs()
 
         if not running_jobs:
             self._reconciled = True
             return stats
 
-        # Limit entries to check
-        jobs_to_check = running_jobs[:max_entries]
-        if len(running_jobs) > max_entries:
+        jobs_to_check = running_jobs if max_entries is None else running_jobs[:max_entries]
+        if max_entries is not None and len(running_jobs) > max_entries:
             logger.warning(
                 "Found %d running local jobs, only checking first %d",
                 len(running_jobs),

tests/test_local_gpu_allocator.py

@@ -204,6 +204,27 @@ def test_reconcile_marks_preboot_job_failed(self):
         self.assertEqual(repo.failed, [("job-preboot", "Process ended after system reboot")])
         self.assertEqual(repo.released, ["job-preboot"])
 
+    def test_reconcile_processes_all_running_jobs_by_default(self):
+        allocator = LocalGPUAllocator()
+        jobs = [
+            _StubJob(
+                f"job-{index}",
+                started_at="2026-01-27T00:00:00+00:00",
+            )
+            for index in range(35)
+        ]
+        repo = _StubJobRepo(jobs)
+        allocator._job_repo = repo
+
+        with patch.object(allocator, "_get_boot_time_utc", return_value=None):
+            stats = asyncio.run(allocator.reconcile_on_startup())
+
+        self.assertEqual(stats["orphaned"], 0)
+        self.assertEqual(stats["adopted"], 0)
+        self.assertEqual(stats["no_pid"], 35)
+        self.assertEqual(len(repo.failed), 35)
+        self.assertEqual(len(repo.released), 35)
+
 
 class TestLocalGPUAllocatorAvailability(unittest.TestCase):
     """Tests for GPU availability checking."""

tests/test_training_service.py

@@ -65,6 +65,15 @@ def load_config(self, name: str):
         return dict(self._config), {}
 
 
+class DummyJobRepository:
+    def __init__(self):
+        self.jobs = []
+
+    async def add(self, job):
+        self.jobs.append(job)
+        return job
+
+
 def _mock_is_truthy(value: Any) -> bool:
     if value is None:
         return False
@@ -79,8 +88,14 @@ def _mock_is_truthy(value: Any) -> bool:
 class TrainingServiceTests(unittest.TestCase):
     def setUp(self) -> None:
         self._saved_state = copy.deepcopy(training_service.APIState.state)
+        self.job_repo = DummyJobRepository()
         self._save_state_patch = patch.object(training_service.APIState, "save_state", return_value=None)
         self._save_state_patch.start()
+        self._job_repo_patch = patch(
+            "simpletuner.simpletuner_sdk.server.services.cloud.storage.job_repository.get_job_repository",
+            return_value=self.job_repo,
+        )
+        self._job_repo_patch.start()
         # Mock cache/scan service checks so stale singleton state doesn't block training
         mock_cache_svc = MagicMock()
         mock_cache_svc.get_active_status.return_value = None
@@ -99,6 +114,7 @@ def setUp(self) -> None:
 
     def tearDown(self) -> None:
         training_service.APIState.state = self._saved_state
+        self._job_repo_patch.stop()
         self._save_state_patch.stop()
         self._cache_svc_patch.stop()
         self._scan_svc_patch.stop()