pivot 0.2.2__tar.gz → 0.2.4__tar.gz

{pivot-0.2.2 → pivot-0.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pivot
-Version: 0.2.2
+Version: 0.2.4
 Summary: High-performance Python pipeline tool with automatic code change detection
 Author: Sami Jawhar
 Author-email: Sami Jawhar <sami@metr.org>

{pivot-0.2.2 → pivot-0.2.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pivot"
-version = "0.2.2"
+version = "0.2.4"
 description = "High-performance Python pipeline tool with automatic code change detection"
 requires-python = ">=3.13,<3.14"
 readme = "README.md"

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/cli/checkout.py

@@ -89,6 +89,9 @@ def _restore_path_sync(
 ) -> tuple[RestoreResult, str]:
     """Restore a file or directory from cache (sync version).
 
+    Holds the restore lock for path across the existence check and the restore, so a
+    concurrent checkout never compares against, or deletes, a half-restored output.
+
     Returns:
         Tuple of (result, path_name) for the caller to handle output.
 
@@ -97,6 +100,22 @@ def _restore_path_sync(
             "already exists" without --force). Cache misses return ("missing", name)
             instead of raising.
     """
+    with cache.restore_lock(path, state_dir) as lock_token:
+        return _restore_path_locked(
+            path, output_hash, cache_dir, checkout_modes, behavior, state_dir, lock_token
+        )
+
+
+def _restore_path_locked(
+    path: pathlib.Path,
+    output_hash: HashInfo,
+    cache_dir: pathlib.Path,
+    checkout_modes: list[cache.CheckoutMode],
+    behavior: CheckoutBehavior,
+    state_dir: pathlib.Path | None,
+    lock_token: str | None,
+) -> tuple[RestoreResult, str]:
+    """Body of _restore_path_sync; the caller holds the restore lock for path."""
     keep_ignored = False
     if path.exists():
         match behavior:
@@ -146,6 +165,7 @@ def _restore_path_sync(
         cache_dir,
         checkout_modes=checkout_modes,
         state_dir=state_dir,
+        lock_token=lock_token,
         keep_ignored=keep_ignored,
     )
     if not success:

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/engine/engine.py

@@ -1306,26 +1306,6 @@ class Engine:
                                                 run_id=run_id,
                                             )
 
-                            # Check cancellation
-                            if self._cancel_event.is_set():
-                                cancelled = self._scheduler.apply_cancel()
-                                for name, old_state in cancelled:
-                                    await self.emit(
-                                        StageStateChanged(
-                                            type="stage_state_changed",
-                                            stage=name,
-                                            state=StageExecutionState.COMPLETED,
-                                            previous_state=old_state,
-                                        )
-                                    )
-                                    await self._emit_terminal_stage(
-                                        name,
-                                        StageStatus.CANCELLED,
-                                        "cancelled",
-                                        results,
-                                        run_id=run_id,
-                                    )
-
                             # Start more stages if slots available
                             if not self._scheduler.stop_starting_new:
                                 await self._start_ready_stages(
@@ -1389,6 +1369,9 @@ class Engine:
                                 "Sentinel send failed (will rely on shutdown_event): %s", exc
                             )
 
+            if self._cancel_event.is_set():
+                await self._finalize_cancelled_stages(results, run_id)
+
         finally:
             # Close StateDBs first (lightweight), then shut down worker pool.
             with anyio.CancelScope(shield=True):
@@ -1682,6 +1665,19 @@ class Engine:
                     skipped = False
                 if skipped:
                     continue
+                # A watch event may cancel this run while its coordinator skip check awaited.
+                # Do not submit a stage to the pool after cancellation stopped admissions.
+                if self._cancel_event.is_set():
+                    self._scheduler.release_resources(stage_name)
+                    await self._set_stage_state(stage_name, StageExecutionState.COMPLETED)
+                    await self._emit_terminal_stage(
+                        stage_name,
+                        StageStatus.CANCELLED,
+                        "cancelled",
+                        results,
+                        run_id=run_id,
+                    )
+                    continue
 
                 # Stage must run — lazily create pool on first dispatch
                 started += 1
@@ -1839,7 +1835,7 @@ class Engine:
                 stage_info["outs"],
                 lock_data["output_hashes"],
                 state_db,
-                state_dir=stage_state_dir,
+                state_dir=state_dir,
                 file_hash_entries=restored_entries,
             )
             if restored and restored_entries:
@@ -2162,6 +2158,29 @@ class Engine:
 
         return duration_ms
 
+    async def _finalize_cancelled_stages(
+        self,
+        results: dict[str, executor_core.ExecutionSummary],
+        run_id: str,
+    ) -> None:
+        """Mark every ready or pending stage cancelled exactly once."""
+        for stage_name, old_state in self._scheduler.apply_cancel():
+            await self.emit(
+                StageStateChanged(
+                    type="stage_state_changed",
+                    stage=stage_name,
+                    state=StageExecutionState.COMPLETED,
+                    previous_state=old_state,
+                )
+            )
+            await self._emit_terminal_stage(
+                stage_name,
+                StageStatus.CANCELLED,
+                "cancelled",
+                results,
+                run_id=run_id,
+            )
+
     async def _emit_terminal_stage(
         self,
         stage_name: str,

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/executor/commit.py

@@ -150,7 +150,7 @@ def commit_stages(
                 identity_key = types.identity_key(out.identity)
                 try:
                     out_path = store.resolve_display_path(out)
-                    hash_info = store.commit(out, out_path)
+                    hash_info = store.commit(out, out_path, state_dir=default_state_dir)
                     file_hash_entries.extend(worker.output_file_hash_entries(out_path, hash_info))
                     output_hashes_by_key[identity_key] = hash_info
                     output_hashes_by_identity[out.identity] = hash_info

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/executor/core.py

@@ -337,6 +337,7 @@ def prepare_worker_info(
         params_arg_name=stage_info["params_arg_name"],
         project_root=project_root,
         state_dir=state_dir,
+        project_state_dir=default_state_dir,
         collection_params=stage_info["collection_params"],
     )
 

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/executor/worker.py

@@ -154,6 +154,7 @@ class WorkerStageInfo(TypedDict):
     params_arg_name: str | None
     project_root: pathlib.Path
     state_dir: pathlib.Path
+    project_state_dir: pathlib.Path
     collection_params: dict[str, str]
     incremental_inputs: dict[str, types.ArtifactIdentity]
 
@@ -194,6 +195,7 @@ def execute_stage(
     _ = cache_dir
     ring_buffer = _OutputRingBuffer()
     state_dir = stage_info["state_dir"]
+    project_state_dir = stage_info["project_state_dir"]
     project_root = stage_info["project_root"]
     store = store_mod.store_from_spec(stage_info["store_spec"])
 
@@ -276,7 +278,7 @@ def execute_stage(
                                 stage_outs,
                                 lock_data["output_hashes"],
                                 state_db,
-                                state_dir=stage_info["state_dir"],
+                                state_dir=project_state_dir,
                                 file_hash_entries=restored_entries,
                             )
                             if restored:
@@ -331,7 +333,7 @@ def execute_stage(
                             stage_outs,
                             lock_data["output_hashes"],
                             state_db,
-                            state_dir=stage_info["state_dir"],
+                            state_dir=project_state_dir,
                             file_hash_entries=restored_entries,
                         )
                         if restored:
@@ -352,7 +354,7 @@ def execute_stage(
                             stage_outs,
                             store,
                             state_db,
-                            state_dir=stage_info["state_dir"],
+                            state_dir=project_state_dir,
                         )
                         if run_cache_skip is not None:
                             if no_commit:
@@ -405,12 +407,15 @@ def execute_stage(
                     stage_info["incremental_inputs"],
                     lock_data,
                     state_dir,
+                    project_state_dir,
                     project_root,
                 )
                 if no_commit:
                     output_hashes = _hash_output_paths(output_paths)
                 else:
-                    output_hashes = _commit_outputs_with_store(output_paths, store)
+                    output_hashes = _commit_outputs_with_store(
+                        output_paths, store, project_state_dir
+                    )
                     for ref, path in output_paths.items():
                         file_hash_entries.extend(
                             output_file_hash_entries(
@@ -922,7 +927,8 @@ def _prepare_incremental_inputs(
     incremental_inputs: dict[str, types.ArtifactIdentity],
     lock_data: LockData | None,
     store: store_mod.Store,
-    state_dir: pathlib.Path,
+    stage_state_dir: pathlib.Path,
+    project_state_dir: pathlib.Path,
 ) -> dict[str, Any]:
     """Restore previous incremental outputs and load them for the stage call."""
     if not incremental_inputs:
@@ -940,7 +946,7 @@ def _prepare_incremental_inputs(
             continue
 
         output_hash = lock_data["output_hashes"].get(identity)
-        stage_lock = lock.StageLock(stage_name, lock.get_stages_dir(state_dir))
+        stage_lock = lock.StageLock(stage_name, lock.get_stages_dir(stage_state_dir))
         if output_hash is None:
             raise exceptions.CacheRestoreError(
                 f"Cache missing for incremental output '{types.identity_key(identity)}'. "
@@ -952,7 +958,7 @@ def _prepare_incremental_inputs(
         if not store.restore_from_cache(
             ref,
             output_hash,
-            state_dir=state_dir,
+            state_dir=project_state_dir,
             checkout_modes=[config.CheckoutMode.COPY],
         ):
             raise exceptions.CacheRestoreError(
@@ -978,6 +984,7 @@ def _run_stage_function_with_store(
     incremental_inputs: dict[str, types.ArtifactIdentity],
     lock_data: LockData | None,
     state_dir: pathlib.Path,
+    project_state_dir: pathlib.Path,
     project_root: pathlib.Path,
 ) -> dict[types.ArtifactRef, pathlib.Path]:
     with (
@@ -1006,6 +1013,7 @@ def _run_stage_function_with_store(
                 lock_data,
                 store,
                 state_dir,
+                project_state_dir,
             )
         )
 
@@ -1051,12 +1059,17 @@ def _hash_output_paths(
 def _commit_outputs_with_store(
     outputs_by_ref: dict[types.ArtifactRef, pathlib.Path],
     store: store_mod.Store,
+    project_state_dir: pathlib.Path,
 ) -> dict[str, HashInfo]:
     output_hashes = dict[str, HashInfo]()
     for ref, path in outputs_by_ref.items():
         if not path.exists():
             raise exceptions.OutputMissingError(f"Stage did not produce output: {path}")
-        output_hashes[types.identity_key(ref.identity)] = store.commit(ref, path)
+        if isinstance(store, store_mod.WorkspaceStore):
+            output_hash = store.commit(ref, path, state_dir=project_state_dir)
+        else:
+            output_hash = store.commit(ref, path)
+        output_hashes[types.identity_key(ref.identity)] = output_hash
     return output_hashes
 
 

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/fingerprint.py

@@ -1217,6 +1217,11 @@ def _resolve_dotted_path(parts: tuple[str, ...], ns: dict[str, Any]) -> Any:
     return obj
 
 
+def _method_function(method: Callable[..., Any]) -> Callable[..., Any]:
+    """Return the underlying function behind a bound method."""
+    return cast("Callable[..., Any]", getattr(method, "__func__", method))
+
+
 def get_loader_fingerprint(loader: "loaders.Writer[Any] | loaders.Reader[Any]") -> dict[str, str]:
     """Generate fingerprint manifest for a loader instance.
 
@@ -1244,13 +1249,14 @@ def get_loader_fingerprint(loader: "loaders.Writer[Any] | loaders.Reader[Any]")
         if method is not None:
             _process_callable_dependency(f"{class_name}.{method_name}", method, manifest, visited)
 
-    # Fingerprint save() if this is a Writer
+    # Use each method definition rather than a temporary bound method so the
+    # weak AST hash cache retains the entry across loader fingerprint calls.
     if isinstance(loader, loaders_module.Writer):
-        manifest[f"loader:{class_name}:save"] = hash_function_ast(loader.save)
+        manifest[f"loader:{class_name}:save"] = hash_function_ast(_method_function(loader.save))
 
     # Fingerprint load() if this is a Reader
     if isinstance(loader, loaders_module.Reader):
-        manifest[f"loader:{class_name}:load"] = hash_function_ast(loader.load)
+        manifest[f"loader:{class_name}:load"] = hash_function_ast(_method_function(loader.load))
 
         # Only fingerprint empty() if it's overridden in a Loader subclass
         if isinstance(loader, loaders_module.Loader):
@@ -1262,7 +1268,9 @@ def get_loader_fingerprint(loader: "loaders.Writer[Any] | loaders.Reader[Any]")
                     break
                 if "empty" in cls.__dict__:
                     # Found override in a subclass
-                    manifest[f"loader:{class_name}:empty"] = hash_function_ast(typed_loader.empty)
+                    manifest[f"loader:{class_name}:empty"] = hash_function_ast(
+                        _method_function(typed_loader.empty)
+                    )
                     break
 
     loader_any = cast("Any", loader)

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/remote/storage.py

@@ -2,12 +2,15 @@ from __future__ import annotations
 
 import asyncio
 import contextlib
+import errno
 import logging
 import os
 import re
 import tempfile
 from typing import TYPE_CHECKING, Protocol, TypedDict
 
+import xxhash
+
 from pivot import config, exceptions, metrics
 from pivot.remote import config as remote_config
 from pivot.types import TransferResult
@@ -35,6 +38,9 @@ BULK_EXISTS_LIST_THRESHOLD = 50
 
 _HEX_PATTERN = re.compile(r"^[a-f0-9]{16}$")
 
+# link() errors on filesystems without hard links (vfat, exFAT, many SMB and FUSE mounts).
+_NO_HARDLINK_ERRNO = frozenset({errno.EPERM, errno.ENOTSUP, errno.EOPNOTSUPP})
+
 # Cached S3 config to avoid re-parsing config on every call.
 # Safe without locking - asyncio is single-threaded, worst case race just
 # creates duplicate config (harmless).
@@ -188,8 +194,9 @@ async def _write_all_async(fd: int, data: bytes) -> None:
 async def _stream_download_to_fd(
     response: GetObjectOutputTypeDef,
     fd: int,
-) -> None:
-    """Stream S3 response body to file descriptor in chunks with timeout."""
+) -> str:
+    """Stream S3 response body to fd and return its xxh64 integrity digest."""
+    hasher = xxhash.xxh64()
     stream = response["Body"]
     try:
         while True:
@@ -199,10 +206,12 @@ async def _stream_download_to_fd(
             )
             if not chunk:
                 break
+            hasher.update(chunk)
             await _write_all_async(fd, chunk)
     finally:
         with contextlib.suppress(Exception):
             stream.close()  # pyright: ignore[reportUnknownMemberType]
+    return hasher.hexdigest()
 
 
 async def _atomic_download(
@@ -210,12 +219,19 @@ async def _atomic_download(
     bucket: str,
     key: str,
     local_path: Path,
+    expected_hash: str,
     *,
+    overwrite: bool,
     readonly: bool = False,
 ) -> None:
-    """Download S3 object to local path atomically via temp file with streaming.
+    """Download an expected cache blob atomically and verify its xxh64 integrity.
+
+    The digest detects accidental corruption during storage or transfer. It does not
+    authenticate remote content.
 
     Args:
+        overwrite: Replace an existing file at local_path. When False, an existing file is
+            kept: a content-addressed cache entry another process may be hardlinking.
         readonly: If True, set file permissions to 0o444 (for cache files).
     """
     local_path.parent.mkdir(parents=True, exist_ok=True)
@@ -235,13 +251,40 @@ async def _atomic_download(
                     f"Access denied to {bucket}/{key}. Check AWS credentials."
                 ) from exc
             raise exceptions.RemoteError(f"Failed to download from S3: {exc}") from exc
-        await _stream_download_to_fd(response, fd)
+        actual_hash = await _stream_download_to_fd(response, fd)
         os.close(fd)
         fd = -1  # Mark as closed
+        if actual_hash != expected_hash:
+            raise exceptions.RemoteError(
+                f"Hash mismatch: expected {expected_hash}, got {actual_hash} "
+                + f"for {bucket}/{key}. Delete {bucket}/{key} from the remote, then run "
+                + "`pivot push` on a machine whose cache has a good copy and has not run "
+                + "`pivot push`, `pivot pull`, `pivot fetch`, or `pivot status --remote` "
+                + "against this remote for this object. Pivot records hashes those commands "
+                + "upload, download, or find, then skips recorded hashes. If every machine "
+                + "with a good copy has recorded it, delete `.pivot/state.lmdb/` before "
+                + "pushing; this also clears run history and local caches."
+            )
         if readonly:
             os.chmod(tmp_path, 0o444)
-        os.replace(tmp_path, local_path)
-        move_succeeded = True
+        if overwrite:
+            os.replace(tmp_path, local_path)
+            move_succeeded = True
+        else:
+            # link() refuses an existing name, so a concurrent reader never sees the
+            # entry disappear. The temp file is removed below either way.
+            try:
+                os.link(tmp_path, local_path)
+            except FileExistsError:
+                pass
+            except OSError as exc:
+                if exc.errno not in _NO_HARDLINK_ERRNO:
+                    raise
+                # No hard links here, so nothing can be hardlinking the entry either:
+                # rename unless another process already promoted it.
+                if not os.path.exists(local_path):
+                    os.replace(tmp_path, local_path)
+                    move_succeeded = True
     finally:
         if fd >= 0:
             os.close(fd)
@@ -555,6 +598,8 @@ class S3Remote:
                     self._bucket,
                     _hash_to_key(self._prefix, cache_hash),
                     local_path,
+                    cache_hash,
+                    overwrite=True,
                     readonly=readonly,
                 )
         except botocore_exc.ClientError as e:
@@ -614,7 +659,10 @@ class S3Remote:
         *,
         readonly: bool = False,
     ) -> list[TransferResult]:
-        """Download multiple files in parallel with atomic writes and streaming.
+        """Download multiple cache entries in parallel with atomic writes and streaming.
+
+        A destination that already exists is kept, never replaced: destinations are
+        content-addressed cache entries that other processes may be hardlinking.
 
         Args:
             readonly: If True, set file permissions to 0o444 (for cache files).
@@ -643,6 +691,8 @@ class S3Remote:
                                 self._bucket,
                                 _hash_to_key(self._prefix, hash_),
                                 local_path,
+                                hash_,
+                                overwrite=False,
                                 readonly=readonly,
                             )
                             completed += 1

{pivot-0.2.2 → pivot-0.2.4}/src/pivot/storage/AGENTS.md

@@ -56,6 +56,16 @@ When multiple `pivot` processes share the same StateDB:
   - LMDB's internal write serialization (via `lock.mdb`) handles actual mutual exclusion
 - On timeout: raises `PivotDBWriteTimeoutError` with diagnostic message
 
+## Restore Locking
+
+- Restores with a `state_dir` take an exclusive per-target `flock` in `state_dir/locks/`.
+- Restore lock files persist. Do not unlink or age-delete them: a waiter can hold the
+  old inode while a new opener locks a replacement inode, defeating mutual exclusion.
+- File restores stage their link or copy beside the target, under a name tagged with the
+  target's token, and atomically replace it. Directory restores build in a private
+  temporary directory and swap it. Both run under the target's restore lock, and each
+  first removes the token-tagged entries a killed restore of the same target left.
+
 ## Path Storage
 
 All paths stored in the database must be **relative** to ensure portability across machines and correct cache behavior.