Spaces:
Running
Running
File size: 5,268 Bytes
e2d54c9 | 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 159 160 161 162 163 164 165 166 | """Evidence writers.
Every stage writes machine-readable raw output under `.openresearch/artifacts/`
so a run is inspectable after the fact, and prints the same numbers to stdout
because in local orx mode the run log is the only evidence channel.
"""
from __future__ import annotations
import csv
import hashlib
import json
import os
import platform
import subprocess
import time
from pathlib import Path
from typing import Any, Iterable
ARTIFACTS = Path(".openresearch/artifacts")
def artifact_dir(*parts: str) -> Path:
d = ARTIFACTS.joinpath(*parts)
d.mkdir(parents=True, exist_ok=True)
return d
def write_json(path: Path, obj: Any) -> Path:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(obj, indent=2, sort_keys=True, default=_default) + "\n")
_emit(path)
return path
def write_csv(path: Path, rows: Iterable[dict], fieldnames: list[str] | None = None) -> Path:
rows = list(rows)
if not rows:
return path
path.parent.mkdir(parents=True, exist_ok=True)
fieldnames = fieldnames or list(rows[0].keys())
with path.open("w", newline="") as fh:
w = csv.DictWriter(fh, fieldnames=fieldnames)
w.writeheader()
w.writerows(rows)
_emit(path)
return path
# --------------------------------------------------------------------------- #
# Lossless artifact channel.
#
# In orx local mode `orx artifacts` does not exist: files written under
# .openresearch/artifacts/ inside a compute job are unreachable once the job ends,
# and the run log is the only wire out. So every artifact is also emitted to stdout
# between framed markers carrying its exact byte count and SHA-256. A consumer can
# reconstruct the file byte-for-byte and verify it, which means a truncated or
# interleaved log fails the hash check loudly instead of yielding plausible-looking
# wrong numbers.
#
# Emission is base64 of the raw bytes: artifacts contain newlines, and a payload that
# cannot collide with the marker lines is what makes the frame unambiguous.
EMIT_BEGIN = "<<<ORX-ARTIFACT"
EMIT_END = ">>>ORX-ARTIFACT-END"
EMIT_LIMIT = int(os.environ.get("ORX_ARTIFACT_EMIT_LIMIT", 4_000_000))
def _emit(path: Path) -> None:
import base64
raw = path.read_bytes()
rel = path.as_posix()
digest = hashlib.sha256(raw).hexdigest()
if len(raw) > EMIT_LIMIT:
# Never truncate silently: say so, and say exactly how much was withheld.
print(f"{EMIT_BEGIN} path={rel} bytes={len(raw)} sha256={digest} "
f"encoding=none status=OMITTED-TOO-LARGE limit={EMIT_LIMIT}")
print(f"{EMIT_END} path={rel}")
return
payload = base64.b64encode(raw).decode("ascii")
print(f"{EMIT_BEGIN} path={rel} bytes={len(raw)} sha256={digest} "
f"encoding=base64 status=OK")
for i in range(0, len(payload), 76):
print(payload[i:i + 76])
print(f"{EMIT_END} path={rel}")
def _default(o: Any) -> Any:
import numpy as np
if isinstance(o, (np.integer,)):
return int(o)
if isinstance(o, (np.floating,)):
return float(o)
if isinstance(o, (np.bool_,)):
return bool(o)
if isinstance(o, np.ndarray):
return o.tolist()
return str(o)
def git_sha() -> str:
for cmd in (["git", "rev-parse", "HEAD"],):
try:
return subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL).strip()
except Exception:
pass
return os.environ.get("GIT_COMMIT", "unknown")
def cpu_info() -> dict:
"""CPU allocation actually granted to this run (recorded in every artifact)."""
info = {
"platform": platform.platform(),
"machine": platform.machine(),
"processor": platform.processor(),
"python": platform.python_version(),
"os_cpu_count": os.cpu_count(),
}
try:
info["sched_getaffinity"] = len(os.sched_getaffinity(0)) # type: ignore[attr-defined]
except (AttributeError, OSError):
info["sched_getaffinity"] = None
try:
with open("/proc/cpuinfo") as fh:
models = [l.split(":", 1)[1].strip() for l in fh if l.startswith("model name")]
if models:
info["model_name"] = models[0]
info["n_logical_from_cpuinfo"] = len(models)
except OSError:
pass
for path, key in (
("/sys/fs/cgroup/cpu.max", "cgroup_cpu_max"),
("/sys/fs/cgroup/memory.max", "cgroup_memory_max"),
):
try:
info[key] = Path(path).read_text().strip()
except OSError:
pass
return info
class Section:
"""Times a stage, prints a banner, and returns runtime for the artifact."""
def __init__(self, title: str) -> None:
self.title = title
def __enter__(self) -> "Section":
self.t0 = time.time()
print(f"\n{'=' * 78}\n== {self.title}\n{'=' * 78}", flush=True)
return self
def __exit__(self, *exc: Any) -> None:
self.seconds = time.time() - self.t0
print(f"-- {self.title}: {self.seconds:.2f}s", flush=True)
def banner(msg: str) -> None:
print(f"\n--- {msg}", flush=True)
def kv(label: str, value: Any, width: int = 46) -> None:
print(f" {label:<{width}} {value}", flush=True)
|