furu 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

furu/storage/state.py

@@ -9,13 +9,12 @@ from collections.abc import Generator
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Annotated, Any, Callable, Literal, Mapping, TypedDict, TypeAlias
+from typing import Annotated, Any, Callable, Literal, Mapping, TypeAlias, TypedDict
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationError, model_validator
 
 from ..errors import FuruLockNotAcquired, FuruWaitTimeout
 
-
 # Type alias for scheduler-specific metadata. Different schedulers (SLURM, LSF, PBS, local)
 # return different fields, so this must remain dynamic.
 SchedulerMetadata = dict[str, Any]
@@ -167,7 +166,6 @@ class _StateAttemptBase(BaseModel):
     number: int = 1
     backend: str
     started_at: str
-    heartbeat_at: str
     lease_duration_sec: float
     lease_expires_at: str
     owner: StateOwner
@@ -228,7 +226,6 @@ class StateAttempt(BaseModel):
     backend: str
     status: str
     started_at: str
-    heartbeat_at: str
     lease_duration_sec: float
     lease_expires_at: str
     owner: StateOwner
@@ -246,7 +243,6 @@ class StateAttempt(BaseModel):
             backend=attempt.backend,
             status=attempt.status,
             started_at=attempt.started_at,
-            heartbeat_at=attempt.heartbeat_at,
             lease_duration_sec=attempt.lease_duration_sec,
             lease_expires_at=attempt.lease_expires_at,
             owner=attempt.owner,
@@ -286,9 +282,9 @@ class StateManager:
     EVENTS_FILE = "events.jsonl"
     SUCCESS_MARKER = "SUCCESS.json"
 
-    COMPUTE_LOCK = ".compute.lock"
-    SUBMIT_LOCK = ".submit.lock"
-    STATE_LOCK = ".state.lock"
+    COMPUTE_LOCK = "compute.lock"
+    SUBMIT_LOCK = "submit.lock"
+    STATE_LOCK = "state.lock"
 
     TERMINAL_STATUSES = {
         "success",
@@ -302,6 +298,12 @@ class StateManager:
     def get_internal_dir(cls, directory: Path) -> Path:
         return directory / cls.INTERNAL_DIR
 
+    @classmethod
+    def ensure_internal_dir(cls, directory: Path) -> Path:
+        internal_dir = cls.get_internal_dir(directory)
+        internal_dir.mkdir(parents=True, exist_ok=True)
+        return internal_dir
+
     @classmethod
     def get_state_path(cls, directory: Path) -> Path:
         return cls.get_internal_dir(directory) / cls.STATE_FILE
@@ -366,7 +368,6 @@ class StateManager:
     @classmethod
     def _write_state_unlocked(cls, directory: Path, state: _FuruState) -> None:
         state_path = cls.get_state_path(directory)
-        state_path.parent.mkdir(parents=True, exist_ok=True)
         tmp_path = state_path.with_suffix(".tmp")
         tmp_path.write_text(json.dumps(state.model_dump(mode="json"), indent=2))
         os.replace(tmp_path, state_path)
@@ -385,7 +386,6 @@ class StateManager:
     @classmethod
     def try_lock(cls, lock_path: Path) -> int | None:
         try:
-            lock_path.parent.mkdir(parents=True, exist_ok=True)
             fd = os.open(lock_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o644)
             payload = {
                 "pid": os.getpid(),
@@ -400,15 +400,41 @@ class StateManager:
 
     @classmethod
     def release_lock(cls, fd: int | None, lock_path: Path) -> None:
-        if fd is not None:
+        """Release a lock acquired via :meth:`try_lock`.
+
+        We best-effort avoid deleting a lock that was broken and replaced by
+        another process by verifying the inode of the open fd matches the
+        current lock_path inode before unlinking.
+        """
+        if fd is None:
+            return
+        try:
+            fd_stat = os.fstat(fd)
+        except OSError:
+            fd_stat = None
+        try:
+            path_stat = lock_path.stat()
+        except FileNotFoundError:
+            path_stat = None
+        try:
+            if (
+                fd_stat is not None
+                and path_stat is not None
+                and fd_stat.st_ino == path_stat.st_ino
+                and fd_stat.st_dev == path_stat.st_dev
+            ):
+                lock_path.unlink(missing_ok=True)
+        finally:
             os.close(fd)
-        lock_path.unlink(missing_ok=True)
 
     @classmethod
     def _read_lock_info(cls, lock_path: Path) -> _LockInfoDict | None:
         if not lock_path.is_file():
             return None
-        text = lock_path.read_text().strip()
+        try:
+            text = lock_path.read_text().strip()
+        except FileNotFoundError:
+            return None
         if not text:
             return None
         lines = text.splitlines()
@@ -459,19 +485,23 @@ class StateManager:
 
     @classmethod
     def update_state(
-        cls, directory: Path, mutator: Callable[[_FuruState], None]
+        cls, directory: Path, mutator: Callable[[_FuruState], bool | None]
     ) -> _FuruState:
         lock_path = cls.get_lock_path(directory, cls.STATE_LOCK)
         fd: int | None = None
         try:
             fd = cls._acquire_lock_blocking(lock_path)
+            state_path = cls.get_state_path(directory)
+            force_write = not state_path.is_file()
             state = cls.read_state(directory)
-            mutator(state)
-            state.schema_version = cls.SCHEMA_VERSION
-            state.updated_at = cls._iso_now()
-            validated = _FuruState.model_validate(state)
-            cls._write_state_unlocked(directory, validated)
-            return validated
+            changed = mutator(state)
+            if force_write or changed is not False:
+                state.schema_version = cls.SCHEMA_VERSION
+                state.updated_at = cls._iso_now()
+                validated = _FuruState.model_validate(state)
+                cls._write_state_unlocked(directory, validated)
+                return validated
+            return state
         finally:
             cls.release_lock(fd, lock_path)
 
@@ -484,14 +514,12 @@ class StateManager:
             "host": socket.gethostname(),
             **event,
         }
-        path.parent.mkdir(parents=True, exist_ok=True)
         with path.open("a", encoding="utf-8") as f:
             f.write(json.dumps(enriched) + "\n")
 
     @classmethod
     def write_success_marker(cls, directory: Path, *, attempt_id: str) -> None:
         marker = cls.get_success_marker_path(directory)
-        marker.parent.mkdir(parents=True, exist_ok=True)
         payload = {"attempt_id": attempt_id, "created_at": cls._iso_now()}
         tmp = marker.with_suffix(".tmp")
         tmp.write_text(json.dumps(payload, indent=2))
@@ -510,6 +538,26 @@ class StateManager:
             return True
         return cls._utcnow() >= expires
 
+    @classmethod
+    def last_heartbeat_mtime(cls, directory: Path) -> float | None:
+        lock_path = cls.get_lock_path(directory, cls.COMPUTE_LOCK)
+        try:
+            return lock_path.stat().st_mtime
+        except FileNotFoundError:
+            return None
+
+    @classmethod
+    def _running_heartbeat_reason(
+        cls, directory: Path, attempt: _StateAttemptRunning
+    ) -> str | None:
+        last_heartbeat = cls.last_heartbeat_mtime(directory)
+        if last_heartbeat is None:
+            return "missing_heartbeat"
+        expires_at = last_heartbeat + float(attempt.lease_duration_sec)
+        if time.time() >= expires_at:
+            return "lease_expired"
+        return None
+
     @classmethod
     def start_attempt_queued(
         cls,
@@ -578,7 +626,6 @@ class StateManager:
 
             owner_state = StateOwner.model_validate(owner)
             started_at = now.isoformat(timespec="seconds")
-            heartbeat_at = started_at
             lease_duration = float(lease_duration_sec)
             lease_expires_at = expires.isoformat(timespec="seconds")
             scheduler_state: SchedulerMetadata = scheduler or {}
@@ -588,7 +635,6 @@ class StateManager:
                 number=int(number),
                 backend=backend,
                 started_at=started_at,
-                heartbeat_at=heartbeat_at,
                 lease_duration_sec=lease_duration,
                 lease_expires_at=lease_expires_at,
                 owner=owner_state,
@@ -635,49 +681,9 @@ class StateManager:
         return attempt.id
 
     @classmethod
-    def heartbeat(
-        cls, directory: Path, *, attempt_id: str, lease_duration_sec: float
-    ) -> bool:
-        ok = False
-
-        def mutate(state: _FuruState) -> None:
-            nonlocal ok
-            attempt = state.attempt
-            if not isinstance(attempt, _StateAttemptRunning):
-                return
-            if attempt.id != attempt_id:
-                return
-            now = cls._utcnow()
-            expires = now + _dt.timedelta(seconds=float(lease_duration_sec))
-            attempt.heartbeat_at = now.isoformat(timespec="seconds")
-            attempt.lease_duration_sec = float(lease_duration_sec)
-            attempt.lease_expires_at = expires.isoformat(timespec="seconds")
-            ok = True
-
-        cls.update_state(directory, mutate)
-        return ok
-
-    @classmethod
-    def set_attempt_fields(
-        cls, directory: Path, *, attempt_id: str, fields: SchedulerMetadata
-    ) -> bool:
-        ok = False
-
-        def mutate(state: _FuruState) -> None:
-            nonlocal ok
-            attempt = state.attempt
-            if attempt is None or attempt.id != attempt_id:
-                return
-            for key, value in fields.items():
-                if key == "scheduler" and isinstance(value, dict):
-                    attempt.scheduler.update(value)
-                    continue
-                if hasattr(attempt, key):
-                    setattr(attempt, key, value)
-            ok = True
-
-        cls.update_state(directory, mutate)
-        return ok
+    def heartbeat(cls, directory: Path) -> None:
+        lock_path = cls.get_lock_path(directory, cls.COMPUTE_LOCK)
+        os.utime(lock_path)
 
     @classmethod
     def finish_attempt_success(cls, directory: Path, *, attempt_id: str) -> None:
@@ -691,7 +697,6 @@ class StateManager:
                     number=attempt.number,
                     backend=attempt.backend,
                     started_at=attempt.started_at,
-                    heartbeat_at=attempt.heartbeat_at,
                     lease_duration_sec=attempt.lease_duration_sec,
                     lease_expires_at=attempt.lease_expires_at,
                     owner=attempt.owner,
@@ -728,7 +733,6 @@ class StateManager:
                     number=attempt.number,
                     backend=attempt.backend,
                     started_at=attempt.started_at,
-                    heartbeat_at=attempt.heartbeat_at,
                     lease_duration_sec=attempt.lease_duration_sec,
                     lease_expires_at=attempt.lease_expires_at,
                     owner=attempt.owner,
@@ -766,7 +770,6 @@ class StateManager:
                     number=attempt.number,
                     backend=attempt.backend,
                     started_at=attempt.started_at,
-                    heartbeat_at=attempt.heartbeat_at,
                     lease_duration_sec=attempt.lease_duration_sec,
                     lease_expires_at=attempt.lease_expires_at,
                     owner=attempt.owner,
@@ -816,10 +819,10 @@ class StateManager:
           to lease expiry.
         """
 
-        def mutate(state: _FuruState) -> None:
+        def mutate(state: _FuruState) -> bool:
             attempt = state.attempt
             if not isinstance(attempt, (_StateAttemptQueued, _StateAttemptRunning)):
-                return
+                return False
 
             # Fast promotion if we can see a durable success marker.
             if cls.success_marker_exists(directory):
@@ -829,7 +832,6 @@ class StateManager:
                     number=attempt.number,
                     backend=attempt.backend,
                     started_at=attempt.started_at,
-                    heartbeat_at=attempt.heartbeat_at,
                     lease_duration_sec=attempt.lease_duration_sec,
                     lease_expires_at=attempt.lease_expires_at,
                     owner=attempt.owner,
@@ -839,7 +841,7 @@ class StateManager:
                 state.result = _coerce_result(
                     state.result, status="success", created_at=ended
                 )
-                return
+                return True
 
             backend = attempt.backend
             now = cls._iso_now()
@@ -852,6 +854,10 @@ class StateManager:
                 if alive is False:
                     terminal_status = "crashed"
                     reason = "pid_dead"
+                elif isinstance(attempt, _StateAttemptRunning):
+                    reason = cls._running_heartbeat_reason(directory, attempt)
+                    if reason is not None:
+                        terminal_status = "crashed"
                 elif cls._lease_expired(attempt):
                     terminal_status = "crashed"
                     reason = "lease_expired"
@@ -864,16 +870,25 @@ class StateManager:
                         attempt.scheduler.update(
                             {k: v for k, v in verdict.items() if k != "terminal_status"}
                         )
-                if terminal_status is None and cls._lease_expired(attempt):
-                    terminal_status = "crashed"
-                    reason = "lease_expired"
+                if terminal_status is None:
+                    if isinstance(attempt, _StateAttemptRunning):
+                        reason = cls._running_heartbeat_reason(directory, attempt)
+                        if reason is not None:
+                            terminal_status = "crashed"
+                    elif cls._lease_expired(attempt):
+                        terminal_status = "crashed"
+                        reason = "lease_expired"
             else:
-                if cls._lease_expired(attempt):
+                if isinstance(attempt, _StateAttemptRunning):
+                    reason = cls._running_heartbeat_reason(directory, attempt)
+                    if reason is not None:
+                        terminal_status = "crashed"
+                elif cls._lease_expired(attempt):
                     terminal_status = "crashed"
                     reason = "lease_expired"
 
             if terminal_status is None:
-                return
+                return False
             if terminal_status == "success":
                 terminal_status = "crashed"
                 reason = reason or "scheduler_success_no_success_marker"
@@ -884,7 +899,6 @@ class StateManager:
                     number=attempt.number,
                     backend=attempt.backend,
                     started_at=attempt.started_at,
-                    heartbeat_at=attempt.heartbeat_at,
                     lease_duration_sec=attempt.lease_duration_sec,
                     lease_expires_at=attempt.lease_expires_at,
                     owner=attempt.owner,
@@ -901,7 +915,6 @@ class StateManager:
                         number=attempt.number,
                         backend=attempt.backend,
                         started_at=attempt.started_at,
-                        heartbeat_at=attempt.heartbeat_at,
                         lease_duration_sec=attempt.lease_duration_sec,
                         lease_expires_at=attempt.lease_expires_at,
                         owner=attempt.owner,
@@ -916,7 +929,6 @@ class StateManager:
                         number=attempt.number,
                         backend=attempt.backend,
                         started_at=attempt.started_at,
-                        heartbeat_at=attempt.heartbeat_at,
                         lease_duration_sec=attempt.lease_duration_sec,
                         lease_expires_at=attempt.lease_expires_at,
                         owner=attempt.owner,
@@ -931,7 +943,6 @@ class StateManager:
                         number=attempt.number,
                         backend=attempt.backend,
                         started_at=attempt.started_at,
-                        heartbeat_at=attempt.heartbeat_at,
                         lease_duration_sec=attempt.lease_duration_sec,
                         lease_expires_at=attempt.lease_expires_at,
                         owner=attempt.owner,
@@ -944,6 +955,7 @@ class StateManager:
                 state.result,
                 status="failed" if terminal_status == "failed" else "incomplete",
             )
+            return True
 
         state = cls.update_state(directory, mutate)
         attempt = state.attempt
@@ -978,6 +990,7 @@ def compute_lock(
     wait_log_every_sec: float = 10.0,
     reconcile_fn: Callable[[Path], None] | None = None,
     allow_failed: bool = False,
+    allow_success: bool = False,
 ) -> Generator[ComputeLockContext, None, None]:
     """
     Context manager that atomically acquires lock + records attempt + starts heartbeat.
@@ -1002,6 +1015,7 @@ def compute_lock(
         wait_log_every_sec: Interval between "waiting for lock" log messages
         reconcile_fn: Optional function to call to reconcile stale attempts
         allow_failed: Allow recomputation even if state is failed
+        allow_success: Allow recomputation even if state is successful
 
     Yields:
         ComputeLockContext with attempt_id and stop_heartbeat callable
@@ -1039,16 +1053,28 @@ def compute_lock(
         return ", ".join(parts)
 
     def _describe_wait(attempt: _StateAttempt, waited_sec: float) -> str:
-        label = "last heartbeat"
-        timestamp = attempt.heartbeat_at
         if attempt.status == "queued":
             label = "queued at"
             timestamp = attempt.started_at
-        parsed = StateManager._parse_time(timestamp)
-        timestamp_info = timestamp
-        if parsed is not None:
-            age = (StateManager._utcnow() - parsed).total_seconds()
-            timestamp_info = f"{timestamp} ({_format_wait_duration(age)} ago)"
+            parsed = StateManager._parse_time(timestamp)
+            timestamp_info = timestamp
+            if parsed is not None:
+                age = (StateManager._utcnow() - parsed).total_seconds()
+                timestamp_info = f"{timestamp} ({_format_wait_duration(age)} ago)"
+        else:
+            label = "last heartbeat"
+            last_heartbeat = StateManager.last_heartbeat_mtime(directory)
+            if last_heartbeat is None:
+                timestamp_info = "missing"
+            else:
+                heartbeat_dt = _dt.datetime.fromtimestamp(
+                    last_heartbeat, tz=_dt.timezone.utc
+                )
+                age = time.time() - last_heartbeat
+                timestamp_info = (
+                    f"{heartbeat_dt.isoformat(timespec='seconds')} "
+                    f"({_format_wait_duration(age)} ago)"
+                )
         return (
             "waited "
             f"{_format_wait_duration(waited_sec)}, {label} {timestamp_info}, "
@@ -1097,7 +1123,7 @@ def compute_lock(
         lock_fd = StateManager.try_lock(lock_path)
         if lock_fd is not None:
             state = StateManager.read_state(directory)
-            if isinstance(state.result, _StateResultSuccess):
+            if isinstance(state.result, _StateResultSuccess) and not allow_success:
                 StateManager.release_lock(lock_fd, lock_path)
                 raise FuruLockNotAcquired(
                     "Cannot acquire lock: experiment already succeeded"
@@ -1117,7 +1143,7 @@ def compute_lock(
                 if reconcile_fn is not None:
                     reconcile_fn(directory)
                     state = StateManager.read_state(directory)
-                if isinstance(state.result, _StateResultSuccess):
+                if isinstance(state.result, _StateResultSuccess) and not allow_success:
                     raise FuruLockNotAcquired(
                         "Cannot acquire lock: experiment already succeeded"
                     )
@@ -1151,7 +1177,7 @@ def compute_lock(
         attempt = state.attempt
 
         # If result is terminal, no point waiting
-        if isinstance(state.result, _StateResultSuccess):
+        if isinstance(state.result, _StateResultSuccess) and not allow_success:
             raise FuruLockNotAcquired(
                 "Cannot acquire lock: experiment already succeeded"
             )
@@ -1200,11 +1226,7 @@ def compute_lock(
         # Start heartbeat IMMEDIATELY
         def heartbeat() -> None:
             while not stop_event.wait(heartbeat_interval_sec):
-                StateManager.heartbeat(
-                    directory,
-                    attempt_id=attempt_id,  # type: ignore[arg-type]
-                    lease_duration_sec=lease_duration_sec,
-                )
+                StateManager.heartbeat(directory)
 
         thread = threading.Thread(target=heartbeat, daemon=True)
         thread.start()