8f43927575
examples/ is based on KasperskyLab/knp at 55470ed6 with local changes to the mnist-learn altai/blifat models, cuda-backend CMake and model_loader. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
54 lines
2.2 KiB
Python
54 lines
2.2 KiB
Python
"""Bounded parallel execution of independent real KNP runs, with separate outputs."""
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument('plan', type=Path)
|
|
ap.add_argument('--jobs', type=int, default=min(4, os.cpu_count() or 1))
|
|
ap.add_argument('--summary', type=Path, required=True)
|
|
a = ap.parse_args()
|
|
if not 1 <= a.jobs <= (os.cpu_count() or 1):
|
|
ap.error('--jobs must be between 1 and available logical CPUs')
|
|
plan = json.loads(a.plan.read_text())
|
|
jobs = plan['jobs']
|
|
names = [job['name'] for job in jobs]
|
|
if len(set(names)) != len(names) or not all(re.fullmatch(r'[A-Za-z0-9_-]+', name) for name in names):
|
|
ap.error('Run names must be unique and contain only letters, digits, underscores or hyphens')
|
|
root = Path('/home/int/run/cats-dogs')
|
|
if any((root / name).exists() for name in names):
|
|
ap.error('A planned run directory already exists; choose new names')
|
|
|
|
|
|
def execute(job):
|
|
if not isinstance(job['args'], list) or not all(isinstance(x, str) for x in job['args']):
|
|
raise ValueError('args must be a string list')
|
|
result = subprocess.run([sys.executable, str(Path(__file__).with_name('run.py')),
|
|
job['name'], *job['args']], capture_output=True, text=True)
|
|
row = {'name': job['name'], 'returncode': result.returncode}
|
|
if result.returncode:
|
|
row['error'] = result.stderr[-3000:]
|
|
else:
|
|
row['metrics'] = json.loads((root / job['name'] / 'metrics.json').read_text())
|
|
return row
|
|
|
|
|
|
started = time.monotonic()
|
|
rows = []
|
|
with ThreadPoolExecutor(max_workers=a.jobs) as pool:
|
|
for future in as_completed([pool.submit(execute, job) for job in jobs]):
|
|
row = future.result()
|
|
rows.append(row)
|
|
print(json.dumps(row), flush=True)
|
|
summary = {'jobs': a.jobs, 'logical_cpus': os.cpu_count(), 'runs': rows,
|
|
'elapsed_seconds': round(time.monotonic() - started, 3),
|
|
'plan': str(a.plan), 'all_succeeded': all(row['returncode'] == 0 for row in rows)}
|
|
a.summary.write_text(json.dumps(summary, indent=2) + '\n')
|
|
print('BATCH', json.dumps(summary), flush=True)
|
|
sys.exit(0 if summary['all_succeeded'] else 1)
|