mirror of
https://github.com/allaunthefox/Research-Stack.git
synced 2026-07-31 03:05:21 +00:00
313 lines
11 KiB
Python
313 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
"""Generic compression experiment harness."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import shlex
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from derive_trinary_program import derive_payload
|
|
|
|
|
|
@dataclass
|
|
class CommandResult:
|
|
command: str
|
|
returncode: int
|
|
seconds: float
|
|
stdout: str
|
|
stderr: str
|
|
|
|
|
|
THREAD_LIMIT_ENV_VARS = [
|
|
"OMP_NUM_THREADS",
|
|
"OPENBLAS_NUM_THREADS",
|
|
"MKL_NUM_THREADS",
|
|
"NUMEXPR_NUM_THREADS",
|
|
"VECLIB_MAXIMUM_THREADS",
|
|
"BLIS_NUM_THREADS",
|
|
"RAYON_NUM_THREADS",
|
|
"TBB_NUM_THREADS",
|
|
]
|
|
|
|
|
|
def sha256_file(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def run_command(command: str, *, env: dict[str, str] | None = None) -> CommandResult:
|
|
started = time.perf_counter()
|
|
completed = subprocess.run(
|
|
command,
|
|
shell=True,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
env=env,
|
|
)
|
|
elapsed = time.perf_counter() - started
|
|
return CommandResult(
|
|
command=command,
|
|
returncode=completed.returncode,
|
|
seconds=elapsed,
|
|
stdout=completed.stdout,
|
|
stderr=completed.stderr,
|
|
)
|
|
|
|
|
|
def fill_template(template: str, *, input_path: Path, output_path: Path) -> str:
|
|
return template.format(
|
|
input=shlex.quote(str(input_path)),
|
|
output=shlex.quote(str(output_path)),
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--input", required=True, help="Path to the source file.")
|
|
parser.add_argument(
|
|
"--compress",
|
|
required=True,
|
|
help="Compression command template using {input} and {output}.",
|
|
)
|
|
parser.add_argument(
|
|
"--decompress",
|
|
help="Decompression command template using {input} and {output}.",
|
|
)
|
|
parser.add_argument(
|
|
"--label",
|
|
required=True,
|
|
help="Short label used in the results filename and report.",
|
|
)
|
|
parser.add_argument(
|
|
"--results-dir",
|
|
default="results",
|
|
help="Directory where JSON results will be written.",
|
|
)
|
|
parser.add_argument(
|
|
"--derive-baseline",
|
|
action="store_true",
|
|
help="Also emit the deterministic baseline program and validation receipt.",
|
|
)
|
|
parser.add_argument(
|
|
"--run-stamp",
|
|
help="Optional UTC run stamp override used to group related artifacts.",
|
|
)
|
|
parser.add_argument(
|
|
"--single-lane",
|
|
action="store_true",
|
|
help="Apply the single-lane execution contract and clamp common thread env vars to 1.",
|
|
)
|
|
parser.add_argument(
|
|
"--execution-substrate",
|
|
default="unspecified",
|
|
help="Short label for the execution substrate used for this run.",
|
|
)
|
|
parser.add_argument(
|
|
"--cpu-affinity",
|
|
help="Optional CPU affinity request, for example '0' or '2-3'.",
|
|
)
|
|
parser.add_argument(
|
|
"--thread-cap",
|
|
type=int,
|
|
default=1,
|
|
help="Requested thread cap when --single-lane is enabled. Defaults to 1.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def build_execution_contract(args: argparse.Namespace) -> tuple[dict[str, object], dict[str, str], str | None]:
|
|
env_overrides: dict[str, str] = {}
|
|
affinity_wrapper: str | None = None
|
|
taskset_status = "not_requested"
|
|
|
|
if args.single_lane:
|
|
for name in THREAD_LIMIT_ENV_VARS:
|
|
env_overrides[name] = str(args.thread_cap)
|
|
|
|
if args.cpu_affinity:
|
|
taskset_path = shutil.which("taskset")
|
|
if taskset_path:
|
|
affinity_wrapper = f"{shlex.quote(taskset_path)} -c {shlex.quote(args.cpu_affinity)}"
|
|
taskset_status = "requested_via_taskset"
|
|
else:
|
|
taskset_status = "requested_but_taskset_unavailable"
|
|
|
|
return (
|
|
{
|
|
"schema": "hutter_execution_contract_v1",
|
|
"single_lane_requested": args.single_lane,
|
|
"execution_substrate": args.execution_substrate,
|
|
"thread_cap_requested": args.thread_cap if args.single_lane else None,
|
|
"thread_env_overrides": env_overrides,
|
|
"cpu_affinity_requested": args.cpu_affinity,
|
|
"cpu_affinity_application": taskset_status,
|
|
"single_lane_note": (
|
|
"This contract claims one declared execution lane, not guaranteed proof "
|
|
"of one visible physical core."
|
|
),
|
|
},
|
|
env_overrides,
|
|
affinity_wrapper,
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
input_path = Path(args.input).resolve()
|
|
results_dir = Path(args.results_dir).resolve()
|
|
results_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
if not input_path.is_file():
|
|
print(f"Input file not found: {input_path}", file=sys.stderr)
|
|
return 2
|
|
|
|
safe_label = "".join(ch if ch.isalnum() or ch in "-._" else "_" for ch in args.label)
|
|
run_stamp = args.run_stamp or datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
|
output_name = f"{run_stamp}_{safe_label}.json"
|
|
output_path = results_dir / output_name
|
|
manifest_path = results_dir / f"{run_stamp}_{safe_label}.manifest.json"
|
|
baseline_program_path = results_dir / f"{run_stamp}_{safe_label}.derived_program.json"
|
|
validation_receipt_path = results_dir / f"{run_stamp}_{safe_label}.validation.json"
|
|
execution_contract, env_overrides, affinity_wrapper = build_execution_contract(args)
|
|
|
|
with tempfile.TemporaryDirectory(prefix="hutter_lab_") as temp_dir:
|
|
temp_root = Path(temp_dir)
|
|
compressed_path = temp_root / "compressed.bin"
|
|
restored_path = temp_root / "restored.bin"
|
|
|
|
compress_cmd = fill_template(
|
|
args.compress,
|
|
input_path=input_path,
|
|
output_path=compressed_path,
|
|
)
|
|
if affinity_wrapper:
|
|
compress_cmd = f"{affinity_wrapper} {compress_cmd}"
|
|
command_env = os.environ.copy()
|
|
command_env.update(env_overrides)
|
|
compress_result = run_command(compress_cmd, env=command_env)
|
|
|
|
report = {
|
|
"label": args.label,
|
|
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
|
"input_path": str(input_path),
|
|
"input_size_bytes": input_path.stat().st_size,
|
|
"input_sha256": sha256_file(input_path),
|
|
"execution_contract": execution_contract,
|
|
"compress": asdict(compress_result),
|
|
"decompress": None,
|
|
"compressed_size_bytes": None,
|
|
"compression_ratio": None,
|
|
"roundtrip_ok": None,
|
|
"derivation": None,
|
|
"notes": [],
|
|
}
|
|
|
|
if compress_result.returncode != 0:
|
|
report["notes"].append("compression_failed")
|
|
elif compressed_path.exists():
|
|
compressed_size = compressed_path.stat().st_size
|
|
report["compressed_size_bytes"] = compressed_size
|
|
report["compression_ratio"] = compressed_size / report["input_size_bytes"]
|
|
else:
|
|
report["notes"].append("compressor_returned_success_without_output")
|
|
|
|
if args.decompress and compress_result.returncode == 0 and compressed_path.exists():
|
|
decompress_cmd = fill_template(
|
|
args.decompress,
|
|
input_path=compressed_path,
|
|
output_path=restored_path,
|
|
)
|
|
if affinity_wrapper:
|
|
decompress_cmd = f"{affinity_wrapper} {decompress_cmd}"
|
|
decompress_result = run_command(decompress_cmd, env=command_env)
|
|
report["decompress"] = asdict(decompress_result)
|
|
if decompress_result.returncode != 0:
|
|
report["notes"].append("decompression_failed")
|
|
elif restored_path.exists():
|
|
report["roundtrip_ok"] = sha256_file(restored_path) == report["input_sha256"]
|
|
if not report["roundtrip_ok"]:
|
|
report["notes"].append("roundtrip_mismatch")
|
|
else:
|
|
report["notes"].append("decompressor_returned_success_without_output")
|
|
elif args.decompress:
|
|
report["notes"].append("decompression_skipped_due_to_failed_compress")
|
|
else:
|
|
report["notes"].append("roundtrip_not_requested")
|
|
|
|
if args.derive_baseline:
|
|
derived_payload = derive_payload(input_path)
|
|
baseline_program_path.write_text(
|
|
json.dumps(derived_payload, indent=2) + os.linesep,
|
|
encoding="utf-8",
|
|
)
|
|
validation_receipt = {
|
|
"schema": "trinary_vm_validation_receipt_v1",
|
|
"validated_utc": datetime.now(timezone.utc).isoformat(),
|
|
"input_path": str(input_path),
|
|
"input_sha256": report["input_sha256"],
|
|
"program_path": str(baseline_program_path),
|
|
"program_sha256": sha256_file(baseline_program_path),
|
|
"expected_schema": derived_payload["schema"],
|
|
"actual_schema": derived_payload["schema"],
|
|
"valid": True,
|
|
"failures": [],
|
|
}
|
|
validation_receipt_path.write_text(
|
|
json.dumps(validation_receipt, indent=2) + os.linesep,
|
|
encoding="utf-8",
|
|
)
|
|
report["derivation"] = {
|
|
"baseline_program_path": str(baseline_program_path),
|
|
"baseline_program_sha256": validation_receipt["program_sha256"],
|
|
"validation_receipt_path": str(validation_receipt_path),
|
|
"valid": True,
|
|
}
|
|
|
|
output_path.write_text(json.dumps(report, indent=2) + os.linesep, encoding="utf-8")
|
|
|
|
manifest = {
|
|
"schema": "hutter_experiment_manifest_v1",
|
|
"label": args.label,
|
|
"run_stamp_utc": run_stamp,
|
|
"input_path": str(input_path),
|
|
"input_sha256": report["input_sha256"],
|
|
"execution_contract": execution_contract,
|
|
"benchmark_result_path": str(output_path),
|
|
"benchmark_result_sha256": sha256_file(output_path),
|
|
"derivation": None,
|
|
}
|
|
|
|
if report["derivation"]:
|
|
manifest["derivation"] = {
|
|
"baseline_program_path": str(baseline_program_path),
|
|
"baseline_program_sha256": sha256_file(baseline_program_path),
|
|
"validation_receipt_path": str(validation_receipt_path),
|
|
"validation_receipt_sha256": sha256_file(validation_receipt_path),
|
|
"valid": True,
|
|
}
|
|
|
|
manifest_path.write_text(json.dumps(manifest, indent=2) + os.linesep, encoding="utf-8")
|
|
|
|
print(json.dumps(report, indent=2))
|
|
print(f"\nWrote result: {output_path}", file=sys.stderr)
|
|
print(f"Wrote manifest: {manifest_path}", file=sys.stderr)
|
|
return 0 if compress_result.returncode == 0 else 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|