furu 0.0.4__py3-none-any.whl → 0.0.6__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(),
@@ -485,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)
 
@@ -510,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))
@@ -536,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,
@@ -604,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 {}
@@ -614,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,
@@ -661,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:
@@ -717,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,
@@ -754,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,
@@ -792,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,
@@ -842,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):
@@ -855,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,
@@ -865,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()
@@ -878,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"
@@ -890,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"
@@ -910,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,
@@ -927,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,
@@ -942,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,
@@ -957,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,
@@ -970,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
@@ -1067,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}, "
@@ -1228,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()

furu/testing.py

@@ -0,0 +1,232 @@
+from collections.abc import Generator, Mapping, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import cast, overload
+
+import chz
+import pytest
+
+from .config import FURU_CONFIG, RecordGitMode
+from .core import Furu
+from .runtime.overrides import OverrideValue, override_furu_hashes
+
+OverrideKey = Furu | str
+_PATH_MISSING = object()
+
+
+@dataclass(frozen=True)
+class _FuruConfigSnapshot:
+    base_root: Path
+    version_controlled_root_override: Path | None
+    record_git: RecordGitMode
+    allow_no_git_origin: bool
+    poll_interval: float
+    stale_timeout: float
+    max_wait_time_sec: float | None
+    lease_duration_sec: float
+    heartbeat_interval_sec: float
+    cancelled_is_preempted: bool
+    retry_failed: bool
+    submitit_root: Path
+
+    @classmethod
+    def capture(cls) -> "_FuruConfigSnapshot":
+        return cls(
+            base_root=FURU_CONFIG.base_root,
+            version_controlled_root_override=FURU_CONFIG.version_controlled_root_override,
+            record_git=FURU_CONFIG.record_git,
+            allow_no_git_origin=FURU_CONFIG.allow_no_git_origin,
+            poll_interval=FURU_CONFIG.poll_interval,
+            stale_timeout=FURU_CONFIG.stale_timeout,
+            max_wait_time_sec=FURU_CONFIG.max_wait_time_sec,
+            lease_duration_sec=FURU_CONFIG.lease_duration_sec,
+            heartbeat_interval_sec=FURU_CONFIG.heartbeat_interval_sec,
+            cancelled_is_preempted=FURU_CONFIG.cancelled_is_preempted,
+            retry_failed=FURU_CONFIG.retry_failed,
+            submitit_root=FURU_CONFIG.submitit_root,
+        )
+
+    def restore(self) -> None:
+        FURU_CONFIG.base_root = self.base_root
+        FURU_CONFIG.version_controlled_root_override = (
+            self.version_controlled_root_override
+        )
+        FURU_CONFIG.record_git = self.record_git
+        FURU_CONFIG.allow_no_git_origin = self.allow_no_git_origin
+        FURU_CONFIG.poll_interval = self.poll_interval
+        FURU_CONFIG.stale_timeout = self.stale_timeout
+        FURU_CONFIG.max_wait_time_sec = self.max_wait_time_sec
+        FURU_CONFIG.lease_duration_sec = self.lease_duration_sec
+        FURU_CONFIG.heartbeat_interval_sec = self.heartbeat_interval_sec
+        FURU_CONFIG.cancelled_is_preempted = self.cancelled_is_preempted
+        FURU_CONFIG.retry_failed = self.retry_failed
+        FURU_CONFIG.submitit_root = self.submitit_root
+
+
+def _apply_test_config(base_root: Path) -> Path:
+    root = base_root.resolve()
+    FURU_CONFIG.base_root = root
+    FURU_CONFIG.version_controlled_root_override = root / "furu-data" / "artifacts"
+    FURU_CONFIG.record_git = "ignore"
+    FURU_CONFIG.allow_no_git_origin = False
+    FURU_CONFIG.poll_interval = 0.01
+    FURU_CONFIG.stale_timeout = 0.1
+    FURU_CONFIG.max_wait_time_sec = None
+    FURU_CONFIG.lease_duration_sec = 0.05
+    FURU_CONFIG.heartbeat_interval_sec = 0.01
+    FURU_CONFIG.cancelled_is_preempted = True
+    FURU_CONFIG.retry_failed = True
+    FURU_CONFIG.submitit_root = root / "submitit"
+    return root
+
+
+@contextmanager
+def furu_test_env(base_root: Path) -> Generator[Path, None, None]:
+    snapshot = _FuruConfigSnapshot.capture()
+    root = _apply_test_config(base_root)
+    try:
+        yield root
+    finally:
+        snapshot.restore()
+
+
+@overload
+def override_results(
+    overrides: Mapping[Furu, OverrideValue],
+) -> Generator[None, None, None]: ...
+
+
+@overload
+def override_results(
+    overrides: Mapping[str, OverrideValue],
+) -> Generator[None, None, None]: ...
+
+
+@contextmanager
+def override_results(
+    overrides: Mapping[OverrideKey, OverrideValue],
+) -> Generator[None, None, None]:
+    """Override specific Furu results within the context.
+
+    Overrides are keyed by furu_hash, so identical configs share a stub.
+    Keys may be a Furu instance or a furu_hash string.
+    """
+    hash_overrides: dict[str, OverrideValue] = {}
+    for key, value in overrides.items():
+        hash_key: str
+        if isinstance(key, Furu):
+            hash_key = cast(str, key.furu_hash)
+        elif isinstance(key, str):
+            if not key:
+                raise ValueError("override furu_hash must be non-empty")
+            hash_key = key
+        else:
+            raise TypeError(
+                "override_results keys must be Furu instances or furu_hash strings"
+            )
+        hash_overrides[hash_key] = value
+    with override_furu_hashes(hash_overrides):
+        yield
+
+
+@overload
+def override_results_for(
+    root: Furu,
+    overrides: Mapping[str, OverrideValue],
+) -> Generator[None, None, None]: ...
+
+
+@contextmanager
+def override_results_for(
+    root: Furu,
+    overrides: Mapping[str, OverrideValue],
+) -> Generator[None, None, None]:
+    """Override results by dotted field path relative to a root Furu object.
+
+    Paths follow chz traversal: use numeric segments for list indices and
+    key segments for mappings (e.g. "deps.0" or "deps.key").
+    """
+    hash_overrides: dict[str, OverrideValue] = {}
+    for path, value in overrides.items():
+        target = _resolve_override_path(root, path)
+        hash_overrides[cast(str, target.furu_hash)] = value
+    with override_furu_hashes(hash_overrides):
+        yield
+
+
+def _resolve_override_path(root: Furu, path: str) -> Furu:
+    if not path:
+        raise ValueError("override path must be non-empty")
+    current: object = root
+    for segment in path.split("."):
+        if not segment:
+            raise ValueError(f"override path has empty segment: {path!r}")
+        current = _resolve_path_segment(current, segment, path)
+    if not isinstance(current, Furu):
+        raise TypeError(
+            f"override path {path!r} must resolve to a Furu instance; got {type(current).__name__}"
+        )
+    return current
+
+
+def _resolve_path_segment(current: object, segment: str, path: str) -> object:
+    if chz.is_chz(current):
+        value = getattr(current, segment, _PATH_MISSING)
+        if value is _PATH_MISSING:
+            raise AttributeError(f"override path {path!r} has no attribute {segment!r}")
+        return value
+    if isinstance(current, Mapping):
+        mapping = cast(Mapping[object, OverrideValue], current)
+        return _resolve_mapping_segment(mapping, segment, path)
+    if _is_indexable_sequence(current):
+        sequence = cast(Sequence[OverrideValue], current)
+        return _resolve_sequence_segment(sequence, segment, path)
+    raise TypeError(
+        f"override path {path!r} cannot traverse into {type(current).__name__}"
+    )
+
+
+def _resolve_mapping_segment(
+    mapping: Mapping[object, OverrideValue],
+    segment: str,
+    path: str,
+) -> OverrideValue:
+    if segment in mapping:
+        return mapping[segment]
+    index = _parse_index_segment(segment)
+    if index is not None and index in mapping:
+        return mapping[index]
+    raise KeyError(f"override path {path!r} has no key {segment!r}")
+
+
+def _resolve_sequence_segment(
+    sequence: Sequence[OverrideValue],
+    segment: str,
+    path: str,
+) -> OverrideValue:
+    index = _parse_index_segment(segment)
+    if index is None:
+        raise TypeError(f"override path {path!r} index {segment!r} is not an integer")
+    if index < 0:
+        raise ValueError(f"override path {path!r} index must be non-negative")
+    return sequence[index]
+
+
+def _parse_index_segment(segment: str) -> int | None:
+    if not segment.isdigit():
+        return None
+    return int(segment)
+
+
+def _is_indexable_sequence(value: object) -> bool:
+    if isinstance(value, (str, bytes, bytearray)):
+        return False
+    return isinstance(value, Sequence)
+
+
+@pytest.fixture()
+def furu_tmp_root(tmp_path: Path) -> Generator[Path, None, None]:
+    """Configure furu to use a temporary root for the test."""
+    with furu_test_env(tmp_path) as root:
+        yield root