pcb-afm-simulation-system/scripts/test_p3_adaptive_execution.py

"""P3-M6 regression: web AdaptiveLoop closed loop + executor bridge.

Validates the complete adaptive closed loop on the web side (AdaptiveLoop),
including the new submit_batch_to_executor bridge that wraps the current
pending batch into an adaptive_batch task for the local executor:

    fake plan -> initialize_search (initial batch)
    -> submit_batch_to_executor (create task)
    -> fake executor runs the task and reports results
    -> report_results feeds the search back
    -> get_next_batch -> ... -> until budget exhausted / converged

Run:  python scripts/test_p3_adaptive_execution.py   (exit 0 = PASS)

Uses an isolated temp SQLite DB and KIMI_API_KEY="" so no AI backend is
contacted. No real Motor-CAD is involved.
"""
import json
import os
import sys
import tempfile

_TMP = os.path.join(tempfile.mkdtemp(), "test_afm.db")
os.environ["AFM_DB_PATH"] = _TMP
os.environ["KIMI_API_KEY"] = ""
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "web", "backend"))
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from app.database import init_db  # noqa: E402

init_db()

from app.services.adaptive_loop import create_loop  # noqa: E402
from app.services.task_manager import get_task_manager  # noqa: E402

tm = get_task_manager()

# ------------------------------------------------------------------ 1. create
loop = create_loop(
    user_requirement="maximize average torque within feasibility constraints",
    total_budget=8,
    batch_size=4,
)
# Inject a fake plan to bypass AI plan generation (parameters are the L0-known
# names so the feasibility pre-screening matches).
loop.plan = {
    "plan_name": "p3-fake",
    "topology": "SSSR",
    "scan_variables": [
        {"name": "airgap_mm", "min_value": 0.8, "max_value": 2.0, "step": 0.1, "unit": "mm"},
        {"name": "current_a", "min_value": 5.0, "max_value": 20.0, "step": 0.5, "unit": "A"},
    ],
    "search_strategy": {"max_solver_calls": 8, "batch_size": 4, "initial_samples": 4},
    "acceptance_criteria": {
        "objective_metric": "tavg_nm",
        "objective_direction": "maximize",
        "hard_constraints": [],
    },
}

# ---------------------------------------------------------- 2. init search
res = loop.initialize_search()
assert res["phase"] == "search_initialized", res
init_points = res["initial_batch"]
assert len(init_points) > 0, res
init_ids = sorted(p["id"] for p in init_points)
print("[1] initialize_search OK: initial_batch=%d ids=%s" % (len(init_points), init_ids))

# --------------------------------------------- 3. submit-batch (new bridge)
sub = loop.submit_batch_to_executor()
tid = sub["task_id"]
assert tid, sub
assert sub["n_points"] == len(init_points), sub
assert loop.phase.value == "simulation_running", loop.phase
task = tm.get_task(tid)
assert task is not None, tid
assert task["task_type"] == "adaptive_batch", task
assert task["loop_id"] == loop.loop_id, task
assert sorted(task["point_ids"]) == init_ids, (task["point_ids"], init_ids)
print("[2] submit_batch_to_executor OK: task=%s points=%d type=%s"
      % (tid, sub["n_points"], task["task_type"]))


def run_batch(tid):
    """Fake local executor: read the task parameters and fabricate metrics."""
    task = tm.get_task(tid)
    with open(task["task_file"], "r", encoding="utf-8") as f:
        payload = json.load(f)
    results = []
    for i, params in enumerate(payload["parameters"]):
        pid = params.get("point_id")
        results.append({
            "point_id": pid,
            "point_index": i,
            "params": params,
            "metrics": {"tavg_nm": round(8.0 + 0.5 * pid, 3), "efficiency_pct": 90.0 + (pid % 5)},
            "status": "OK",
        })
    tm.report_results(tid, results, status="completed")
    return results


# --------------------------------------------------------- 4. drive the loop
steps = 0
max_steps = 12
batches = set()
all_tids = []
while steps < max_steps:
    steps += 1
    sub = loop.submit_batch_to_executor()
    if not sub["task_id"]:
        # nothing pending: ask the search for the next batch
        nb = loop.get_next_batch()
        if not nb.get("points"):
            break
        sub = loop.submit_batch_to_executor()
        if not sub["task_id"]:
            break
    batches.add(sub["batch_id"])
    all_tids.append(sub["task_id"])
    results = run_batch(sub["task_id"])
    point_results = [
        {"point_id": r["point_id"], "metrics": r["metrics"], "status": "ok"}
        for r in results
    ]
    loop.report_results(point_results)
    print("[3] step %d batch=%s n_points=%d phase=%s"
          % (steps, sub["batch_id"], len(point_results), loop.phase.value))
    comp = loop.check_completion()
    if comp["completed"]:
        break

final_phase = loop.phase.value
state = loop.search.get_state_summary()
print("[4] FINAL phase=%s batches=%s completed_points=%s used_budget=%s"
      % (final_phase, sorted(batches), state.get("completed_points"), state.get("used_budget")))
assert final_phase in ("budget_exhausted", "converged", "completed"), final_phase
assert state.get("completed_points", 0) >= 4, state
assert state.get("used_budget", 0) > 0, state
assert len(batches) >= 1, batches

# ----------------------------------------------------------------- 5. checks
# every submitted batch task must be persisted with the adaptive fields
adaptive_tasks = [tm.get_task(t) for t in all_tids]
assert len(adaptive_tasks) == len(batches) == len(all_tids), (len(adaptive_tasks), len(batches))
for t in adaptive_tasks:
    assert t is not None, "task missing"
    assert t.get("task_type") == "adaptive_batch", t
    assert t.get("loop_id") == loop.loop_id, t
    assert t.get("dynamic") is True, t
print("[5] %d adaptive_batch tasks persisted with loop_id/dynamic fields" % len(adaptive_tasks))

# submitted-batch bridge must be idempotent: second call with no pending
# points returns task_id=None instead of creating a duplicate task.
dup = loop.submit_batch_to_executor()
assert dup.get("task_id") is None, dup
print("[6] submit-batch idempotency OK (no duplicate task on no pending batch)")

print("\nALL P3-M6 ADAPTIVE EXECUTION BRIDGE TESTS PASSED")