File size: 5,674 Bytes
9422383
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
#!/usr/bin/env python3
"""Record review evidence around an unchanged Ultralytics benchmark command.

Ultralytics remains responsible for loading the dataset, running validation,
and reporting accuracy and stage timings. This command retains that output with
the measured runtime and model identities, source revision, and Raspberry Pi state needed to
interpret the result, then exits with the benchmark's status.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import resource
import shutil
import subprocess
import sys
import tempfile
from collections.abc import Mapping, Sequence
from datetime import UTC, datetime
from pathlib import Path

from .host_metrics import capture_host_snapshot

__all__ = ["record_benchmark"]


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for block in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def _run(command: Sequence[str], log_path: Path) -> tuple[int, dict[str, object]]:
    """Run once, preserving output and descriptive state from both boundaries."""
    started_at = datetime.now(UTC)
    before = capture_host_snapshot()
    with log_path.open("x", encoding="utf-8") as log:
        process = subprocess.Popen(
            command,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            encoding="utf-8",
            errors="replace",
            bufsize=1,
        )
        if process.stdout is None:
            raise RuntimeError("benchmark output stream is unavailable")
        for line in process.stdout:
            sys.stdout.write(line)
            sys.stdout.flush()
            log.write(line)
            log.flush()
        return_code = process.wait()
    after = capture_host_snapshot()
    finished_at = datetime.now(UTC)
    child_usage = resource.getrusage(resource.RUSAGE_CHILDREN)
    report: dict[str, object] = {
        "version": 1,
        "kind": "yolo26n-ultralytics-benchmark-result",
        "command": list(command),
        "started_at": started_at.isoformat(),
        "finished_at": finished_at.isoformat(),
        "exit_code": return_code,
        # Linux reports ru_maxrss in KiB. The fresh recorder process has only
        # one child command, so this is the benchmark process tree's peak.
        "child_peak_rss_bytes": child_usage.ru_maxrss * 1024,
        "host": {"before": before.to_json(), "after": after.to_json()},
    }
    return return_code, report


def _write_json(path: Path, value: Mapping[str, object]) -> None:
    with path.open("x", encoding="utf-8") as stream:
        json.dump(value, stream, indent=2, sort_keys=True)
        stream.write("\n")


def record_benchmark(
    output: Path,
    runtime_library: Path,
    model_manifest: Path,
    revision: str,
    command: Sequence[str],
) -> int:
    """Publish one complete measurement directory without replacing prior evidence."""
    if output.exists():
        raise FileExistsError(f"output already exists: {output}")
    if not runtime_library.is_file():
        raise FileNotFoundError(f"runtime library does not exist: {runtime_library}")
    if not model_manifest.is_file():
        raise FileNotFoundError(f"model manifest does not exist: {model_manifest}")
    manifest_document = json.loads(model_manifest.read_text(encoding="utf-8"))
    if not isinstance(manifest_document, Mapping):
        raise ValueError("model manifest must be a JSON object")

    output.parent.mkdir(parents=True, exist_ok=True)
    staging = Path(tempfile.mkdtemp(prefix=f".{output.name}.", dir=output.parent))
    try:
        return_code, result = _run(command, staging / "benchmark.log")
        result["source_revision"] = revision
        result["runtime_library"] = {
            "path": str(runtime_library),
            "bytes": runtime_library.stat().st_size,
            "sha256": _sha256(runtime_library),
        }
        result["model_manifest"] = {
            "path": str(model_manifest),
            "sha256": _sha256(model_manifest),
            "document": manifest_document,
        }
        _write_json(staging / "results.json", result)
        staging.rename(output)
    finally:
        shutil.rmtree(staging, ignore_errors=True)
    return return_code


def _main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--runtime-library", type=Path, required=True)
    parser.add_argument("--model-manifest", type=Path, required=True)
    parser.add_argument("--revision", required=True)
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    command = args.command[1:] if args.command[:1] == ["--"] else args.command
    if not command:
        parser.error("a benchmark command is required after --")
    revision = args.revision.strip()
    if not revision or "\n" in revision:
        parser.error("--revision must be one non-empty line")
    output = args.output_dir.expanduser().resolve()
    if output.exists():
        parser.error(f"output already exists: {output}")

    return_code = record_benchmark(
        output,
        args.runtime_library.expanduser().resolve(),
        args.model_manifest.expanduser().resolve(),
        revision,
        command,
    )
    print(f"benchmark evidence: {output}")
    if return_code:
        raise SystemExit(return_code)


if __name__ == "__main__":
    try:
        _main()
    except (OSError, RuntimeError, ValueError) as error:
        raise SystemExit(f"error: {error}") from None