pivot 0.2.1__tar.gz → 0.2.3__tar.gz

{pivot-0.2.1 → pivot-0.2.3}/PKG-INFO

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

{pivot-0.2.1 → pivot-0.2.3}/pyproject.toml

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

{pivot-0.2.1 → pivot-0.2.3}/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.1 → pivot-0.2.3}/src/pivot/engine/engine.py

@@ -1839,7 +1839,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:

{pivot-0.2.1 → pivot-0.2.3}/src/pivot/engine/sinks.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import dataclasses
 import sys
 import time
 from typing import TYPE_CHECKING
@@ -10,8 +11,8 @@ import rich.live
 import rich.markup
 import rich.text
 
-from pivot.engine.types import StageCompleted
-from pivot.types import DisplayCategory, categorize_stage_result
+from pivot.engine.types import LogLine, StageCompleted
+from pivot.types import LOSSY_OUTPUT_PHRASE, DisplayCategory, categorize_stage_result
 
 if TYPE_CHECKING:
     from pivot.engine.types import OutputEvent
@@ -84,6 +85,20 @@ def _categorize(event: StageCompleted) -> DisplayCategory:
     return categorize_stage_result(event["status"])
 
 
+@dataclasses.dataclass
+class _StageWarnings:
+    """A stage's ``WARNING:`` lines for the report.
+
+    Pivot's own output-write warnings are always kept and print first. The stage's own
+    warnings are counted per distinct line (in first-seen order); the report shows the
+    first ``_MAX_LOG_LINES_PER_STAGE`` distinct lines and counts the rest.
+    """
+
+    stage: dict[str, int] = dataclasses.field(default_factory=dict[str, int])
+    lossy_outputs: list[str] = dataclasses.field(default_factory=list[str])
+    dropped: int = 0
+
+
 def _print_completions(
     console: rich.console.Console,
     events: list[StageCompleted],
@@ -91,13 +106,19 @@ def _print_completions(
     total: int,
     max_name_width: int,
     stage_logs: dict[str, list[str]] | None = None,
+    stage_warnings: dict[str, _StageWarnings] | None = None,
 ) -> None:
-    """Print buffered completions sorted by index with skip collapsing."""
+    """Print buffered completions sorted by index with skip collapsing.
+
+    A failed stage shows its captured output; any other stage shows the ``WARNING:``
+    lines in ``stage_warnings`` (pass None when those already streamed).
+    """
     if not events:
         return
 
     sorted_events = sorted(events, key=lambda e: e["index"])
     logs = stage_logs or {}
+    warnings_by_stage = stage_warnings or {}
 
     i = 0
     while i < len(sorted_events):
@@ -152,12 +173,53 @@ def _print_completions(
                 if event["reason"]:
                     for detail in _format_error_detail(event["reason"], total=total):
                         console.print(detail)
+            elif event["stage"] in warnings_by_stage:
+                buffered = warnings_by_stage[event["stage"]]
+                stage_lines = [
+                    line if count == 1 else f"{line} (×{count})"
+                    for line, count in buffered.stage.items()
+                ]
+                for warning_line in [*buffered.lossy_outputs, *stage_lines]:
+                    escaped = rich.markup.escape(warning_line)
+                    console.print(f"          [yellow]{escaped}[/yellow]")
+                if buffered.dropped:
+                    console.print(
+                        f"          [yellow]… {buffered.dropped} more warnings "
+                        + "(run with --show-output)[/yellow]"
+                    )
             i += 1
 
 
 _MAX_LOG_LINES_PER_STAGE = 50
 
 
+def _buffer_log_line(
+    event: LogLine,
+    stage_logs: dict[str, list[str]],
+    stage_warnings: dict[str, _StageWarnings],
+) -> None:
+    """Keep a stage's output for a failure report and its warnings for every report.
+
+    Warnings have their own buffer: output writes log them after the stage's own output,
+    which can already fill the log buffer.
+    """
+    line = event["line"]
+    buf = stage_logs.setdefault(event["stage"], [])
+    if len(buf) < _MAX_LOG_LINES_PER_STAGE:
+        buf.append(line)
+    if not line.startswith("WARNING:"):
+        return
+    buffered = stage_warnings.setdefault(event["stage"], _StageWarnings())
+    if LOSSY_OUTPUT_PHRASE in line and line.startswith("WARNING: Stage '"):
+        buffered.lossy_outputs.append(line)
+    elif line in buffered.stage:
+        buffered.stage[line] += 1
+    elif len(buffered.stage) < _MAX_LOG_LINES_PER_STAGE:
+        buffered.stage[line] = 1
+    else:
+        buffered.dropped += 1
+
+
 class StaticConsoleSink:
     """Pipe/CI sink: buffers completions, prints a sorted report when each run ends."""
 
@@ -166,6 +228,7 @@ class StaticConsoleSink:
     _max_name_width: int
     _completion_buffer: list[StageCompleted]
     _stage_logs: dict[str, list[str]]
+    _stage_warnings: dict[str, _StageWarnings]
     _total: int
 
     def __init__(self, *, console: rich.console.Console, show_output: bool = False) -> None:
@@ -174,6 +237,7 @@ class StaticConsoleSink:
         self._max_name_width = 0
         self._completion_buffer = list[StageCompleted]()
         self._stage_logs = dict[str, list[str]]()
+        self._stage_warnings = dict[str, _StageWarnings]()
         self._total = 0
 
     async def handle(self, event: OutputEvent) -> None:
@@ -187,11 +251,9 @@ class StaticConsoleSink:
                 self._max_name_width = max(self._max_name_width, len(event["stage"]))
                 self._completion_buffer.append(event)
             case "log_line":
-                # Always buffer for failed-stage output; stream if --show-output
+                # Buffer for the report; stream if --show-output
                 stage = event["stage"]
-                buf = self._stage_logs.setdefault(stage, [])
-                if len(buf) < _MAX_LOG_LINES_PER_STAGE:
-                    buf.append(event["line"])
+                _buffer_log_line(event, self._stage_logs, self._stage_warnings)
                 if self._show_output:
                     line = rich.markup.escape(event["line"])
                     if event["is_stderr"]:
@@ -218,10 +280,12 @@ class StaticConsoleSink:
             total=self._total,
             max_name_width=self._max_name_width,
             stage_logs=self._stage_logs,
+            stage_warnings=None if self._show_output else self._stage_warnings,
         )
         self._max_name_width = 0
         self._completion_buffer.clear()
         self._stage_logs.clear()
+        self._stage_warnings.clear()
         self._total = 0
 
     async def close(self) -> None:
@@ -254,6 +318,7 @@ class LiveConsoleSink:
     _max_name_width: int
     _completion_buffer: list[StageCompleted]
     _stage_logs: dict[str, list[str]]
+    _stage_warnings: dict[str, _StageWarnings]
     _total: int
     _live: rich.live.Live | None
 
@@ -268,6 +333,7 @@ class LiveConsoleSink:
         self._max_name_width = 0
         self._completion_buffer = list[StageCompleted]()
         self._stage_logs = dict[str, list[str]]()
+        self._stage_warnings = dict[str, _StageWarnings]()
         self._total = 0
         self._live = None
 
@@ -372,9 +438,7 @@ class LiveConsoleSink:
                 self._update_live()
             case "log_line":
                 stage = event["stage"]
-                buf = self._stage_logs.setdefault(stage, [])
-                if len(buf) < _MAX_LOG_LINES_PER_STAGE:
-                    buf.append(event["line"])
+                _buffer_log_line(event, self._stage_logs, self._stage_warnings)
                 if self._show_output:
                     line = rich.markup.escape(event["line"])
                     if event["is_stderr"]:
@@ -405,6 +469,7 @@ class LiveConsoleSink:
             total=self._total,
             max_name_width=self._max_name_width,
             stage_logs=self._stage_logs,
+            stage_warnings=None if self._show_output else self._stage_warnings,
         )
         self._running.clear()
         self._completed_count = 0
@@ -412,6 +477,7 @@ class LiveConsoleSink:
         self._max_name_width = 0
         self._completion_buffer.clear()
         self._stage_logs.clear()
+        self._stage_warnings.clear()
         self._total = 0
 
     async def close(self) -> None:

{pivot-0.2.1 → pivot-0.2.3}/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.1 → pivot-0.2.3}/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.1 → pivot-0.2.3}/src/pivot/executor/worker.py

@@ -18,6 +18,7 @@ import random
 import sys
 import threading
 import traceback
+import warnings
 from typing import IO, TYPE_CHECKING, Any, Literal, NamedTuple, TypedDict, cast, overload, override
 
 import pydantic
@@ -153,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]
 
@@ -193,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"])
 
@@ -275,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:
@@ -330,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:
@@ -351,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:
@@ -404,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(
@@ -819,6 +825,42 @@ def _write_output_with_store(
         _reject_path_inside_project(ref, value, project_root, dependency_paths)
     output_path = store.prepare_output(ref)
     writer = cast("Any", ref.format)
+
+    output_label = (
+        f"its {type(writer).__name__} output"
+        if ref.identity.key is None
+        else f"output {ref.identity.key!r}"
+    )
+
+    def save_output(output_value: Any, destination: pathlib.Path) -> None:
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always", loaders.LossyOutputWarning)
+            try:
+                writer.save(output_value, destination)
+            except TypeError as e:
+                if isinstance(writer, loaders.YAML):
+                    raise TypeError(
+                        f"Stage '{ref.identity.producer}' cannot write {output_label}: {e}"
+                    ) from e
+                raise
+        for caught_warning in caught:
+            if issubclass(caught_warning.category, loaders.LossyOutputWarning):
+                logger.warning(
+                    "Stage '%s' wrote %s %s: %s",
+                    ref.identity.producer,
+                    output_label,
+                    types.LOSSY_OUTPUT_PHRASE,
+                    caught_warning.message,
+                )
+            else:
+                warnings.warn_explicit(
+                    caught_warning.message,
+                    caught_warning.category,
+                    caught_warning.filename,
+                    caught_warning.lineno,
+                    source=caught_warning.source,
+                )
+
     if ref.tag is types.ArtifactTag.DIRECTORY:
         if not isinstance(value, dict):
             raise RuntimeError(
@@ -845,11 +887,11 @@ def _write_output_with_store(
             members.append((output_path / normalized, item_value))
         for full_path, item_value in members:
             full_path.parent.mkdir(parents=True, exist_ok=True)
-            writer.save(item_value, full_path)
+            save_output(item_value, full_path)
         return output_path
 
     output_path.parent.mkdir(parents=True, exist_ok=True)
-    writer.save(value, output_path)
+    save_output(value, output_path)
     return output_path
 
 
@@ -885,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:
@@ -903,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)}'. "
@@ -915,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(
@@ -941,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 (
@@ -969,6 +1013,7 @@ def _run_stage_function_with_store(
                 lock_data,
                 store,
                 state_dir,
+                project_state_dir,
             )
         )
 
@@ -1014,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