"""Adaptive simulation loop orchestrator (P3-M2). Bridges the web-side feasibility-first search (FeasibilityFirstSearch) to the local executor through the task system. Lifecycle: start_loop() build search from plan parameters -> create the initial adaptive_batch task (task_type="adaptive_batch") advance_loop() when the current batch task is completed, read its results (point_id -> params -> metrics), feed them back into the search via report_result(), run AI analysis, then either converge or create the next batch task. The orchestrator is pull-driven (advance_loop is invoked by a caller / route / scheduler), matching the existing poll-based executor model and keeping the task system free of new completion hooks. Loop state is persisted as JSON under output/adaptive_loops/ so a restart can at least recover loop metadata and the current batch task. All source is ASCII only. """ import json import os from datetime import datetime from typing import Any, Dict, List, Optional from ..services.feasibility_search import FeasibilityFirstSearch, ParameterRange from ..services.l0_prescreening import L0PreScreeningEngine from ..services.task_manager import get_task_manager, TaskManager from ..services.result_analyst import AIResultAnalyst from ..config import KIMI_API_KEY # Valid loop phases (mirror the executor/task vocabulary). LOOP_PHASE_INIT = "initializing" LOOP_PHASE_RUNNING = "running" LOOP_PHASE_CONVERGED = "converged" LOOP_PHASE_BUDGET_EXHAUSTED = "budget_exhausted" LOOP_PHASE_FAILED = "failed" def _loop_state_path(state_dir: str, loop_id: str) -> str: return os.path.join(state_dir, "%s_loop.json" % loop_id) class AdaptiveOrchestrator: """Coordinates adaptive search <-> task system <-> local executor.""" def __init__( self, task_manager: Optional[TaskManager] = None, state_dir: Optional[str] = None, l0_engine: Optional[L0PreScreeningEngine] = None, ): self.tm = task_manager or get_task_manager() self.l0_engine = l0_engine or L0PreScreeningEngine() self.state_dir = state_dir or os.path.join( os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))), "output", "adaptive_loops", ) os.makedirs(self.state_dir, exist_ok=True) self._analyst = AIResultAnalyst() # loop_id -> in-memory search + runtime metadata self._searches: Dict[str, FeasibilityFirstSearch] = {} self._loops: Dict[str, Dict[str, Any]] = {} self._load_state() # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def start_loop( self, loop_id: str, parameters: List[Dict[str, Any]], plan_id: Optional[int] = None, plan_data: Optional[Dict[str, Any]] = None, objective_metric: str = "tavg_nm", objective_direction: str = "maximize", total_budget: int = 80, batch_size: int = 4, initial_samples: int = 16, seed: int = 42, ) -> Dict[str, Any]: """Start an adaptive loop from explicit search parameters. parameters: [{"name", "min_value", "max_value", "step"?, "unit"?}] Returns loop status including the initial batch task_id. """ if loop_id in self._loops: raise ValueError("Loop %s already exists" % loop_id) search = self._build_search( parameters, objective_metric, objective_direction, total_budget, batch_size, initial_samples, seed, ) self._searches[loop_id] = search meta = { "loop_id": loop_id, "phase": LOOP_PHASE_INIT, "plan_id": plan_id, "plan_data": plan_data or {}, "parameters": parameters, "objective_metric": objective_metric, "objective_direction": objective_direction, "total_budget": total_budget, "batch_size": batch_size, "initial_samples": initial_samples, "current_batch": 0, "current_task_id": None, "n_results": 0, "created_at": datetime.now().isoformat(), "updated_at": datetime.now().isoformat(), } self._loops[loop_id] = meta # First batch batch = search.generate_initial_batch() task = self._create_batch_task(loop_id, search, batch, batch_id=0) meta["current_batch"] = 0 meta["current_task_id"] = task.get("task_id") meta["phase"] = LOOP_PHASE_RUNNING self._save_state(loop_id) return self._loop_view(loop_id, include_task=task) def advance_loop(self, loop_id: str) -> Dict[str, Any]: """Advance one adaptive step: if the current batch task completed, feed results back, analyze, then converge or create next batch.""" meta = self._loops.get(loop_id) if meta is None: raise ValueError("Loop %s not found" % loop_id) search = self._searches.get(loop_id) if search is None: raise ValueError("Loop %s has no search (process restarted?)" % loop_id) if meta["phase"] in (LOOP_PHASE_CONVERGED, LOOP_PHASE_BUDGET_EXHAUSTED, LOOP_PHASE_FAILED): return self._loop_view(loop_id) task_id = meta.get("current_task_id") if not task_id: return self._loop_view(loop_id, message="no current task") task = self.tm.get_task(task_id) if task is None: raise ValueError("Loop %s current task %s missing" % (loop_id, task_id)) if task.get("status") not in ("completed", "failed", "cancelled"): # still running - nothing to do yet return self._loop_view(loop_id, message="batch still running") # ---- batch finished: pull results ---- results = self.tm.get_task_results(task_id) point_results = (results or {}).get("results", []) if task.get("status") == "failed": meta["phase"] = LOOP_PHASE_FAILED meta["updated_at"] = datetime.now().isoformat() self._save_state(loop_id) return self._loop_view(loop_id, message="batch task failed") # Feed results back into the search by point_id n_fed = 0 for r in point_results: pid = r.get("point_id") if pid is None: continue metrics = r.get("metrics") or {} if not metrics: # metrics may be flattened on the result top level for k, v in r.items(): if k not in ("point_id", "params", "status", "error", "point_index", "solve_time_s"): if isinstance(v, (int, float)): metrics[k] = float(v) status = "ok" if r.get("status") == "OK" else "failed" search.report_result(int(pid), metrics, status) n_fed += 1 meta["n_results"] += n_fed # Optional AI analysis on accumulated results (best-effort; only # when the AI backend is configured, else keep quantitative only). try: analysis = None if KIMI_API_KEY: analysis = self._analyst.analyze( results=self._collect_results(search), targets=None, fidelity="L3", scan_parameters=[p.name for p in search.parameters], ) meta["latest_analysis"] = analysis except Exception as exc: # noqa: BLE001 - analysis is best-effort meta["latest_analysis_error"] = str(exc) # ---- decide next step ---- state = search.get_state_summary() if state.get("convergence_status") == "converged": meta["phase"] = LOOP_PHASE_CONVERGED elif state.get("remaining_budget", 0) <= 0: meta["phase"] = LOOP_PHASE_BUDGET_EXHAUSTED else: nxt = search.select_next_batch() if not nxt: meta["phase"] = LOOP_PHASE_CONVERGED else: next_batch_id = int(meta.get("current_batch", 0)) + 1 task = self._create_batch_task(loop_id, search, nxt, batch_id=next_batch_id) meta["current_batch"] = next_batch_id meta["current_task_id"] = task.get("task_id") meta["phase"] = LOOP_PHASE_RUNNING meta["updated_at"] = datetime.now().isoformat() self._save_state(loop_id) return self._loop_view(loop_id) def get_loop_status(self, loop_id: str) -> Dict[str, Any]: if loop_id not in self._loops: raise ValueError("Loop %s not found" % loop_id) return self._loop_view(loop_id) def list_loops(self) -> List[Dict[str, Any]]: return [ { "loop_id": m["loop_id"], "phase": m["phase"], "current_batch": m.get("current_batch"), "n_results": m.get("n_results"), "updated_at": m.get("updated_at"), } for m in self._loops.values() ] # ------------------------------------------------------------------ # Internals # ------------------------------------------------------------------ def _build_search( self, parameters, objective_metric, objective_direction, total_budget, batch_size, initial_samples, seed, ) -> FeasibilityFirstSearch: ranges = [] for p in parameters: ranges.append(ParameterRange( name=p["name"], min_value=float(p["min_value"]), max_value=float(p["max_value"]), step=float(p["step"]) if p.get("step") else None, unit=p.get("unit", ""), description=p.get("description", ""), )) if not ranges: raise ValueError("at least one search parameter required") return FeasibilityFirstSearch( parameters=ranges, l0_engine=self.l0_engine, total_budget=int(total_budget), batch_size=int(batch_size), initial_samples=int(initial_samples), objective_metric=objective_metric, objective_direction=objective_direction, seed=int(seed), ) def _create_batch_task(self, loop_id, search, batch, batch_id: int) -> Dict[str, Any]: meta = self._loops[loop_id] parameters = [] point_ids = [] for p in batch: pid = getattr(p, "id", None) params = getattr(p, "params", None) if params is None and isinstance(p, dict): params = p.get("params") pid = p.get("point_id", p.get("id")) point_ids.append(pid) item = dict(params or {}) item["point_id"] = pid parameters.append(item) task = self.tm.create_task( plan_id=meta.get("plan_id"), plan_data=meta.get("plan_data") or {}, parameters=parameters, task_name="adaptive-%s-b%s" % (loop_id, batch_id), task_type="adaptive_batch", loop_id=loop_id, batch_id=batch_id, point_ids=point_ids, dynamic=True, ) return task def _collect_results(self, search) -> List[Dict[str, Any]]: """Flatten search completed points to [{point_id, **metrics}].""" out = [] for p in search.state.get_completed_points(): entry = {"point_id": p.id, "status": p.status} entry.update(p.metrics) out.append(entry) return out def _loop_view(self, loop_id, include_task=None, message=None) -> Dict[str, Any]: meta = self._loops[loop_id] search = self._searches.get(loop_id) view = { "loop_id": loop_id, "phase": meta["phase"], "plan_id": meta.get("plan_id"), "current_batch": meta.get("current_batch"), "current_task_id": meta.get("current_task_id"), "n_results": meta.get("n_results", 0), "updated_at": meta.get("updated_at"), "created_at": meta.get("created_at"), "search_state": search.get_state_summary() if search else None, } if include_task: view["batch_task"] = include_task if message: view["message"] = message if meta.get("latest_analysis"): view["latest_analysis"] = meta["latest_analysis"] return view # ------------------------------------------------------------------ # Persistence (best-effort) # ------------------------------------------------------------------ def _save_state(self, loop_id: str) -> None: meta = self._loops.get(loop_id) if meta is None: return search = self._searches.get(loop_id) payload = dict(meta) try: if search is not None and hasattr(search, "export_state"): payload["search_export"] = search.export_state() except Exception: pass path = _loop_state_path(self.state_dir, loop_id) with open(path, "w", encoding="utf-8") as f: json.dump(payload, f, ensure_ascii=False, indent=2) def _load_state(self) -> None: if not os.path.isdir(self.state_dir): return for fname in os.listdir(self.state_dir): if not fname.endswith("_loop.json"): continue path = os.path.join(self.state_dir, fname) try: with open(path, "r", encoding="utf-8") as f: meta = json.load(f) if isinstance(meta, dict) and meta.get("loop_id"): # metadata recovered; search object is rebuilt on demand self._loops[meta["loop_id"]] = meta except Exception: continue # Module-level singleton (mirrors get_task_manager pattern). _orchestrator: Optional[AdaptiveOrchestrator] = None def get_orchestrator() -> AdaptiveOrchestrator: global _orchestrator if _orchestrator is None: _orchestrator = AdaptiveOrchestrator() return _orchestrator