commit - release 1

src/kernbench/runtime_api/bench_runner.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from enum import Enum
+from typing import Any
+
+from kernbench.common.types import Completion, SimEngine, Trace
+
+from .context import RuntimeContext
+from .types import BenchResult, DeviceSelector
+
+
+class CompletionPolicy(str, Enum):
+    LAST_SUBMITTED = "last_submitted"
+    LAST_COMPLETED = "last_completed"  # requires trace/timestamps or engine support; stub for now
+    ALL_OK_FAIL_FAST = "all_ok_fail_fast"
+
+
+BenchFn = Callable[[RuntimeContext], Any]
+EngineFactory = Callable[[object, DeviceSelector], SimEngine]
+
+
+def run_bench(
+    *,
+    topology: object,
+    bench_fn: BenchFn,
+    device: DeviceSelector,
+    engine_factory: EngineFactory,
+    correlation_id: str = "bench0",
+    completion_policy: CompletionPolicy = CompletionPolicy.LAST_SUBMITTED,
+) -> BenchResult:
+    """
+    Minimal bench runner.
+
+    - topology: compiled topology object (opaque to runtime here)
+    - bench_fn: callable that receives RuntimeContext and submits requests
+    - device: DeviceSelector ("all" or "sip:<N>")
+    - engine_factory: builds sim_engine for given topology & device
+    - completion_policy: how to determine overall completion/result
+    """
+    engine = engine_factory(topology, device)
+    # Extract spec from TopologyHandle or TopologyGraph
+    topo_obj = getattr(topology, "topology_obj", topology)
+    spec = getattr(topo_obj, "spec", None)
+    ctx = RuntimeContext(
+        engine=engine, target_device=device,
+        correlation_id=correlation_id, spec=spec,
+    )
+
+    bench_fn(ctx)
+
+    ctx.wait_all()
+
+    collected_traces = ctx._traces or None
+
+    handles = ctx.handles()
+    if not handles:
+        return BenchResult(
+            completion=Completion(
+                ok=False, error_code="NO_REQUESTS", error_message="Bench submitted no requests"
+            ),
+            correlation_id=correlation_id,
+            trace=None,
+            traces=collected_traces,
+        )
+
+    if completion_policy == CompletionPolicy.LAST_SUBMITTED:
+        last = handles[-1]
+        completion, trace = engine.get_completion(last)
+        return BenchResult(
+            completion=completion, correlation_id=correlation_id,
+            trace=trace, traces=collected_traces,
+        )
+
+    if completion_policy == CompletionPolicy.ALL_OK_FAIL_FAST:
+        last_trace: Trace | None = None
+        for h in handles:
+            c, t = engine.get_completion(h)
+            last_trace = t if t is not None else last_trace
+            if not c.ok:
+                return BenchResult(
+                    completion=c, correlation_id=correlation_id,
+                    trace=last_trace, traces=collected_traces,
+                )
+        return BenchResult(
+            completion=Completion(ok=True), correlation_id=correlation_id,
+            trace=last_trace, traces=collected_traces,
+        )
+
+    # LAST_COMPLETED placeholder (needs engine support for timing). Fall back.
+    last = handles[-1]
+    completion, trace = engine.get_completion(last)
+    return BenchResult(
+        completion=completion, correlation_id=correlation_id,
+        trace=trace, traces=collected_traces,
+    )