Filesystem state persistence with flock to survive @spaces.GPU fork

game/scheduler.py

@@ -19,18 +19,21 @@ Reference: https://huggingface.co/docs/hub/spaces-zerogpu
 
 from __future__ import annotations
 
+import fcntl
+import json
 import random
 import threading
 import time
 import traceback
 from collections import deque
+from contextlib import contextmanager
 from pathlib import Path
 
 from smolagents import CodeAgent
 from smolagents.monitoring import LogLevel
 
 from .actions import advance, build_tools
-from .character import Character
+from .character import Action, Character, default_traits
 from .dreams import run_dream
 from .interpreter import SharedInterpreter
 from .perception import system_prompt_for, task_prompt_for
@@ -38,6 +41,9 @@ from .sentiment import classify_board
 from .smol_adapter import BackendModel
 from .world import GRID_H, GRID_W, PLACES, BoardPost, ShellLine, WorldState
 
+STATE_PATH = Path("/tmp/townlet/state.json")
+LOCK_PATH = Path("/tmp/townlet/state.lock")
+
 TICK_SECONDS = 0.4
 GOAL_CHECK_INTERVAL_TICKS = 60
 MAX_AGENT_STEPS = 2
@@ -58,6 +64,93 @@ class Townlet:
         # thread (no LLM work), drained by the @spaces.GPU endpoint.
         self._pending_decisions: deque[str] = deque()
 
+    # ----- shared state via filesystem -----
+    #
+    # `@spaces.GPU` forks the endpoint into a subprocess, so the parent
+    # process (handling non-GPU endpoints like /world_snapshot) and the
+    # child (handling /tick_decisions) do NOT share Python memory. We
+    # serialise the whole world on every endpoint to keep them in sync.
+
+    @contextmanager
+    def _state_lock(self):
+        STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
+        lf = open(LOCK_PATH, "w")
+        fcntl.flock(lf.fileno(), fcntl.LOCK_EX)
+        try:
+            self._load_state_unlocked()
+            yield
+            self._save_state_unlocked()
+        finally:
+            fcntl.flock(lf.fileno(), fcntl.LOCK_UN)
+            lf.close()
+
+    def _save_state_unlocked(self) -> None:
+        state = {
+            "tick": self.world.tick,
+            "paused": self.world.paused,
+            "characters": [c.to_full_dict() for c in self.world.characters.values()],
+            "board": [{"author": p.author, "text": p.text, "tick": p.tick} for p in self.world.board],
+            "shell_lines": [
+                {"author": s.author, "kind": s.kind, "text": s.text, "tick": s.tick}
+                for s in self.world.shell_lines
+            ],
+            "shell_lock_holder": self.world.shell_lock_holder,
+            "shell_queue": list(self.world.shell_queue),
+            "board_sentiment": self.world.board_sentiment,
+            "pending": list(self._pending_decisions),
+        }
+        tmp = STATE_PATH.with_suffix(".json.tmp")
+        tmp.write_text(json.dumps(state))
+        tmp.rename(STATE_PATH)
+
+    def _load_state_unlocked(self) -> bool:
+        if not STATE_PATH.exists():
+            return False
+        try:
+            state = json.loads(STATE_PATH.read_text())
+        except Exception:
+            return False
+        new_chars: dict[str, Character] = {}
+        for d in state.get("characters", []):
+            c = Character(
+                name=d["name"],
+                personality=d["personality"],
+                model_id=d["model_id"],
+                sprite_id=d["sprite_id"],
+                archetype=d.get("archetype", "wanderer"),
+                journal=list(d.get("journal", [])),
+                traits=dict(d.get("traits", default_traits())),
+                goal=d.get("goal", ""),
+                pos=tuple(d["pos"]),
+                inventory=dict(d.get("inventory", {"electricity": 0, "water": 0})),
+                locker=dict(d.get("locker", {"electricity": 0, "water": 0})),
+                alive=bool(d.get("alive", True)),
+            )
+            a = d.get("current_action")
+            if a:
+                c.current_action = Action(
+                    verb=a["verb"],
+                    params=dict(a.get("params", {})),
+                    ticks_remaining=int(a.get("ticks_remaining", 0)),
+                )
+            new_chars[c.name] = c
+        self.world.characters = new_chars
+        self.world.tick = int(state.get("tick", 0))
+        self.world.paused = bool(state.get("paused", False))
+        self.world.board = [
+            BoardPost(author=p["author"], text=p["text"], tick=p["tick"])
+            for p in state.get("board", [])
+        ]
+        self.world.shell_lines = [
+            ShellLine(author=s["author"], kind=s["kind"], text=s["text"], tick=s["tick"])
+            for s in state.get("shell_lines", [])
+        ]
+        self.world.shell_lock_holder = state.get("shell_lock_holder")
+        self.world.shell_queue = list(state.get("shell_queue", []))
+        self.world.board_sentiment = state.get("board_sentiment", "neutral")
+        self._pending_decisions = deque(state.get("pending", []))
+        return True
+
     # ----- lifecycle -----
 
     def start(self) -> None:
@@ -83,12 +176,12 @@ class Townlet:
     # ----- mutation API -----
 
     def add_character(self, char: Character) -> None:
-        with self._mutex:
+        with self._state_lock(), self._mutex:
             self.world.characters[char.name] = char
             self._persist_character(char)
 
     def set_personality(self, name: str, personality: str) -> bool:
-        with self._mutex:
+        with self._state_lock(), self._mutex:
             c = self.world.characters.get(name)
             if c is None:
                 return False
@@ -101,7 +194,7 @@ class Townlet:
 
         if model_id not in ROSTER:
             return False
-        with self._mutex:
+        with self._state_lock(), self._mutex:
             c = self.world.characters.get(name)
             if c is None:
                 return False
@@ -110,20 +203,17 @@ class Townlet:
             return True
 
     def set_paused(self, paused: bool) -> None:
-        self.world.paused = paused
+        with self._state_lock():
+            self.world.paused = paused
 
     # ----- snapshot -----
 
     def snapshot(self) -> dict:
-        # Advance the world clock on every snapshot poll too — keeps the
-        # animation smooth between tick_decisions calls. Cheap (no LLM).
-        import os
-        print(
-            f"[townlet] SNAPSHOT pid={os.getpid()} tid={threading.get_ident()} "
-            f"T[id={id(self):x}] tick={self.world.tick}",
-            flush=True,
-        )
-        self.tick_world(2)
+        with self._state_lock():
+            self.tick_world(2)
+            return self._snapshot_dict()
+
+    def _snapshot_dict(self) -> dict:
         with self._mutex:
             return {
                 "tick": self.world.tick,
@@ -192,19 +282,13 @@ class Townlet:
 
     def drain_decisions(self, limit: int = DECISIONS_PER_DRAIN, world_ticks: int = 5) -> int:
         """Advance the world clock then pop up to `limit` ready characters
-        and run their CodeAgent.
-
-        Called from app.py's @spaces.GPU-decorated endpoint so all LLM work
-        AND world-clock work happens inside the GPU boundary on Zero-GPU.
-        Keeping both in the same call ensures they share object graph (which
-        a background daemon thread on Gradio Server does not reliably do).
+        and run their CodeAgent. Wrapped in _state_lock so mutations
+        survive the @spaces.GPU subprocess fork.
         """
-        import os
-        print(
-            f"[townlet] DRAIN-ENTRY pid={os.getpid()} tid={threading.get_ident()} "
-            f"T[id={id(self):x}] tick={self.world.tick}",
-            flush=True,
-        )
+        with self._state_lock():
+            return self._drain_decisions_inner(limit, world_ticks)
+
+    def _drain_decisions_inner(self, limit: int, world_ticks: int) -> int:
         self.tick_world(world_ticks)
         with self._mutex:
             start_pending = list(self._pending_decisions)