fix: persist desktop training state so a restart survives the live run

src/praisonai-desktop/engine/test_training.py

@@ -590,6 +590,117 @@ def test_a_run_not_stopped_in_that_window_still_runs_normally(self):
         self.assertEqual(run.state, training.DONE)
 
 
+class SurvivingARestart(unittest.TestCase):
+    """State must outlive the engine process, or a restart lies.
+
+    The engine kept the live run and the history in memory only, so a crash,
+    relaunch or kill left `/train/runs` empty, `stop` with nothing to stop, and
+    `start` willing to launch a second trainer beside a live one -- the OOM
+    that "one GPU runs one job" exists to refuse. A second Trainer over the
+    same home stands in for the process that comes back after the restart.
+    """
+
+    def setUp(self):
+        self.home = tempfile.mkdtemp(prefix="praison-restart-")
+        self.first = training.Trainer(self.home, sys.executable)
+        self.config = {"model_name": "unsloth/tiny", "dataset": "d.json"}
+
+    def tearDown(self):
+        self.first.stop()
+        shutil.rmtree(self.home, ignore_errors=True)
+
+    def _start_a_live_run(self):
+        self.first.command_builder = _script(
+            "import time\nprint('up', flush=True)\ntime.sleep(120)")
+        run = self.first.start(self.config, "run-live")
+        self.assertTrue(_wait(lambda: run.state == training.RUNNING), run.state)
+        return run
+
+    def test_a_live_run_reappears_after_a_restart(self):
+        run = self._start_a_live_run()
+        second = training.Trainer(self.home, sys.executable)
+        self.assertIn("run-live", [r.id for r in second.history],
+                      "the run vanished from history after the restart")
+        adopted = second.get("run-live")
+        self.assertEqual(adopted.state, training.RUNNING,
+                         "a run still on the GPU was not shown as running")
+
+    def test_a_second_finetune_is_refused_after_a_restart(self):
+        self._start_a_live_run()
+        second = training.Trainer(self.home, sys.executable)
+        second.command_builder = _script("pass")
+        with self.assertRaises(RuntimeError):
+            second.start(self.config, "run-two")
+
+    def test_stop_after_a_restart_actually_kills_the_live_pid(self):
+        run = self._start_a_live_run()
+        pid = run.pid
+        self.assertTrue(_alive(pid))
+        second = training.Trainer(self.home, sys.executable)
+        self.assertTrue(second.stop(), "stop() found nothing to stop after restart")
+        self.assertTrue(_wait(lambda: not _alive(pid), 15),
+                        "the adopted run's pid outlived stop()")
+        self.assertEqual(second.get("run-live").state, training.CANCELLED)
+
+    def test_an_interrupted_run_is_reported_failed_not_missing(self):
+        # No live process: a run.json left as RUNNING by a killed engine whose
+        # pid is now gone must read as failed, not adopted.
+        run_dir = os.path.join(self.home, "runs", "run-dead")
+        os.makedirs(run_dir, exist_ok=True)
+        import json
+        with open(os.path.join(run_dir, "run.json"), "w") as fh:
+            json.dump({"id": "run-dead", "state": training.RUNNING,
+                       "started": time.time(), "pid": 2 ** 31 - 1}, fh)
+        second = training.Trainer(self.home, sys.executable)
+        dead = second.get("run-dead")
+        self.assertIsNotNone(dead, "the interrupted run was dropped entirely")
+        self.assertEqual(dead.state, training.FAILED)
+        self.assertIsNone(second.current, "a dead run was adopted as live")
+
+    def test_a_finished_run_is_not_adopted_after_a_restart(self):
+        self.first.command_builder = _script("print('done')")
+        run = self.first.start(self.config, "run-done")
+        self.assertTrue(_wait(lambda: run.state in training.TERMINAL), run.state)
+        second = training.Trainer(self.home, sys.executable)
+        self.assertEqual(second.get("run-done").state, training.DONE)
+        self.assertIsNone(second.current, "a finished run was adopted as live")
+        second.command_builder = _script("pass")
+        again = second.start(self.config, "run-after")
+        self.assertTrue(_wait(lambda: again.state in training.TERMINAL), again.state)
+
+    def test_a_half_written_state_file_never_reaches_the_reader(self):
+        # _persist writes a temp file and os.replaces it, so a reader only ever
+        # sees the old whole file or the new whole file -- never the truncated
+        # middle a plain open("w") would leave if the engine were killed mid
+        # write. Were the write non-atomic, an interrupted _reload would skip
+        # the record, the live run would vanish, and stop() and the single-GPU
+        # guard would both lose it.
+        run = self._start_a_live_run()
+        state = pathlib.Path(self.home, "runs", "run-live", "run.json")
+        self.assertTrue(state.exists(), "the live run was never persisted")
+        import json
+        json.loads(state.read_text())          # complete JSON, not a fragment
+        siblings = list(state.parent.glob("run.json.*.tmp"))
+        self.assertEqual(siblings, [], f"a temp file was left behind: {siblings}")
+
+    def test_a_run_persisted_before_spawn_is_not_lost_on_an_early_restart(self):
+        # start() persists the run before the supervisor thread exists, so an
+        # engine that dies in that window still leaves a record. Simulate it: a
+        # pending run.json with no pid must reload as failed and not be adopted,
+        # rather than vanishing and letting a second trainer start.
+        run_dir = os.path.join(self.home, "runs", "run-early")
+        os.makedirs(run_dir, exist_ok=True)
+        import json
+        with open(os.path.join(run_dir, "run.json"), "w") as fh:
+            json.dump({"id": "run-early", "state": training.PENDING,
+                       "started": time.time(), "pid": None}, fh)
+        second = training.Trainer(self.home, sys.executable)
+        early = second.get("run-early")
+        self.assertIsNotNone(early, "a run persisted before spawn was dropped")
+        self.assertEqual(early.state, training.FAILED)
+        self.assertIsNone(second.current, "a pidless pending run was adopted")
+
+
 class RunLifecycle(unittest.TestCase):
     def setUp(self):
         self.home = tempfile.mkdtemp(prefix="praison-train-test-")

src/praisonai-desktop/engine/training.py

@@ -158,6 +158,11 @@ def __init__(self, run_id, config_path, log_path):
         self.events = []             # [(cursor, kind, payload)]
         self._next_cursor = 0        # never derived from len(events): see emit
         self._proc = None
+        # The child's pid, recorded on disk so a *later* engine process can
+        # find a run it never spawned: _proc is unpicklable and dies with us,
+        # but the pid outlives the restart and is how stop() reaches an
+        # adopted run.
+        self.pid = None
         self._lock = threading.Lock()
 
     # -- event history --------------------------------------------------
@@ -223,6 +228,86 @@ def __init__(self, home, python, command_builder=None):
         # without limit -- each retains its event ring and metric series.
         self.history = collections.deque(maxlen=MAX_HISTORY)
         self._lock = threading.Lock()
+        # Rebuild what an earlier engine process left on disk. Without this a
+        # restart during a run reports the run as never having happened: it
+        # vanishes from history, stop() has nothing to stop, and start()
+        # cheerfully launches a second trainer beside the live one -- the OOM
+        # that "one GPU runs one job" exists to refuse.
+        self._reload()
+
+    # -- surviving a restart --------------------------------------------
+    def _state_path(self, run_id):
+        return os.path.join(self.dir, run_id, "run.json")
+
+    def _persist(self, run):
+        """Enough on disk to answer 'is a job on the GPU right now'.
+
+        Only config.yaml and train.log were ever written; neither records the
+        state or the pid, so a new process could not tell a live run from a
+        finished one. Best-effort: a run that trains but cannot write its state
+        file is still better than one that refuses to start.
+
+        The write is atomic -- a temp file in the same directory, fsynced, then
+        os.replace. A plain open("w") truncates first, so an engine killed mid
+        json.dump leaves empty or half-written JSON; _reload then skips that
+        record and the live run disappears from history, becomes unreachable
+        through stop(), and no longer refuses a second trainer. os.replace is
+        atomic on POSIX and Windows, so a reader only ever sees the old file or
+        the whole new one.
+        """
+        path = self._state_path(run.id)
+        tmp = f"{path}.{os.getpid()}.tmp"
+        try:
+            with open(tmp, "w", encoding="utf-8") as fh:
+                json.dump({**run.summary(), "pid": run.pid}, fh)
+                fh.flush()
+                os.fsync(fh.fileno())
+            os.replace(tmp, path)
+        except OSError:
+            try:
+                os.remove(tmp)
+            except OSError:
+                pass
+
+    def _existing_ids(self):
+        """Run directories on disk, whether or not this process created them."""
+        try:
+            return [name for name in os.listdir(self.dir)
+                    if os.path.isdir(os.path.join(self.dir, name))]
+        except OSError:
+            return []
+
+    def _reload(self):
+        """Rebuild history, adopting a run whose process is still alive.
+
+        A run the new process cannot see is not 'missing' -- it is either still
+        on the GPU (adopt it, so stop() can reach it) or it was interrupted
+        (say so). Reporting neither is how a second fine-tune got started
+        beside a live one.
+        """
+        for name in sorted(self._existing_ids()):
+            try:
+                with open(self._state_path(name), encoding="utf-8") as fh:
+                    saved = json.load(fh)
+            except (OSError, ValueError):
+                continue
+            run_dir = os.path.join(self.dir, name)
+            run = Run(name, os.path.join(run_dir, "config.yaml"),
+                      os.path.join(run_dir, "train.log"))
+            run.state = saved.get("state", FAILED)
+            run.step, run.total = saved.get("step", 0), saved.get("total", 0)
+            run.started = saved.get("started", run.started)
+            run.ended, run.error = saved.get("ended"), saved.get("error")
+            run.pid = saved.get("pid")
+            if run.state not in TERMINAL:
+                if run.pid and _pid_alive(run.pid):
+                    run.state = RUNNING
+                    self.current = run          # stop() signals run.pid's group
+                else:
+                    run.finish(FAILED,
+                               "the engine restarted while this run was live")
+                    self._persist(run)
+            self.history.appendleft(run)
 
     # -- starting --------------------------------------------------------
     def start(self, config, run_id=None):
@@ -266,6 +351,14 @@ def start(self, config, run_id=None):
             # far end, which is the oldest run -- what the cap is for.
             self.history.appendleft(run)
 
+        # Record the run as pending *before* the supervisor thread exists.
+        # start() returns the moment the thread is created, so an engine that
+        # exits between here and the first post-spawn _persist would otherwise
+        # leave a child alive with no run.json: _reload skips it, stop() cannot
+        # reach it, and start() launches a second trainer beside it. A pending
+        # record with no pid is adopted as failed on the next boot, which is
+        # the honest reading of "spawned, engine died, pid unknown".
+        self._persist(run)
         run.emit("start", {"id": run.id, "config": config_path})
         threading.Thread(target=self._supervise, args=(run,), daemon=True).start()
         return run
@@ -333,8 +426,17 @@ def _supervise(self, run):
             proc = self._spawn(run)
         except Exception as exc:                       # noqa: BLE001
             run.finish(FAILED, f"could not start the trainer: {exc}")
+            # Persist the failure: start() left a pending record on disk, and
+            # without this a restart would read that stale pending run (no pid)
+            # as freshly interrupted rather than as the spawn failure it was.
+            self._persist(run)
             return
         run._proc = proc
+        run.pid = proc.pid
+        # Written now, before the RUNNING transition, so a crash in the next
+        # instruction still leaves a pid on disk for the next process to adopt
+        # or reap rather than a run that looks like it never spawned.
+        self._persist(run)
         # Honour a stop that arrived before the process existed.
         #
         # start() returns as soon as this thread is created, so there is a
@@ -351,6 +453,7 @@ def _supervise(self, run):
             _terminate_group(proc)
         else:
             run.emit("state", {"state": RUNNING})
+        self._persist(run)          # RUNNING (or STOPPING) is now on disk
         try:
             with open(run.log_path, "a", encoding="utf-8") as log:
                 for line in proc.stdout:
@@ -365,6 +468,9 @@ def _supervise(self, run):
             run.finish(DONE)
         else:
             run.finish(FAILED, _last_meaningful_line(run.log_path) or f"exit {code}")
+        # Record the ending, so a restart after the run finishes reads it as
+        # terminal rather than adopting a dead pid.
+        self._persist(run)
 
     def _consume(self, run, line):
         run.emit("log", {"line": line})
@@ -395,8 +501,20 @@ def stop(self, run_id=None):
         proc = run._proc
         if proc and proc.poll() is None:
             _terminate_group(proc)
-        # If proc is still None the run has not spawned yet; _supervise sees
-        # STOPPING under the same lock and terminates it on arrival.
+            # _supervise is reading proc.stdout; it sees STOPPING, calls
+            # finish(CANCELLED) and persists when the pipe closes.
+        elif proc is None and run.pid:
+            # An adopted run from a previous process: there is no _proc and no
+            # supervisor reading its output, so this process owns the ending.
+            # Signal the pid's group -- the group is what the original spawn
+            # created -- then record the ending ourselves.
+            _terminate_pid_group(run.pid)
+            run.finish(CANCELLED)
+            self._persist(run)
+        # If proc is still None *and* there is no pid the run has not spawned
+        # yet; _supervise sees STOPPING under the same lock and terminates it
+        # on arrival.
+        self._persist(run)
         return True
 
     def get(self, run_id):
@@ -523,6 +641,92 @@ def _terminate_group(proc):
         proc.terminate()
 
 
+def _pid_alive(pid):
+    """Whether `pid` is a live process, without touching it.
+
+    Only asked at startup, to decide whether a run from a previous engine is
+    still on the GPU (adopt it) or was interrupted (fail it). On POSIX,
+    `os.kill(pid, 0)` is the liveness idiom -- signal 0 is not delivered, it
+    only checks. On Windows signal 0 is *not* a query: CPython maps every
+    signal but CTRL_C/CTRL_BREAK to TerminateProcess, so it would kill a live
+    pid; Windows therefore gets an OpenProcess probe that only observes.
+
+    This does not guard against pid reuse -- a minimal version, as the audit
+    noted. The window is a restart landing on a recycled pid, which is narrow
+    on the desktop; the cost of getting it wrong is one adopted-then-reaped
+    run, not a lost GPU.
+    """
+    try:
+        pid = int(pid)
+    except (TypeError, ValueError):
+        return False
+    if pid <= 0:
+        return False
+    if IS_WINDOWS:
+        return _pid_alive_windows(pid)
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True             # exists, owned by someone else
+    except OSError:
+        return False
+    return True
+
+
+def _pid_alive_windows(pid):
+    """Observe, never signal: open the process and read its exit code."""
+    import ctypes
+    from ctypes import wintypes
+
+    PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
+    STILL_ACTIVE = 259
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+    kernel32.OpenProcess.restype = wintypes.HANDLE
+    kernel32.OpenProcess.argtypes = (wintypes.DWORD, wintypes.BOOL, wintypes.DWORD)
+    handle = kernel32.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, False, pid)
+    if not handle:
+        return False
+    try:
+        code = wintypes.DWORD()
+        if not kernel32.GetExitCodeProcess(handle, ctypes.byref(code)):
+            return False
+        return code.value == STILL_ACTIVE
+    finally:
+        kernel32.CloseHandle(handle)
+
+
+def _terminate_pid_group(pid):
+    """Stop an adopted run reached only by its pid, tree and all.
+
+    The sibling of `_terminate_group`, for a run this process did not spawn and
+    so has no Popen for. On POSIX the pid *is* its own group leader -- `_spawn`
+    started it with a new session -- so killpg on the pid reaches the trainer
+    and everything it spawned; on Windows taskkill /T walks the tree. The same
+    self-preservation guard applies: never signal our own group.
+    """
+    try:
+        pid = int(pid)
+    except (TypeError, ValueError):
+        return
+    if IS_WINDOWS:
+        _taskkill_tree(pid)
+        return
+    try:
+        group = os.getpgid(pid)
+    except (ProcessLookupError, PermissionError, OSError):
+        return
+    if group == os.getpgid(0):
+        # The adopted pid is not in its own group -- refuse rather than take
+        # the engine down with it.
+        return
+    try:
+        os.killpg(group, signal.SIGTERM)
+    except (ProcessLookupError, PermissionError, OSError):
+        pass
+
+
 def _write_config(path, config):
     """Write the config as YAML, falling back to JSON.