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

furu/core/furu.py

@@ -10,12 +10,29 @@ import traceback
 from abc import ABC, abstractmethod
 from pathlib import Path
 from types import FrameType
-from typing import Any, Callable, ClassVar, Self, TypedDict, TypeVar, cast, overload
+from typing import (
+    AbstractSet,
+    Any,
+    Callable,
+    ClassVar,
+    Hashable,
+    Mapping,
+    Protocol,
+    Self,
+    Sequence,
+    TypedDict,
+    TypeAlias,
+    TypeVar,
+    cast,
+    overload,
+)
 
 import chz
 import submitit
 from typing_extensions import dataclass_transform
 
+from chz.field import Field as ChzField
+
 from ..adapters import SubmititAdapter
 from ..adapters.submitit import SubmititJob
 from ..config import FURU_CONFIG
@@ -41,7 +58,6 @@ from ..storage import (
 from ..storage.state import (
     _FuruState,
     _OwnerDict,
-    _StateAttemptFailed,
     _StateAttemptQueued,
     _StateAttemptRunning,
     _StateResultAbsent,
@@ -177,6 +193,29 @@ class Furu[T](ABC):
         """Validate that result is complete and correct (override if needed)."""
         return True
 
+    def _dependencies(self: Self) -> "DependencySpec | None":
+        """Return extra dependencies not captured by fields."""
+        return None
+
+    def _get_dependencies(self: Self, *, recursive: bool = True) -> list["Furu"]:
+        """Collect Furu dependencies from fields and `_dependencies()`."""
+        seen = {self._furu_hash}
+        dependencies: list[Furu] = []
+        _collect_dependencies(self, dependencies, seen, recursive=recursive)
+        return dependencies
+
+    def _dependency_hashes(self: Self) -> list[str]:
+        dependencies = _direct_dependencies(self)
+        if not dependencies:
+            return []
+
+        digests: set[str] = set()
+        for dependency in dependencies:
+            if dependency is self:
+                raise ValueError("Furu dependencies cannot include self")
+            digests.add(dependency._furu_hash)
+        return sorted(digests)
+
     def _invalidate_cached_success(self: Self, directory: Path, *, reason: str) -> None:
         logger = get_logger()
         logger.warning(
@@ -275,20 +314,33 @@ class Furu[T](ABC):
         return MigrationManager.read_migration(self._base_furu_dir())
 
     @overload
-    def load_or_create(self, executor: submitit.Executor) -> T | submitit.Job[T]: ...
+    def load_or_create(
+        self,
+        executor: submitit.Executor,
+        *,
+        retry_failed: bool | None = None,
+    ) -> T | submitit.Job[T]: ...
 
     @overload
-    def load_or_create(self, executor: None = None) -> T: ...
+    def load_or_create(
+        self,
+        executor: None = None,
+        *,
+        retry_failed: bool | None = None,
+    ) -> T: ...
 
     def load_or_create(
         self: Self,
         executor: submitit.Executor | None = None,
+        *,
+        retry_failed: bool | None = None,
     ) -> T | submitit.Job[T]:
         """
         Load result if it exists, computing if necessary.
 
         Args:
             executor: Optional executor for batch submission (e.g., submitit.Executor)
+            retry_failed: Whether to retry failed results (default uses FURU_RETRY_FAILED)
 
         Returns:
             Result if wait=True, job handle if wait=False, or None if already exists
@@ -299,6 +351,9 @@ class Furu[T](ABC):
         logger = get_logger()
         parent_holder = current_holder()
         has_parent = parent_holder is not None and parent_holder is not self
+        retry_failed_effective = (
+            retry_failed if retry_failed is not None else FURU_CONFIG.retry_failed
+        )
         if has_parent:
             logger.debug(
                 "dep: begin %s %s %s",
@@ -380,6 +435,16 @@ class Furu[T](ABC):
 
                 state0 = StateManager.read_state(directory)
 
+                if (
+                    isinstance(state0.result, _StateResultFailed)
+                    and not retry_failed_effective
+                ):
+                    raise self._build_failed_state_error(
+                        directory,
+                        state0,
+                        message="Computation previously failed",
+                    )
+
                 needs_reconcile = True
                 if isinstance(state0.result, _StateResultSuccess):
                     # Double check logic if we fell through to here (e.g. race condition or invalidation above)
@@ -462,7 +527,8 @@ class Furu[T](ABC):
                 # Synchronous execution
                 if executor is None:
                     status, created_here, result = self._run_locally(
-                        start_time=start_time
+                        start_time=start_time,
+                        allow_failed=retry_failed_effective,
                     )
                     if status == "success":
                         ok = True
@@ -478,19 +544,10 @@ class Furu[T](ABC):
                         )
                         return self._load()
 
-                    state = StateManager.read_state(directory)
-                    attempt = state.attempt
-                    message = (
-                        attempt.error.message
-                        if isinstance(attempt, _StateAttemptFailed)
-                        else None
-                    )
-                    suffix = (
-                        f": {message}" if isinstance(message, str) and message else ""
-                    )
-                    raise FuruComputeError(
-                        f"Computation {status}{suffix}",
-                        StateManager.get_state_path(directory),
+                    raise self._build_failed_state_error(
+                        directory,
+                        None,
+                        message="Computation previously failed",
                     )
 
                 # Asynchronous execution with submitit
@@ -504,7 +561,12 @@ class Furu[T](ABC):
                     "load_or_create: %s -> submitit submit_once()",
                     self.__class__.__name__,
                 )
-                job = self._submit_once(adapter, directory, None)
+                job = self._submit_once(
+                    adapter,
+                    directory,
+                    None,
+                    allow_failed=retry_failed_effective,
+                )
                 ok = True
                 return cast(submitit.Job[T], job)
         finally:
@@ -547,12 +609,56 @@ class Furu[T](ABC):
             },
         )
 
+    def _add_exception_breadcrumbs(self, exc: BaseException, directory: Path) -> None:
+        if not hasattr(exc, "add_note"):
+            return
+        state_path = StateManager.get_state_path(directory)
+        log_path = StateManager.get_internal_dir(directory) / "furu.log"
+        note = (
+            f"Furu directory: {directory}\n"
+            f"State file: {state_path}\n"
+            f"Log file: {log_path}"
+        )
+        exc.add_note(note)
+
+    @staticmethod
+    def _failed_state_hints() -> list[str]:
+        return [
+            "To retry this failed artifact: set FURU_RETRY_FAILED=1 or call load_or_create(retry_failed=True).",
+            "To inspect details: open the state file and furu.log shown above.",
+        ]
+
+    def _build_failed_state_error(
+        self,
+        directory: Path,
+        state: _FuruState | None,
+        *,
+        message: str,
+    ) -> FuruComputeError:
+        current_state = state or StateManager.read_state(directory)
+        attempt = current_state.attempt
+        error = getattr(attempt, "error", None) if attempt is not None else None
+        return FuruComputeError(
+            message,
+            StateManager.get_state_path(directory),
+            recorded_error_type=getattr(error, "type", None),
+            recorded_error_message=getattr(error, "message", None),
+            recorded_traceback=getattr(error, "traceback", None),
+            hints=self._failed_state_hints(),
+        )
+
+    def _effective_max_wait_time_sec(self) -> float | None:
+        if FURU_CONFIG.max_wait_time_sec is not None:
+            return FURU_CONFIG.max_wait_time_sec
+        return self._max_wait_time_sec
+
     def _check_timeout(self, start_time: float) -> None:
         """Check if operation has timed out."""
-        if self._max_wait_time_sec is not None:
-            if time.time() - start_time > self._max_wait_time_sec:
+        max_wait_time = self._effective_max_wait_time_sec()
+        if max_wait_time is not None:
+            if time.time() - start_time > max_wait_time:
                 raise FuruWaitTimeout(
-                    f"Furu operation timed out after {self._max_wait_time_sec} seconds."
+                    f"Furu operation timed out after {max_wait_time} seconds."
                 )
 
     def _is_migrated_state(self, directory: Path) -> bool:
@@ -613,8 +719,10 @@ class Furu[T](ABC):
         event: dict[str, str | int] = {
             "type": "migration_overwrite",
             "policy": record.policy,
-            "from": f"{record.from_namespace}:{record.from_hash}",
-            "to": f"{record.to_namespace}:{record.to_hash}",
+            "from_namespace": record.from_namespace,
+            "from_hash": record.from_hash,
+            "to_namespace": record.to_namespace,
+            "to_hash": record.to_hash,
             "reason": reason,
         }
         StateManager.append_event(directory, event.copy())
@@ -625,6 +733,8 @@ class Furu[T](ABC):
         adapter: SubmititAdapter,
         directory: Path,
         on_job_id: Callable[[str], None] | None,
+        *,
+        allow_failed: bool,
     ) -> SubmititJob | None:
         """Submit job once without waiting (fire-and-forget mode)."""
         logger = get_logger()
@@ -693,7 +803,7 @@ class Furu[T](ABC):
                 scheduler={},
             )
 
-            job = adapter.submit(lambda: self._worker_entry())
+            job = adapter.submit(lambda: self._worker_entry(allow_failed=allow_failed))
 
             # Save job handle and watch for job ID
             adapter.pickle_job(job, directory)
@@ -729,7 +839,7 @@ class Furu[T](ABC):
         finally:
             StateManager.release_lock(lock_fd, lock_path)
 
-    def _worker_entry(self: Self) -> None:
+    def _worker_entry(self: Self, *, allow_failed: bool | None = None) -> None:
         """Entry point for worker process (called by submitit or locally)."""
         with enter_holder(self):
             logger = get_logger()
@@ -737,6 +847,9 @@ class Furu[T](ABC):
             directory.mkdir(parents=True, exist_ok=True)
 
             env_info = self._collect_submitit_env()
+            allow_failed_effective = (
+                allow_failed if allow_failed is not None else FURU_CONFIG.retry_failed
+            )
 
             try:
                 with compute_lock(
@@ -758,19 +871,23 @@ class Furu[T](ABC):
                     poll_interval_sec=FURU_CONFIG.poll_interval,
                     wait_log_every_sec=FURU_CONFIG.wait_log_every_sec,
                     reconcile_fn=lambda d: self._reconcile(d),
+                    allow_failed=allow_failed_effective,
                 ) as ctx:
-                    # Refresh metadata (now safe - attempt is already recorded)
-                    metadata = MetadataManager.create_metadata(
-                        self, directory, ignore_diff=FURU_CONFIG.ignore_git_diff
-                    )
-                    MetadataManager.write_metadata(metadata, directory)
+                    stage = "metadata"
+                    try:
+                        # Refresh metadata (now safe - attempt is already recorded)
+                        metadata = MetadataManager.create_metadata(
+                            self, directory, ignore_diff=FURU_CONFIG.ignore_git_diff
+                        )
+                        MetadataManager.write_metadata(metadata, directory)
 
-                    # Set up signal handlers
-                    self._setup_signal_handlers(
-                        directory, ctx.stop_heartbeat, attempt_id=ctx.attempt_id
-                    )
+                        # Set up signal handlers
+                        stage = "signal handler setup"
+                        self._setup_signal_handlers(
+                            directory, ctx.stop_heartbeat, attempt_id=ctx.attempt_id
+                        )
 
-                    try:
+                        stage = "_create"
                         # Run computation
                         logger.debug(
                             "_create: begin %s %s %s",
@@ -798,13 +915,23 @@ class Furu[T](ABC):
                             extra={"furu_console_only": True},
                         )
                     except Exception as e:
-                        logger.error(
-                            "_create failed %s %s %s",
-                            self.__class__.__name__,
-                            self._furu_hash,
-                            directory,
-                            extra={"furu_file_only": True},
-                        )
+                        if stage == "_create":
+                            logger.error(
+                                "_create failed %s %s %s",
+                                self.__class__.__name__,
+                                self._furu_hash,
+                                directory,
+                                extra={"furu_file_only": True},
+                            )
+                        else:
+                            logger.error(
+                                "attempt failed (%s) %s %s %s",
+                                stage,
+                                self.__class__.__name__,
+                                self._furu_hash,
+                                directory,
+                                extra={"furu_file_only": True},
+                            )
                         logger.error(
                             "%s", format_traceback(e), extra={"furu_file_only": True}
                         )
@@ -821,6 +948,18 @@ class Furu[T](ABC):
                                 "traceback": tb,
                             },
                         )
+                        self._add_exception_breadcrumbs(e, directory)
+                        if stage != "_create":
+                            message = (
+                                "Failed to create metadata"
+                                if stage == "metadata"
+                                else "Failed to set up signal handlers"
+                            )
+                            raise FuruComputeError(
+                                message,
+                                StateManager.get_state_path(directory),
+                                e,
+                            ) from e
                         raise
             except FuruLockNotAcquired:
                 # Experiment already completed (success or failed), nothing to do
@@ -850,16 +989,22 @@ class Furu[T](ABC):
 
         return info
 
-    def _run_locally(self: Self, start_time: float) -> tuple[str, bool, T | None]:
+    def _run_locally(
+        self: Self,
+        start_time: float,
+        *,
+        allow_failed: bool,
+    ) -> tuple[str, bool, T | None]:
         """Run computation locally, returning (status, created_here, result)."""
         logger = get_logger()
         directory = self._base_furu_dir()
 
         # Calculate remaining time for the lock wait
         max_wait: float | None = None
-        if self._max_wait_time_sec is not None:
+        max_wait_time = self._effective_max_wait_time_sec()
+        if max_wait_time is not None:
             elapsed = time.time() - start_time
-            max_wait = max(0.0, self._max_wait_time_sec - elapsed)
+            max_wait = max(0.0, max_wait_time - elapsed)
 
         try:
             with compute_lock(
@@ -878,26 +1023,23 @@ class Furu[T](ABC):
                 poll_interval_sec=FURU_CONFIG.poll_interval,
                 wait_log_every_sec=FURU_CONFIG.wait_log_every_sec,
                 reconcile_fn=lambda d: self._reconcile(d),
+                allow_failed=allow_failed,
             ) as ctx:
-                # Create metadata (now safe - attempt is already recorded)
+                stage = "metadata"
                 try:
+                    # Create metadata (now safe - attempt is already recorded)
                     metadata = MetadataManager.create_metadata(
                         self, directory, ignore_diff=FURU_CONFIG.ignore_git_diff
                     )
                     MetadataManager.write_metadata(metadata, directory)
-                except Exception as e:
-                    raise FuruComputeError(
-                        "Failed to create metadata",
-                        StateManager.get_state_path(directory),
-                        e,
-                    ) from e
-
-                # Set up preemption handler
-                self._setup_signal_handlers(
-                    directory, ctx.stop_heartbeat, attempt_id=ctx.attempt_id
-                )
 
-                try:
+                    # Set up preemption handler
+                    stage = "signal handler setup"
+                    self._setup_signal_handlers(
+                        directory, ctx.stop_heartbeat, attempt_id=ctx.attempt_id
+                    )
+
+                    stage = "_create"
                     # Run the computation
                     logger.debug(
                         "_create: begin %s %s %s",
@@ -926,13 +1068,23 @@ class Furu[T](ABC):
                     )
                     return "success", True, result
                 except Exception as e:
-                    logger.error(
-                        "_create failed %s %s %s",
-                        self.__class__.__name__,
-                        self._furu_hash,
-                        directory,
-                        extra={"furu_file_only": True},
-                    )
+                    if stage == "_create":
+                        logger.error(
+                            "_create failed %s %s %s",
+                            self.__class__.__name__,
+                            self._furu_hash,
+                            directory,
+                            extra={"furu_file_only": True},
+                        )
+                    else:
+                        logger.error(
+                            "attempt failed (%s) %s %s %s",
+                            stage,
+                            self.__class__.__name__,
+                            self._furu_hash,
+                            directory,
+                            extra={"furu_file_only": True},
+                        )
                     logger.error(
                         "%s", format_traceback(e), extra={"furu_file_only": True}
                     )
@@ -950,6 +1102,18 @@ class Furu[T](ABC):
                             "traceback": tb,
                         },
                     )
+                    self._add_exception_breadcrumbs(e, directory)
+                    if stage != "_create":
+                        message = (
+                            "Failed to create metadata"
+                            if stage == "metadata"
+                            else "Failed to set up signal handlers"
+                        )
+                        raise FuruComputeError(
+                            message,
+                            StateManager.get_state_path(directory),
+                            e,
+                        ) from e
                     raise
         except FuruLockNotAcquired:
             # Lock couldn't be acquired because experiment already completed
@@ -998,4 +1162,201 @@ class Furu[T](ABC):
             signal.signal(sig, handle_signal)
 
 
+class DependencyChzSpec(Protocol):
+    __chz_fields__: dict[str, ChzField]
+
+
+DependencySequence: TypeAlias = Sequence[Furu]
+DependencySet: TypeAlias = AbstractSet[Furu]
+DependencyMapping: TypeAlias = Mapping[str, Furu]
+DependencyCollection: TypeAlias = DependencySequence | DependencySet | DependencyMapping
+DependencyValue: TypeAlias = Furu | DependencyCollection
+DependencySpec: TypeAlias = DependencyValue | DependencyChzSpec
+DependencyLeaf: TypeAlias = str | int | float | bool | None | Path | bytes
+DependencyScanValue: TypeAlias = (
+    DependencyLeaf
+    | Furu
+    | Mapping[Hashable, "DependencyScanValue"]
+    | Sequence["DependencyScanValue"]
+    | AbstractSet["DependencyScanValue"]
+    | DependencyChzSpec
+)
+
+
+def _collect_dependencies(
+    obj: Furu,
+    dependencies: list[Furu],
+    seen: set[str],
+    *,
+    recursive: bool,
+) -> None:
+    for dependency in _direct_dependencies(obj):
+        digest = dependency._furu_hash
+        if digest in seen:
+            continue
+        seen.add(digest)
+        dependencies.append(dependency)
+        if recursive:
+            _collect_dependencies(
+                dependency,
+                dependencies,
+                seen,
+                recursive=recursive,
+            )
+
+
+def _direct_dependencies(obj: Furu) -> list[Furu]:
+    dependencies: list[Furu] = []
+    for field in chz.chz_fields(obj).values():
+        value = cast(DependencyScanValue, getattr(obj, field.logical_name))
+        dependencies.extend(_collect_dependencies_from_value(value))
+    extra = obj._dependencies()
+    if extra is not None:
+        dependencies.extend(_collect_dependencies_from_spec(extra, path="dependencies"))
+    return dependencies
+
+
+def _collect_dependencies_from_value(value: DependencyScanValue) -> list[Furu]:
+    dependencies: list[Furu] = []
+    if isinstance(value, Furu):
+        dependencies.append(value)
+        return dependencies
+    if isinstance(value, dict):
+        mapping = cast(Mapping[Hashable, DependencyScanValue], value)
+        for item in mapping.values():
+            dependencies.extend(_collect_dependencies_from_value(item))
+        return dependencies
+    if isinstance(value, (list, tuple)):
+        sequence = cast(Sequence[DependencyScanValue], value)
+        for item in sequence:
+            dependencies.extend(_collect_dependencies_from_value(item))
+        return dependencies
+    if isinstance(value, (set, frozenset)):
+        items = _sorted_dependency_set(cast(AbstractSet[DependencyScanValue], value))
+        for item in items:
+            dependencies.extend(_collect_dependencies_from_value(item))
+        return dependencies
+    if chz.is_chz(value):
+        for field in chz.chz_fields(value).values():
+            field_value = cast(DependencyScanValue, getattr(value, field.logical_name))
+            dependencies.extend(_collect_dependencies_from_value(field_value))
+    return dependencies
+
+
+def _collect_dependencies_from_spec(value: DependencySpec, path: str) -> list[Furu]:
+    if isinstance(value, Furu):
+        return [value]
+    if isinstance(value, dict):
+        return _collect_dependencies_from_mapping(
+            cast(Mapping[Hashable, DependencyValue], value),
+            path,
+        )
+    if isinstance(value, (list, tuple)):
+        return _collect_dependencies_from_sequence(
+            cast(Sequence[DependencyValue], value),
+            path,
+        )
+    if isinstance(value, (set, frozenset)):
+        return _collect_dependencies_from_set(
+            cast(AbstractSet[DependencyValue], value),
+            path,
+        )
+    if chz.is_chz(value):
+        dependencies: list[Furu] = []
+        for field in chz.chz_fields(value).values():
+            field_value = getattr(value, field.logical_name)
+            field_path = f"{path}.{field.logical_name}"
+            dependencies.extend(
+                _collect_dependencies_from_value_spec(field_value, field_path)
+            )
+        return dependencies
+    raise _dependency_type_error(path, value)
+
+
+def _collect_dependencies_from_value_spec(
+    value: DependencyValue,
+    path: str,
+) -> list[Furu]:
+    if isinstance(value, Furu):
+        return [value]
+    if isinstance(value, dict):
+        return _collect_dependencies_from_mapping(
+            cast(Mapping[Hashable, DependencyValue], value),
+            path,
+        )
+    if isinstance(value, (list, tuple)):
+        return _collect_dependencies_from_sequence(
+            cast(Sequence[DependencyValue], value),
+            path,
+        )
+    if isinstance(value, (set, frozenset)):
+        return _collect_dependencies_from_set(
+            cast(AbstractSet[DependencyValue], value),
+            path,
+        )
+    raise _dependency_type_error(path, value)
+
+
+def _collect_dependencies_from_mapping(
+    mapping: Mapping[Hashable, DependencyValue],
+    path: str,
+) -> list[Furu]:
+    dependencies: list[Furu] = []
+    for key, item in mapping.items():
+        if not isinstance(item, Furu):
+            raise _dependency_type_error(f"{path}[{key!r}]", item)
+        dependencies.append(item)
+    return dependencies
+
+
+def _collect_dependencies_from_sequence(
+    sequence: Sequence[DependencyValue],
+    path: str,
+) -> list[Furu]:
+    dependencies: list[Furu] = []
+    for index, item in enumerate(sequence):
+        if not isinstance(item, Furu):
+            raise _dependency_type_error(f"{path}[{index}]", item)
+        dependencies.append(item)
+    return dependencies
+
+
+def _collect_dependencies_from_set(
+    values: AbstractSet[DependencyValue],
+    path: str,
+) -> list[Furu]:
+    dependencies: list[Furu] = []
+    ordered = sorted(
+        list(cast(AbstractSet[DependencyScanValue], values)),
+        key=_dependency_sort_key,
+    )
+    for index, item in enumerate(ordered):
+        if not isinstance(item, Furu):
+            raise _dependency_type_error(f"{path}[{index}]", item)
+        dependencies.append(item)
+    return dependencies
+
+
+def _sorted_dependency_set(
+    values: AbstractSet[DependencyScanValue],
+) -> list[DependencyScanValue]:
+    return sorted(list(values), key=_dependency_sort_key)
+
+
+def _dependency_sort_key(value: DependencyScanValue) -> tuple[int, str]:
+    if isinstance(value, Furu):
+        return (0, value._furu_hash)
+    return (1, f"{type(value).__name__}:{value!r}")
+
+
+def _dependency_type_error(
+    path: str,
+    value: DependencySpec | DependencyValue | DependencyScanValue,
+) -> TypeError:
+    return TypeError(
+        f"{path} must be a Furu instance or a collection of Furu instances; "
+        f"got {type(value).__name__}"
+    )
+
+
 _H = TypeVar("_H", bound=Furu, covariant=True)