Niccolò Nicolosi
188 lines
8.6 KiB
Python
188 lines
8.6 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from pathlib import Path
|
|
import sys
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools/pim/pimcomp/compare"))
|
|
from raptor_validation import experiment_scheduler as scheduler
|
|
import run_pimcomp_models as suite
|
|
import compare_raptor_pimcomp_model as comparator
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools/pim/ablation"))
|
|
import run_ablation as ablation
|
|
|
|
|
|
def task(root, key, code, dependencies=()):
|
|
command = [sys.executable, "-c", code]
|
|
return scheduler.ExperimentTask(key, lambda: command, dependencies, root / f"{key}.log")
|
|
|
|
|
|
def test_dependency_starts_before_unrelated_job_finishes(tmp_path):
|
|
# The slow job cannot finish until the fast job's dependent has run.
|
|
marker = tmp_path / "released"
|
|
slow = ("import pathlib,time; deadline=time.monotonic()+5; "
|
|
f"p=pathlib.Path({str(marker)!r}); "
|
|
"exec('while not p.exists() and time.monotonic()<deadline: time.sleep(.01)'); "
|
|
"assert p.exists(), 'dependent was blocked by an unrelated job'")
|
|
tasks = [task(tmp_path, "fast", "pass"), task(tmp_path, "slow", slow),
|
|
task(tmp_path, "dependent", f"from pathlib import Path; Path({str(marker)!r}).touch()", ("fast",))]
|
|
results = scheduler.run_tasks(tasks, 2, 1, tmp_path)
|
|
assert all(result.returncode == 0 for result in results.values())
|
|
|
|
|
|
def test_failure_skips_dependents_but_runs_other_jobs(tmp_path):
|
|
marker = tmp_path / "must-not-run"
|
|
tasks = [task(tmp_path, "failed", "raise SystemExit(3)"),
|
|
task(tmp_path, "dependent", f"from pathlib import Path; Path({str(marker)!r}).touch()", ("failed",)),
|
|
task(tmp_path, "independent", "pass")]
|
|
results = scheduler.run_tasks(tasks, 2, 1, tmp_path)
|
|
assert results["failed"].returncode == 3
|
|
assert results["dependent"].returncode != 0
|
|
assert results["independent"].returncode == 0
|
|
assert not marker.exists()
|
|
|
|
|
|
def test_thread_environment_and_execution_record(tmp_path):
|
|
names = ["OMP_NUM_THREADS", "OMP_THREAD_LIMIT", "RAYON_NUM_THREADS",
|
|
"CMAKE_BUILD_PARALLEL_LEVEL", "CARGO_BUILD_JOBS"]
|
|
job = task(tmp_path, "environment", f"import os,json; print(json.dumps({{k:os.environ[k] for k in {names!r}}}))")
|
|
scheduler.run_tasks([job], 1, 3, tmp_path)
|
|
assert set(json.loads(job.log_path.read_text()).values()) == {"3"}
|
|
record = json.loads(job.log_path.with_suffix(".execution.json").read_text())
|
|
assert record["threads_per_job"] == 3
|
|
assert scheduler.thread_environment(3)["OMP_WAIT_POLICY"] == "PASSIVE"
|
|
|
|
|
|
def test_concurrent_process_limit(tmp_path):
|
|
jobs = []
|
|
for index in range(6):
|
|
output = tmp_path / f"interval-{index}.json"
|
|
code = ("import json,time; from pathlib import Path; start=time.monotonic_ns(); "
|
|
f"time.sleep(.05); Path({str(output)!r}).write_text(json.dumps([start,time.monotonic_ns()]))")
|
|
jobs.append(task(tmp_path, str(index), code))
|
|
results = scheduler.run_tasks(jobs, 2, 1, tmp_path)
|
|
assert all(result.returncode == 0 for result in results.values())
|
|
events = []
|
|
for path in tmp_path.glob("interval-*.json"):
|
|
start, finish = json.loads(path.read_text())
|
|
events.extend(((start, 1), (finish, -1)))
|
|
active = 0
|
|
for _, delta in sorted(events):
|
|
active += delta
|
|
assert active <= 2
|
|
assert len(events) == 12 and active == 0
|
|
|
|
|
|
def test_interrupt_terminates_workers(tmp_path, monkeypatch):
|
|
processes = []
|
|
popen = scheduler.subprocess.Popen
|
|
|
|
def start(*args, **kwargs):
|
|
process = popen(*args, **kwargs)
|
|
processes.append(process)
|
|
return process
|
|
|
|
monkeypatch.setattr(scheduler.subprocess, "Popen", start)
|
|
sleep = scheduler.time.sleep
|
|
interrupted = False
|
|
|
|
def interrupt_once(seconds):
|
|
nonlocal interrupted
|
|
if not interrupted:
|
|
interrupted = True
|
|
raise KeyboardInterrupt()
|
|
sleep(seconds)
|
|
|
|
monkeypatch.setattr(scheduler.time, "sleep", interrupt_once)
|
|
with pytest.raises(KeyboardInterrupt):
|
|
scheduler.run_tasks([task(tmp_path, "waiting", "import time; time.sleep(30)")], 1, 1, tmp_path)
|
|
assert processes and all(process.poll() is not None for process in processes)
|
|
|
|
|
|
def test_invalid_dependency_graph_is_rejected(tmp_path):
|
|
with pytest.raises(ValueError, match="unknown dependency"):
|
|
scheduler.run_tasks([task(tmp_path, "a", "pass", ("missing",))], 1, 1, tmp_path)
|
|
with pytest.raises(ValueError, match="cycle"):
|
|
scheduler.run_tasks([task(tmp_path, "a", "pass", ("b",)),
|
|
task(tmp_path, "b", "pass", ("a",))], 1, 1, tmp_path)
|
|
|
|
|
|
def test_defaults_follow_processor_count(monkeypatch):
|
|
monkeypatch.setattr(suite, "available_processors", lambda: 7)
|
|
args = suite.parse_args([])
|
|
assert args.jobs == args.threads_per_job == 7
|
|
args = suite.parse_args(["--jobs", "2", "--threads-per-job", "3"])
|
|
assert (args.jobs, args.threads_per_job) == (2, 3)
|
|
with pytest.raises(SystemExit):
|
|
suite.parse_args(["--threads-per-job", "0"])
|
|
monkeypatch.setattr(ablation, "available_processors", lambda: 7)
|
|
args, forwarded = ablation.parse_args([])
|
|
assert args.jobs == args.threads_per_job == 7
|
|
assert forwarded[forwarded.index("--threads-per-job") + 1] == "7"
|
|
|
|
|
|
def test_combined_plan_reuses_references_and_raptor_baselines(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(suite, "validate_pimcomp_source", lambda: None)
|
|
args = suite.parse_args(["--out-dir", str(tmp_path / "comparison"),
|
|
"--ablation-out-dir", str(tmp_path / "ablation"), "--common-dir", str(tmp_path / "references"),
|
|
"--jobs", "2", "--threads-per-job", "2", "--dry-run"])
|
|
comparison = suite.make_plan(args)
|
|
ablations, _, _ = suite.combined_ablation_plans(args)
|
|
tasks = suite.build_tasks([comparison, *ablations])
|
|
assert len([job for job in tasks if job.key.startswith("reference:")]) == 5
|
|
assert len(tasks) == 5 + 40 + 24 # 16 baseline cases reuse comparison reports.
|
|
assert sum(len(plan.aliases) for plan in ablations) == 16
|
|
assert len({job.key for job in tasks}) == len(tasks)
|
|
for job in tasks:
|
|
if job.key.startswith("reference:"):
|
|
command = job.command()
|
|
assert command[command.index("--jobs") + 1] == "2"
|
|
continue
|
|
command = job.command()
|
|
assert command[command.index("--threads-per-job") + 1] == "2"
|
|
if "--reuse-pimcomp-report" in command:
|
|
assert len(job.dependencies) == 2
|
|
else:
|
|
assert len(job.dependencies) == 1
|
|
|
|
|
|
def test_raptor_receives_thread_limit(tmp_path, monkeypatch):
|
|
received = []
|
|
monkeypatch.setattr(comparator, "compile_with_raptor",
|
|
lambda *args, **kwargs: received.extend(kwargs["raptor_extra_args"]) or {})
|
|
args = argparse.Namespace(threads_per_job=3, pimcomp_config=tmp_path / "config.json",
|
|
raptor_extra_arg=[], raptor_path=tmp_path / "compiler", verbose_raptor_compile=False, timeout_seconds=0)
|
|
comparator.compile_raptor_target(tmp_path / "model.onnx", tmp_path / "output",
|
|
{"core_count": 1, "crossbar_count": 1, "crossbar_size": 128}, args, [])
|
|
assert "-j=3" in received
|
|
|
|
|
|
def test_stale_report_is_not_published_after_missing_execution(tmp_path):
|
|
args = suite.parse_args(["--out-dir", str(tmp_path), "--models", "vgg8",
|
|
"--archs", "arch-a", "--mode", "latency", "--raptor-only"])
|
|
plan = suite.make_plan(args)
|
|
report = suite.report_for(plan, plan.specs[0])
|
|
report.parent.mkdir(parents=True)
|
|
report.write_text(json.dumps({"validation_policy": suite.VALIDATION_POLICY,
|
|
"raptor_validation": {"status": "done", "passed": True},
|
|
"raptor_performance": {"latency_ms": 1}}))
|
|
tasks = suite.build_tasks([plan])
|
|
assert suite.finish_plan(plan, {task.key: suite.TaskResult(0) for task in tasks}) == 1
|
|
import csv
|
|
with (tmp_path / "results_comparison.csv").open() as stream:
|
|
row = next(csv.DictReader(stream))
|
|
assert row["raptor_latency_ms"] == row["raptor_functional_validation"] == "NA"
|
|
|
|
|
|
def test_reference_workers_do_not_create_nested_thread_teams(tmp_path, monkeypatch):
|
|
environments = []
|
|
monkeypatch.setattr(comparator, "save_inputs_to_files", lambda *a: ([], None))
|
|
monkeypatch.setattr(comparator, "run_logged", lambda *a, **k: environments.append(k["env"]))
|
|
comparator.generate_reference_outputs(tmp_path / "runner", tmp_path, tmp_path / "model.onnx", [], [],
|
|
argparse.Namespace(timeout_seconds=0), tmp_path / "output")
|
|
assert environments[0]["OMP_NUM_THREADS"] == environments[0]["OPENBLAS_NUM_THREADS"] == "1"
|