parallel validation jobs

update readme

validation/validate.py

@@ -1,10 +1,15 @@
 #!/usr/bin/env python3
 
 import argparse
+import os
 import signal
 import subprocess
 import sys
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from contextlib import redirect_stderr, redirect_stdout
+from itertools import groupby
 from pathlib import Path
+from tempfile import TemporaryDirectory
 from colorama import Style, Fore
 from raptor_validation.validate_one import (
     MODE_COMPILE_ONLY,
@@ -28,6 +33,49 @@ DEFAULT_PIMCOMP_CONFIG = "arch-a"
 PIMCOMP_CONFIG_CHOICES = ("arch-a", "arch-b", "arch-c")
 
 
+def run_validation_job(job):
+    models, options = job
+    completed = []
+    for index, onnx_path, rel, log_path in models:
+        reporter = ProgressReporter(
+            options["model_total"],
+            stages_per_model=len(MODE_STAGE_TITLES[options["mode"]]),
+            enabled=False,
+            verbose=options["verbose"],
+        )
+        sys.stdout.flush()
+        sys.stderr.flush()
+        saved_stdout = os.dup(1)
+        saved_stderr = os.dup(2)
+        try:
+            with open(log_path, "w", encoding="utf-8", buffering=1) as log:
+                os.dup2(log.fileno(), 1)
+                os.dup2(log.fileno(), 2)
+                with redirect_stdout(log), redirect_stderr(log):
+                    try:
+                        result = validate_network(
+                            onnx_path,
+                            reporter=reporter,
+                            model_index=index,
+                            model_total=options["model_total"],
+                            verbose=options["verbose"],
+                            mode=options["mode"],
+                            **options["validation_kwargs"],
+                        )
+                    except Exception as exc:
+                        result = ValidationResult(False, pimsim_status=PIMSIM_NOT_RUN)
+                        print_validation_error(reporter, rel, exc)
+                    finally:
+                        reporter.finish()
+        finally:
+            os.dup2(saved_stdout, 1)
+            os.dup2(saved_stderr, 2)
+            os.close(saved_stdout)
+            os.close(saved_stderr)
+        completed.append((str(rel), result, str(log_path)))
+    return completed
+
+
 def format_return_status(returncode):
     if returncode < 0:
         signal_num = -returncode
@@ -111,6 +159,8 @@ def main():
                     help="Additional argument to pass through to the Raptor compiler. Repeat as needed.")
     ap.add_argument("--command-timeout-seconds", type=float, default=1000000.0,
                     help="Per-subprocess timeout in seconds for compiler, runner, and simulation commands.")
+    ap.add_argument("-j", "--jobs", type=int, default=os.cpu_count() or 1,
+                    help="Number of model validations to run in parallel (default: all available CPUs).")
     ap.add_argument("--clean", action="store_true",
                     help="Remove generated validation artifacts under each model workspace and exit.")
     mode_group = ap.add_mutually_exclusive_group()
@@ -123,6 +173,8 @@ def main():
     ap.add_argument("--verbose", action="store_true",
                     help="Print per-stage progress and subprocess logs for passing validations too.")
     a = ap.parse_args()
+    if a.jobs < 1:
+        ap.error("--jobs must be at least 1")
 
     operations_dir = Path(a.operations_dir).resolve() if a.operations_dir else script_dir / "operations"
     simulator_dir = Path(a.simulator_dir).resolve() if a.simulator_dir else (
@@ -161,6 +213,7 @@ def main():
 
     print(Style.BRIGHT + f"Found {len(onnx_files)} ONNX file(s) to validate." + Style.RESET_ALL)
     print(f"Operations root: {operations_dir}")
+    print(f"Parallel jobs: {a.jobs}")
     print("=" * 72)
 
     mode = MODE_FULL
@@ -193,43 +246,72 @@ def main():
     pass_timing_counts = {label: 0 for _, label in PIM_PASS_LABELS}
     total_timing_sum = 0.0
     timed_benchmark_count = 0
-    reporter = ProgressReporter(len(onnx_files), stages_per_model=len(MODE_STAGE_TITLES[mode]), verbose=a.verbose)
-    for index, onnx_path in enumerate(onnx_files, start=1):
-        rel = onnx_path.relative_to(operations_dir)
-        try:
-            result = validate_network(
-                onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
-                crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
-                raptor_extra_args=a.raptor_extra_arg,
-                pimsim_nn_build_dir=pimsim_nn_build_dir,
-                pimsim_config_path=selected_pimsim_config,
-                command_timeout_seconds=a.command_timeout_seconds,
-                threshold=a.threshold,
-                rtol=a.relative_threshold,
-                seed=a.seed,
-                reporter=reporter,
-                model_index=index,
-                model_total=len(onnx_files),
-                verbose=a.verbose,
-                mode=mode,
+    reporter = ProgressReporter(len(onnx_files), stages_per_model=1, verbose=a.verbose)
+    validation_kwargs = {
+        "raptor_path": a.raptor_path,
+        "onnx_include_dir": a.onnx_include_dir,
+        "simulator_dir": simulator_dir,
+        "crossbar_size": a.crossbar_size,
+        "crossbar_count": a.crossbar_count,
+        "core_count": a.core_count,
+        "raptor_extra_args": a.raptor_extra_arg,
+        "pimsim_nn_build_dir": pimsim_nn_build_dir,
+        "pimsim_config_path": selected_pimsim_config,
+        "command_timeout_seconds": a.command_timeout_seconds,
+        "threshold": a.threshold,
+        "rtol": a.relative_threshold,
+        "seed": a.seed,
+    }
+    indexed_files = list(enumerate(onnx_files, start=1))
+    workspace_groups = [
+        list(group)
+        for _, group in groupby(indexed_files, key=lambda indexed_path: indexed_path[1].parent)
+    ]
+    with TemporaryDirectory(prefix="raptor-validation-") as log_dir:
+        jobs = [
+            (
+                [
+                    (
+                        index,
+                        onnx_path,
+                        onnx_path.relative_to(operations_dir),
+                        Path(log_dir) / f"{index}.log",
+                    )
+                    for index, onnx_path in workspace_group
+                ],
+                {
+                    "model_total": len(onnx_files),
+                    "mode": mode,
+                    "verbose": a.verbose,
+                    "validation_kwargs": validation_kwargs,
+                },
             )
-            results[str(rel)] = result
-            if result.pim_pass_timings:
-                benchmark_total = 0.0
-                for label, duration in result.pim_pass_timings.items():
-                    pass_timing_sums[label] += duration
-                    pass_timing_counts[label] += 1
-                    benchmark_total += duration
-                total_timing_sum += benchmark_total
-                timed_benchmark_count += 1
-        except subprocess.CalledProcessError as exc:
-            results[str(rel)] = ValidationResult(False, pimsim_status=PIMSIM_NOT_RUN)
-            print_validation_error(reporter, rel, exc)
-        except Exception as exc:
-            results[str(rel)] = ValidationResult(False, pimsim_status=PIMSIM_NOT_RUN)
-            print_validation_error(reporter, rel, exc)
+            for workspace_group in workspace_groups
+        ]
+        with ProcessPoolExecutor(max_workers=a.jobs) as executor:
+            futures = [executor.submit(run_validation_job, job) for job in jobs]
+            for future in as_completed(futures):
+                completed_group = future.result()
+                for rel, result, log_path in completed_group:
+                    reporter.suspend()
+                    output = Path(log_path).read_text(encoding="utf-8", errors="replace")
+                    if output:
+                        print(output, end="" if output.endswith("\n") else "\n")
+                    reporter.resume()
+                    reporter.advance()
+                    reporter.record_result(result.passed)
+                    results[rel] = result
+                    if result.pim_pass_timings:
+                        benchmark_total = 0.0
+                        for label, duration in result.pim_pass_timings.items():
+                            pass_timing_sums[label] += duration
+                            pass_timing_counts[label] += 1
+                            benchmark_total += duration
+                        total_timing_sum += benchmark_total
+                        timed_benchmark_count += 1
 
     reporter.finish()
+    results = dict(sorted(results.items()))
 
     # Summary
     n_passed = sum(1 for result in results.values() if result.passed)