91 lines
2.6 KiB
Python
91 lines
2.6 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)
|
|
"""
|
|
|
|
import time
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
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).
|
|
|
|
Returns:
|
|
Liste av Dask-delayed objects — klar for execute_dag().
|
|
"""
|
|
ids = agent_ids or [getattr(fn, "__name__", f"agent_{i}") for i, fn in enumerate(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.
|
|
"""
|
|
results = dask.compute(*tasks, scheduler=scheduler)
|
|
return list(results)
|