okami.mom/albedo

git clone https://git.nkp.pet/okami.mom/albedo

nekomimi.petfix kernel output truncation with automatic disk-backed pagination1f92415

main
branchesmainall refs →
4.9 KiB132 linesraw
1"""One honest output surface for captured local and SSH command-mode jobs."""
2
3from __future__ import annotations
4
5import os
6from pathlib import Path
7import tempfile
8
9from albedo_api import OUTPUT_PREVIEW, RETAIN, OutputCapture, Text, excerpt
10
11
12def _retention_error(capture: OutputCapture) -> ValueError:
13    spill = getattr(capture, "spill", None)
14    recovery = (
15        f"The retained spill file is {spill}; it does not contain the discarded suffix. "
16        if spill
17        else ""
18    )
19    recovery += (
20        "For job artifacts beyond retention, rerun with "
21        "job.pipe('tee', path) attached before awaiting."
22    )
23    return ValueError(
24        f"output has {capture.seen} bytes, but only the first {capture.retained} "
25        "are retained; refusing incomplete output. "
26        "Paginate the retained prefix with job.read(offset=0, limit=65536) "
27        "or output.read(id, offset=0, limit=65536), advancing the byte offset. "
28        "Bytes beyond retention cannot be recovered by pagination. " + recovery
29    )
30
31
32def read_page(capture: OutputCapture, offset: int, limit: int) -> Text:
33    if offset < 0 or not 0 <= limit <= OUTPUT_PREVIEW:
34        raise ValueError(
35            f"offset >= 0 and 0 <= limit <= {OUTPUT_PREVIEW}; page with offset="
36        )
37    if not limit:
38        return Text("")
39    if min(offset + limit, capture.seen) > capture.retained:
40        raise _retention_error(capture)
41    return Text(capture.read(offset, limit))
42
43
44class JobOutput:
45    capture: OutputCapture
46    duration: float | None
47    exit_code: int | None
48    _read: bool
49
50    def tail(self, n: int = 4000, *, lines: int | None = None) -> Text:
51        """The last n characters or lines, within the 64 KiB preview window."""
52        data = self.capture.tail()
53        result = excerpt(
54            data.decode("utf-8", errors="ignore"),
55            n,
56            lines,
57            end=True,
58            clipped=self.capture.seen > len(data),
59        )
60        if self.exit_code is not None:
61            self._read = True
62        return result
63
64    def head(self, n: int = 4000, *, lines: int | None = None) -> Text:
65        """The first n characters or lines, within the 64 KiB preview window."""
66        data = self.capture.data[:OUTPUT_PREVIEW]
67        result = excerpt(
68            data.decode("utf-8", errors="ignore"),
69            n,
70            lines,
71            end=False,
72            clipped=self.capture.seen > len(data),
73        )
74        if self.exit_code is not None:
75            self._read = True
76        return result
77
78    def _require_complete(self) -> None:
79        if self.duration is None:
80            raise RuntimeError(
81                "output is still running; await job before job.read() or job.save(path)"
82            )
83        if self.capture.seen > self.capture.retained:
84            raise _retention_error(self.capture)
85
86    def read(self, offset: int = 0, limit: int | None = None) -> Text:
87        """Complete text after completion, or a byte-offset page with explicit limit.
88
89        Whole text reads are at most 1 MiB. Pages are at most 64 KiB, read spill
90        files transparently, and may read a running job. save() preserves bytes.
91        """
92        if limit is None:
93            if offset < 0:
94                raise ValueError("offset must be nonnegative")
95            self._require_complete()
96            size = max(0, self.capture.seen - offset)
97            if size > RETAIN:
98                raise ValueError(
99                    "whole text reads are limited to 1 MiB; paginate with "
100                    "job.read(offset=0, limit=65536), advancing the byte offset, "
101                    "or job.save(path) for the complete artifact. "
102                    "Output past 1 MiB is automatically stored on disk (up to 16 MiB)."
103                )
104            result = Text(
105                self.capture.read_bytes(offset, size).decode("utf-8", errors="replace")
106            )
107        else:
108            result = read_page(self.capture, offset, limit)
109        if self.duration is not None:
110            self._read = True
111        return result
112
113    def save(self, path: str | os.PathLike[str]) -> Text:
114        """Atomically replace path with complete output bytes, or leave it untouched.
115
116        The path is on the capture's host (local in SSH command mode). Returns
117        its absolute path. Completion and full retention are required.
118        """
119        self._require_complete()
120        target = Path(path).expanduser().absolute()
121        target.parent.mkdir(parents=True, exist_ok=True)
122        fd, name = tempfile.mkstemp(dir=target.parent)
123        temporary = Path(name)
124        try:
125            with os.fdopen(fd, "wb") as file:
126                for offset in range(0, self.capture.seen, OUTPUT_PREVIEW):
127                    file.write(self.capture.read_bytes(offset, OUTPUT_PREVIEW))
128            os.replace(temporary, target)
129        finally:
130            temporary.unlink(missing_ok=True)
131        self._read = True
132        return Text(str(target))