pyoco 0.1.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

pyoco/server/models.py

@@ -0,0 +1,30 @@
+from typing import Dict, List, Optional, Any
+from pydantic import BaseModel
+from ..core.models import RunStatus, TaskState
+
+class RunSubmitRequest(BaseModel):
+    flow_name: str
+    params: Dict[str, Any] = {}
+    tags: List[str] = []
+
+class RunResponse(BaseModel):
+    run_id: str
+    status: RunStatus
+
+class WorkerPollRequest(BaseModel):
+    worker_id: str
+    tags: List[str] = []
+
+class WorkerPollResponse(BaseModel):
+    run_id: Optional[str] = None
+    flow_name: Optional[str] = None
+    params: Optional[Dict[str, Any]] = None
+
+class WorkerHeartbeatRequest(BaseModel):
+    task_states: Dict[str, TaskState]
+    task_records: Dict[str, Any] = {}
+    logs: List[Dict[str, Any]] = []
+    run_status: RunStatus
+
+class WorkerHeartbeatResponse(BaseModel):
+    cancel_requested: bool

pyoco/server/store.py

@@ -0,0 +1,219 @@
+import uuid
+import time
+import json
+import os
+from pathlib import Path
+from typing import Dict, List, Optional, Any
+from ..core.models import RunContext, RunStatus, TaskState
+from .metrics import metrics
+from .webhook import webhook_notifier
+
+MAX_RUN_HISTORY = int(os.getenv("PYOCO_MAX_RUN_HISTORY", "50"))
+RUN_ARCHIVE_DIR = Path(os.getenv("PYOCO_RUN_ARCHIVE_DIR", "artifacts/runs"))
+MAX_LOG_BYTES_PER_TASK = int(os.getenv("PYOCO_MAX_LOG_BYTES", str(1024 * 1024)))
+
+class StateStore:
+    def __init__(self):
+        self.runs: Dict[str, RunContext] = {}
+        self.queue: List[str] = []
+        self.history: List[str] = []
+        self.max_runs = MAX_RUN_HISTORY
+        self.archive_dir = RUN_ARCHIVE_DIR
+        self.log_limit_bytes = MAX_LOG_BYTES_PER_TASK
+        self.metrics = metrics
+        self.webhook = webhook_notifier
+        
+    def create_run(self, flow_name: str, params: Dict) -> RunContext:
+        run_id = str(uuid.uuid4())
+        run_ctx = RunContext(
+            run_id=run_id,
+            flow_name=flow_name,
+            params=params or {},
+            status=RunStatus.PENDING,
+            start_time=time.time()
+        )
+        
+        self.runs[run_id] = run_ctx
+        self.queue.append(run_id)
+        self.history.append(run_id)
+        self._enforce_retention()
+        self.metrics.record_status_transition(None, run_ctx.status)
+        return run_ctx
+
+    def get_run(self, run_id: str) -> Optional[RunContext]:
+        return self.runs.get(run_id)
+
+    def list_runs(
+        self,
+        status: Optional[RunStatus] = None,
+        flow: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> List[RunContext]:
+        runs = list(self.runs.values())
+        if status:
+            runs = [r for r in runs if r.status == status]
+        if flow:
+            runs = [r for r in runs if r.flow_name == flow]
+        runs.sort(key=lambda r: r.start_time or 0, reverse=True)
+        if limit:
+            runs = runs[:limit]
+        return runs
+
+    def dequeue(self, tags: List[str] = None) -> Optional[RunContext]:
+        # Simple FIFO queue for now. Tags ignored in v0.3.0 MVP.
+        if not self.queue:
+            return None
+            
+        # Find first PENDING run
+        # Note: queue might contain cancelled runs?
+        # We should check status.
+        
+        # Pop from front
+        # We iterate to find a valid candidate
+        for i, run_id in enumerate(self.queue):
+            run = self.runs.get(run_id)
+            if run and run.status == RunStatus.PENDING:
+                self.queue.pop(i)
+                return run
+                
+        return None
+
+    def update_run(self, run_id: str, status: RunStatus = None, task_states: Dict = None, task_records: Dict = None, logs: List[Dict[str, Any]] = None):
+        run = self.runs.get(run_id)
+        if not run:
+            return
+        previous_status = run.status
+
+        if status:
+            # State transition check
+            # If server has CANCELLING, ignore RUNNING from worker
+            if run.status == RunStatus.CANCELLING and status == RunStatus.RUNNING:
+                pass
+            else:
+                run.status = status
+                
+            if status in [RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED]:
+                if not run.end_time:
+                    run.end_time = time.time()
+                    
+        if run.status != previous_status:
+            self.metrics.record_status_transition(previous_status, run.status)
+
+        if task_states:
+            for name, state in task_states.items():
+                run.tasks[name] = TaskState(state) if isinstance(state, str) else state
+        if task_records:
+            for name, record in task_records.items():
+                info = run.ensure_task_record(name)
+                state_val = record.get("state")
+                if state_val:
+                    info.state = TaskState(state_val) if isinstance(state_val, str) else state_val
+                info.started_at = record.get("started_at", info.started_at)
+                info.ended_at = record.get("ended_at", info.ended_at)
+                info.duration_ms = record.get("duration_ms", info.duration_ms)
+                info.error = record.get("error", info.error)
+                info.traceback = record.get("traceback", info.traceback)
+                info.inputs = record.get("inputs", info.inputs)
+                info.output = record.get("output", info.output)
+                info.artifacts = record.get("artifacts", info.artifacts)
+                self._record_task_metrics(run, name, info)
+        if logs:
+            for entry in logs:
+                task_name = entry.get("task") or "unknown"
+                text = entry.get("text", "")
+                encoded_len = len(text.encode("utf-8"))
+                current = run.log_bytes.get(task_name, 0)
+                if current >= self.log_limit_bytes:
+                    continue
+                if current + encoded_len > self.log_limit_bytes:
+                    allowed = max(self.log_limit_bytes - current, 0)
+                    truncated_text = text[:allowed] + "\n[log truncated]"
+                    entry = dict(entry)
+                    entry["text"] = truncated_text
+                    run.log_bytes[task_name] = self.log_limit_bytes
+                else:
+                    run.log_bytes[task_name] = current + encoded_len
+                run.logs.append(entry)
+        if status in [RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED]:
+            self._enforce_retention()
+        if run.end_time and not run.metrics_run_observed:
+            self.metrics.record_run_duration(run.flow_name, run.start_time, run.end_time)
+            run.metrics_run_observed = True
+        if run.status in [RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED]:
+            if run.webhook_notified_status != run.status.value:
+                if self.webhook.notify_run(run):
+                    run.webhook_notified_status = run.status.value
+
+    def cancel_run(self, run_id: str):
+        run = self.runs.get(run_id)
+        if not run:
+            return
+        previous = run.status
+        if run.status in [RunStatus.PENDING, RunStatus.RUNNING]:
+            run.status = RunStatus.CANCELLING
+        if run.status != previous:
+            self.metrics.record_status_transition(previous, run.status)
+
+    def export_run(self, run: RunContext) -> Dict[str, Any]:
+        return {
+            "run_id": run.run_id,
+            "flow_name": run.flow_name,
+            "params": run.params,
+            "status": run.status.value if hasattr(run.status, "value") else run.status,
+            "start_time": run.start_time,
+            "end_time": run.end_time,
+            "tasks": {name: state.value if hasattr(state, "value") else state for name, state in run.tasks.items()},
+            "task_records": run.serialize_task_records(),
+            "logs": run.logs,
+            "metadata": run.metadata,
+            "run_duration_ms": self._run_duration_ms(run),
+            "task_summary": self._build_task_summary(run),
+        }
+
+    def _enforce_retention(self):
+        removable_ids = [rid for rid in self.history if rid in self.runs]
+        while len(self.runs) > self.max_runs and removable_ids:
+            run_id = removable_ids.pop(0)
+            run = self.runs.get(run_id)
+            if not run:
+                continue
+            if run.status not in [RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED]:
+                self.history.append(run_id)
+                continue
+            self._spill_run(run)
+            self.runs.pop(run_id, None)
+            if run_id in self.queue:
+                self.queue.remove(run_id)
+        self.history = [rid for rid in self.history if rid in self.runs]
+
+    def _spill_run(self, run: RunContext):
+        try:
+            self.archive_dir.mkdir(parents=True, exist_ok=True)
+            path = self.archive_dir / f"{run.run_id}.json"
+            with path.open("w", encoding="utf-8") as fp:
+                json.dump(self.export_run(run), fp, indent=2)
+        except Exception:
+            pass
+
+    def _record_task_metrics(self, run: RunContext, task_name: str, record):
+        if task_name in run.metrics_recorded_tasks:
+            return
+        if record.duration_ms is None or record.ended_at is None:
+            return
+        self.metrics.record_task_duration(task_name, record.duration_ms)
+        run.metrics_recorded_tasks.add(task_name)
+
+    def _run_duration_ms(self, run: RunContext) -> Optional[float]:
+        if run.start_time and run.end_time:
+            return (run.end_time - run.start_time) * 1000.0
+        return None
+
+    def _build_task_summary(self, run: RunContext) -> Dict[str, Any]:
+        summary: Dict[str, Any] = {}
+        for name, record in run.task_records.items():
+            summary[name] = {
+                "state": record.state.value if hasattr(record.state, "value") else record.state,
+                "duration_ms": record.duration_ms,
+                "ended_at": record.ended_at,
+            }
+        return summary

pyoco/server/webhook.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import json
+import os
+import time
+from typing import Any, Callable, Dict, Optional
+
+import httpx
+
+from ..core.models import RunContext, RunStatus
+
+
+class WebhookNotifier:
+    def __init__(self) -> None:
+        self.url: Optional[str] = None
+        self.timeout: float = 3.0
+        self.retries: int = 1
+        self.secret: Optional[str] = None
+        self.extra_headers: Dict[str, str] = {}
+        self._sender: Optional[Callable[[str, Dict[str, Any], Dict[str, str], float], None]] = None
+        self.last_error: Optional[str] = None
+        self.load_from_env()
+
+    def load_from_env(self) -> None:
+        self.url = os.getenv("PYOCO_WEBHOOK_URL") or None
+        self.timeout = float(os.getenv("PYOCO_WEBHOOK_TIMEOUT", "3.0"))
+        self.retries = int(os.getenv("PYOCO_WEBHOOK_RETRIES", "1"))
+        self.secret = os.getenv("PYOCO_WEBHOOK_SECRET") or None
+        self.extra_headers = {}
+        self.last_error = None
+        self._sender = None
+
+    def configure(
+        self,
+        *,
+        url: Optional[str] = None,
+        timeout: Optional[float] = None,
+        retries: Optional[int] = None,
+        secret: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
+        sender: Optional[Callable[[str, Dict[str, Any], Dict[str, str], float], None]] = None,
+    ) -> None:
+        if url is not None:
+            self.url = url
+        if timeout is not None:
+            self.timeout = timeout
+        if retries is not None:
+            self.retries = max(1, retries)
+        if secret is not None:
+            self.secret = secret
+        if headers is not None:
+            self.extra_headers = dict(headers)
+        if sender is not None:
+            self._sender = sender
+        self.last_error = None
+
+    def notify_run(self, run: RunContext) -> bool:
+        if not self.url:
+            return False
+        payload = self._build_payload(run)
+        sender = self._sender or self._http_sender
+        headers = {"Content-Type": "application/json", **self.extra_headers}
+        if self.secret:
+            headers.setdefault("X-Pyoco-Token", self.secret)
+
+        last_exc: Optional[Exception] = None
+        for attempt in range(self.retries):
+            try:
+                sender(self.url, payload, headers, self.timeout)
+                self.last_error = None
+                return True
+            except Exception as exc:  # pragma: no cover - retries captured via tests
+                last_exc = exc
+                time.sleep(min(0.5, 0.1 * (attempt + 1)))
+        if last_exc:
+            self.last_error = str(last_exc)
+        return False
+
+    def _http_sender(
+        self,
+        url: str,
+        payload: Dict[str, Any],
+        headers: Dict[str, str],
+        timeout: float,
+    ) -> None:
+        httpx.post(url, json=payload, headers=headers, timeout=timeout)
+
+    def _build_payload(self, run: RunContext) -> Dict[str, Any]:
+        duration_ms = None
+        if run.start_time and run.end_time:
+            duration_ms = (run.end_time - run.start_time) * 1000.0
+        return {
+            "event": f"run.{run.status.value.lower()}",
+            "run_id": run.run_id,
+            "flow_name": run.flow_name,
+            "status": run.status.value if isinstance(run.status, RunStatus) else run.status,
+            "started_at": run.start_time,
+            "ended_at": run.end_time,
+            "duration_ms": duration_ms,
+            "tasks": run.serialize_task_records(),
+            "metadata": run.metadata,
+        }
+
+    def reset(self) -> None:
+        self.load_from_env()
+
+
+webhook_notifier = WebhookNotifier()

pyoco/socketless_reset.py

@@ -0,0 +1,7 @@
+from pyoco.server.api import store
+
+
+def reset_store():
+    store.runs.clear()
+    store.queue.clear()
+    store.history.clear()

pyoco/trace/backend.py

@@ -3,7 +3,7 @@ from typing import Any, Optional
 
 class TraceBackend(ABC):
     @abstractmethod
-    def on_flow_start(self, flow_name: str):
+    def on_flow_start(self, flow_name: str, run_id: str = None):
         pass
 
     @abstractmethod

pyoco/trace/console.py

@@ -2,14 +2,22 @@ import time
 from .backend import TraceBackend
 
 class ConsoleTraceBackend(TraceBackend):
-    def __init__(self, style: str = "cute"):
+    def __init__(self, style: str = None):
+        if style is None:
+            import os
+            env_cute = os.environ.get("PYOCO_CUTE", "true").lower()
+            if env_cute in ["0", "false", "no", "off"]:
+                style = "plain"
+            else:
+                style = "cute"
         self.style = style
 
-    def on_flow_start(self, flow_name: str):
+    def on_flow_start(self, flow_name: str, run_id: str = None):
+        rid_str = f" run_id={run_id}" if run_id else ""
         if self.style == "cute":
-            print(f"🐇 pyoco > start flow={flow_name}")
+            print(f"🐇 pyoco > start flow={flow_name}{rid_str}")
         else:
-            print(f"INFO pyoco start flow={flow_name}")
+            print(f"INFO pyoco start flow={flow_name}{rid_str}")
 
     def on_flow_end(self, flow_name: str):
         if self.style == "cute":

pyoco/worker/client.py

@@ -0,0 +1,43 @@
+import httpx
+from typing import Dict, List, Optional, Any
+from ..core.models import RunStatus, TaskState
+
+class WorkerClient:
+    def __init__(self, server_url: str, worker_id: str):
+        self.server_url = server_url.rstrip("/")
+        self.worker_id = worker_id
+        self.client = httpx.Client(base_url=self.server_url)
+
+    def poll(self, tags: List[str] = []) -> Optional[Dict[str, Any]]:
+        try:
+            resp = self.client.post("/workers/poll", json={
+                "worker_id": self.worker_id,
+                "tags": tags
+            })
+            resp.raise_for_status()
+            data = resp.json()
+            if data.get("run_id"):
+                return data
+            return None
+        except Exception as e:
+            print(f"Poll failed: {e}")
+            return None
+
+    def heartbeat(self, run_id: str, task_states: Dict[str, TaskState], run_status: RunStatus) -> bool:
+        """
+        Sends heartbeat. Returns True if cancellation is requested.
+        """
+        try:
+            # Convert Enums to values
+            states_json = {k: v.value for k, v in task_states.items()}
+            status_value = run_status.value
+            
+            resp = self.client.post(f"/runs/{run_id}/heartbeat", json={
+                "task_states": states_json,
+                "run_status": status_value
+            })
+            resp.raise_for_status()
+            return resp.json().get("cancel_requested", False)
+        except Exception as e:
+            print(f"Heartbeat failed: {e}")
+            return False

pyoco/worker/runner.py

@@ -0,0 +1,166 @@
+import time
+import uuid
+from typing import List, Optional
+from ..core.engine import Engine
+from ..core.models import RunContext, RunStatus, Flow
+from ..trace.backend import TraceBackend
+from ..discovery.loader import TaskLoader
+from ..schemas.config import PyocoConfig
+from ..client import Client
+
+from ..trace.console import ConsoleTraceBackend
+
+class RemoteTraceBackend(TraceBackend):
+    def __init__(self, client: Client, run_ctx: RunContext, cute: bool = True):
+        self.client = client
+        self.run_ctx = run_ctx
+        self.last_heartbeat = 0
+        self.heartbeat_interval = 1.0 # sec
+        self.console = ConsoleTraceBackend(style="cute" if cute else "plain")
+
+    def _send_heartbeat(self, force=False):
+        now = time.time()
+        if force or (now - self.last_heartbeat > self.heartbeat_interval):
+            cancel = self.client.heartbeat(self.run_ctx)
+            if cancel and self.run_ctx.status not in [RunStatus.CANCELLING, RunStatus.CANCELLED]:
+                print(f"🛑 Cancellation requested from server for run {self.run_ctx.run_id}")
+                self.run_ctx.status = RunStatus.CANCELLING
+            self.last_heartbeat = now
+
+    def on_flow_start(self, name: str, run_id: Optional[str] = None):
+        self.console.on_flow_start(name, run_id)
+        self._send_heartbeat(force=True)
+
+    def on_flow_end(self, name: str):
+        self.console.on_flow_end(name)
+        self._send_heartbeat(force=True)
+
+    def on_node_start(self, node_name: str):
+        self.console.on_node_start(node_name)
+        self._send_heartbeat()
+
+    def on_node_end(self, node_name: str, duration: float):
+        self.console.on_node_end(node_name, duration)
+        self._send_heartbeat(force=True)
+
+    def on_node_error(self, node_name: str, error: Exception):
+        self.console.on_node_error(node_name, error)
+        self._send_heartbeat(force=True)
+
+
+class Worker:
+    def __init__(self, server_url: str, config: PyocoConfig, tags: List[str] = []):
+        self.server_url = server_url
+        self.config = config
+        self.tags = tags
+        self.worker_id = f"w-{uuid.uuid4().hex[:8]}"
+        self.client = Client(server_url, self.worker_id)
+        self.loader = TaskLoader(config)
+        self.loader.load() # Load tasks/flows once
+
+    def start(self):
+        print(f"🐜 Worker {self.worker_id} started. Connected to {self.server_url}")
+        try:
+            while True:
+                job = self.client.poll(self.tags)
+                if job:
+                    self._execute_job(job)
+                else:
+                    time.sleep(2.0)
+        except KeyboardInterrupt:
+            print("\n🛑 Worker stopping...")
+
+    def _execute_job(self, job):
+        run_id = job["run_id"]
+        flow_name = job["flow_name"]
+        params = job["params"] or {}
+        
+        print(f"🚀 Received job: {run_id} (Flow: {flow_name})")
+        
+        # Find flow
+        flow_def = self.config.flows.get(flow_name)
+        if not flow_def:
+            print(f"❌ Flow '{flow_name}' not found in local config.")
+            return
+
+        # Build Flow object using exec (same as main.py)
+        from ..core.models import Flow as FlowModel
+        from ..dsl.syntax import TaskWrapper
+        
+        eval_context = {name: TaskWrapper(task) for name, task in self.loader.tasks.items()}
+        
+        try:
+            flow = FlowModel(name=flow_name)
+            for t in self.loader.tasks.values():
+                flow.add_task(t)
+            
+            # Evaluate graph
+            exec(flow_def.graph, {}, eval_context)
+            
+        except Exception as e:
+            print(f"❌ Error building flow: {e}")
+            return
+        
+        # Execute
+        engine = Engine()
+        
+        # We need to inject run_id into Engine.
+        # Engine.run generates run_id if not provided.
+        # We need to pass it.
+        # Engine.run doesn't accept run_id argument currently.
+        # It creates RunContext inside.
+        # I need to modify Engine.run to accept optional run_id or RunContext.
+        
+        # Wait, I modified Engine.run in v0.2.0.
+        # Let's check Engine.run signature.
+        pass 
+        
+        # I will modify Engine.run to accept run_id.
+        # For now, let's assume I will do that.
+        
+        # Create RemoteTraceBackend
+        # But we need run_ctx to create backend.
+        # And Engine creates run_ctx.
+        # Chicken and egg.
+        
+        # Solution: Engine.run should accept an existing RunContext or run_id.
+        # If I pass run_id, Engine creates RunContext with that ID.
+        # Then I can access it?
+        # Or I pass a callback to get it?
+        
+        # Better: Pass run_id to Engine.run.
+        # Engine creates RunContext.
+        # Engine calls trace.on_flow_start(run_id=...).
+        # RemoteTraceBackend receives run_id.
+        # But RemoteTraceBackend needs to know which run_id to report to (it knows from constructor).
+        # Actually, RemoteTraceBackend needs access to the RunContext object that Engine creates.
+        # Because it reads `run_ctx.tasks` and `run_ctx.status`.
+        
+        # If Engine creates RunContext internally, we can't pass it to Backend beforehand.
+        # Unless Engine exposes it.
+        
+        # Alternative:
+        # Modify Engine to accept `run_context` argument.
+        # If provided, use it.
+        
+        # I will modify Engine.run to accept `run_context`.
+        
+        run_ctx = RunContext(run_id=run_id, status=RunStatus.RUNNING)
+        # Initialize tasks as PENDING? Engine does that.
+        
+        backend = RemoteTraceBackend(self.client, run_ctx)
+        engine.trace = backend # Replace default console trace? Or chain?
+        # Maybe chain if we want local logs too.
+        # For now, just replace or use MultiBackend (not implemented).
+        # Let's just use RemoteBackend.
+        
+        try:
+            engine.run(flow, params=params, run_context=run_ctx)
+            print(f"✅ Job {run_id} completed: {run_ctx.status}")
+            # Send final heartbeat
+            self.client.heartbeat(run_ctx)
+        except Exception as e:
+            print(f"💥 Job {run_id} failed: {e}")
+            # Heartbeat one last time
+            run_ctx.status = RunStatus.FAILED
+            self.client.heartbeat(run_ctx)