Chris Christiansen
9fcb9c354a
Some checks are pending
Check Python Version Consistency / Check Python Version (push) Waiting to run
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
ml/task_graph.py — Dask DAG-basert agent-koordinering for OPAX.
|
|
|
|
Konsept (fra ML-dokumentet):
|
|
Dask bygger en task graph (DAG) av operasjoner og avhengigheter.
|
|
Lazy evaluation: beregninger utføres kun når .compute() kalles.
|
|
Agent-funksjoner wrappes som @dask.delayed og kjøres parallelt.
|
|
|
|
PS-analogi:
|
|
OPAX orchestrator = DAG scheduler (server-node)
|
|
Individuelle agenter = Dask-delayed tasks (worker-noder)
|
|
|
|
Fix 2026-06-15: agent_count returneres fra build_agent_dag slik at
|
|
/run/dag kan bruke len(tasks) istedenfor len(agent_ids) (som var
|
|
undefined i feilscenariet der agent_fns og agent_ids hadde ulik lengde).
|
|
"""
|
|
|
|
import time
|
|
from typing import Any, Callable, Dict, List, Optional, Tuple
|
|
|
|
import dask
|
|
|
|
|
|
@dask.delayed
|
|
def run_agent_task(
|
|
agent_fn: Callable,
|
|
payload: Dict[str, Any],
|
|
agent_id: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Wrapper: kjører én agent som en Dask-delayed task.
|
|
Returnerer resultat med metadata for telemetri og state store.
|
|
"""
|
|
_id = agent_id or getattr(agent_fn, "__name__", "unknown")
|
|
start = time.monotonic()
|
|
try:
|
|
result = agent_fn(payload)
|
|
success = True
|
|
error = None
|
|
except Exception as exc: # noqa: BLE001
|
|
result = None
|
|
success = False
|
|
error = str(exc)
|
|
duration = round(time.monotonic() - start, 3)
|
|
return {
|
|
"agent": _id,
|
|
"result": result,
|
|
"success": success,
|
|
"error": error,
|
|
"duration_s": duration,
|
|
"timestamp": time.time(),
|
|
}
|
|
|
|
|
|
def build_agent_dag(
|
|
agent_fns: List[Callable],
|
|
payload: Dict[str, Any],
|
|
agent_ids: Optional[List[str]] = None,
|
|
) -> List:
|
|
"""
|
|
Bygger en liste av Dask-delayed tasks (én per agent).
|
|
Agenter uten avhengigheter kjøres parallelt av scheduleren.
|
|
|
|
Args:
|
|
agent_fns: Liste av callable agentfunksjoner.
|
|
payload: Felles input-dict til alle agenter.
|
|
agent_ids: Valgfrie ID-er som matcher agent_fns (for logging).
|
|
Dersom kortere enn agent_fns fylles resten med auto-navn.
|
|
|
|
Returns:
|
|
Liste av Dask-delayed objects — klar for execute_dag().
|
|
"""
|
|
n = len(agent_fns)
|
|
# Guard: aldri kræsj på lengde-mismatch
|
|
_ids = list(agent_ids or [])
|
|
_ids += [getattr(fn, "__name__", f"agent_{i}") for i, fn in enumerate(agent_fns) if i >= len(_ids)]
|
|
_ids = _ids[:n] # trim hvis agent_ids var lengre enn agent_fns
|
|
return [run_agent_task(fn, payload, aid) for fn, aid in zip(agent_fns, _ids)]
|
|
|
|
|
|
def execute_dag(
|
|
tasks: List,
|
|
scheduler: str = "synchronous",
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Kjører den ferdigbygde DAGen.
|
|
|
|
Scheduler-valg:
|
|
'synchronous' — enkelt-tråd, ingen overhead (dev / unit-test)
|
|
'threads' — lokal multi-tråd, GIL-vennlig for IO-bound agenter
|
|
'processes' — multi-prosess, CPU-bound arbeid
|
|
'distributed' — Dask cluster (full skala, krever dask[distributed])
|
|
|
|
Returns:
|
|
Liste av resultat-dicts fra run_agent_task.
|
|
"""
|
|
if not tasks:
|
|
return []
|
|
results = dask.compute(*tasks, scheduler=scheduler)
|
|
return list(results)
|