expops 0.1.3__py3-none-any.whl

mlops/core/experiment_tracker.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional, Protocol
+
+
+class ExperimentTracker(Protocol):
+    """
+    Minimal interface for an experiment tracker.
+
+    The platform uses this for optional experiment tracking (params/metrics/artifacts/tags)
+    alongside the built-in KV-store metric logging in `mlops.core.step_system`.
+    """
+
+    def start_run(
+        self,
+        run_name: Optional[str] = None,
+        run_id: Optional[str] = None,
+        tags: Optional[Dict[str, Any]] = None,
+    ) -> Any:
+        """Start a new run (returns context manager or run handle)."""
+        ...
+
+    def end_run(self, status: Optional[str] = "FINISHED") -> None:
+        """End the current active run."""
+        ...
+
+
+class NoOpExperimentTracker(ExperimentTracker):
+    """
+    Default tracker: prints a few lifecycle messages but intentionally ignores metrics
+    to avoid noisy logs. Safe fallback when no experiment tracking backend is configured.
+    """
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        self.config = config if config else {}
+        self.run_active = False
+        self.current_run_id = None
+        print(f"[NoOpTracker] Initialized with config: {self.config}")
+
+    def start_run(
+        self,
+        run_name: Optional[str] = None,
+        run_id: Optional[str] = None,
+        tags: Optional[Dict[str, Any]] = None,
+    ) -> "NoOpExperimentTracker":
+        import uuid
+
+        if self.run_active:
+            print(
+                f"[NoOpTracker] Warning: A run (ID: {self.current_run_id}) is already active. "
+                f"Starting a new nested run is not fully supported by NoOpTracker; state will be overridden."
+            )
+
+        self.current_run_id = run_id if run_id else str(uuid.uuid4())
+        self.run_active = True
+
+        run_display_name = run_name if run_name else "default_run"
+        print(f"[NoOpTracker] Started run. Name: '{run_display_name}', ID: '{self.current_run_id}'")
+        if tags:
+            # Intentionally ignore tags in the NoOp tracker.
+            pass
+        return self  # Return self to allow use as a context manager
+
+    def end_run(self, status: Optional[str] = "FINISHED") -> None:
+        if self.run_active:
+            print(f"[NoOpTracker][RunID: {self.current_run_id}] Ended run with status: {status}")
+            self.run_active = False
+            self.current_run_id = None
+        else:
+            print("[NoOpTracker] No active run to end.")
+
+    def __enter__(self):
+        if not self.run_active:
+            self.start_run()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        status = "FAILED" if exc_type else "FINISHED"
+        self.end_run(status=status)
+
+

mlops/core/graph_types.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+from dataclasses import dataclass
+from enum import Enum
+
+
+class NodeType(Enum):
+    """Types of nodes in the execution graph."""
+    PROCESS = "process"
+    STEP = "step"
+
+
+@dataclass
+class ProcessConfig:
+    """Configuration for a process node."""
+    name: str
+    depends_on: list[str] | None = None
+    parallel: bool = True
+    code_function: Optional[str] = None  # Function name to execute (if different from name)
+    process_type: str = "process"  # e.g., "process" or special types like "chart"
+    
+    def __post_init__(self) -> None:
+        if self.depends_on is None:
+            self.depends_on = []
+
+
+@dataclass 
+class StepConfig:
+    """Configuration for a step node."""
+    name: str
+    type: str = "step"
+    process: Optional[str] = None
+    inputs: list[str] | None = None
+    outputs: list[str] | None = None
+    loop_back_to: Optional[str] = None
+    condition: Optional[str] = None
+    parallel: bool = True
+
+
+@dataclass
+class NetworkXGraphConfig:
+    """Configuration for NetworkX-based graph execution."""
+    processes: list[ProcessConfig] | None = None
+    steps: list[StepConfig] | None = None
+    execution: dict[str, Any] | None = None
+    
+    def __post_init__(self) -> None:
+        if self.processes is None:
+            self.processes = []
+        if self.steps is None:
+            self.steps = []
+        if self.execution is None:
+            self.execution = {}
+
+
+@dataclass 
+class ExecutionResult:
+    """Result of executing a node (process or step)."""
+    name: str
+    result: Optional[dict[str, Any]] = None  # Dictionary containing step/process results
+    execution_time: float = 0.0
+    was_cached: bool = False
+    error: Optional[str] = None 

mlops/core/networkx_parser.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+import logging
+
+from .graph_types import NetworkXGraphConfig, ProcessConfig, StepConfig
+from .step_system import get_step_registry
+
+logger = logging.getLogger(__name__)
+
+
+class NetworkXPipelineParser:
+    """Parser for NetworkX-based pipeline configurations with loop support."""
+    
+    def __init__(self) -> None:
+        self.processes: Dict[str, ProcessConfig] = {}
+        self.steps: Dict[str, StepConfig] = {}
+        
+    def parse_pipeline_config(self, pipeline_config: Dict[str, Any]) -> NetworkXGraphConfig:
+        """Parse pipeline config and return NetworkX graph configuration."""
+        self.processes = {}
+        self.steps = {}
+        
+        if "processes" in pipeline_config:
+            self._parse_networkx_format(pipeline_config)
+
+        config = self._generate_networkx_config(pipeline_config)
+        
+        logger.info(f"Parsed pipeline with {len(self.processes)} processes and {len(self.steps)} steps")
+        return config
+        
+    def _parse_networkx_format(self, pipeline_config: Dict[str, Any]) -> None:
+        """Parse NetworkX configuration format with DAG flow support."""
+        
+        if "process_adjlist" in pipeline_config:
+            self._parse_process_adjlist(pipeline_config["process_adjlist"])
+        
+        for process_data in pipeline_config.get("processes", []):
+            process_name = process_data["name"]
+            code_function = process_data.get("code_function")  # Extract code_function if provided
+            proc_type = process_data.get("type", "process")
+            
+            if process_name in self.processes:
+                process_config = self.processes[process_name]
+                process_config.parallel = process_data.get("parallel", process_config.parallel)
+                # Update code_function if provided in config
+                if code_function:
+                    process_config.code_function = code_function
+                # Update process_type if provided (important for chart processes)
+                if proc_type:
+                    process_config.process_type = str(proc_type)
+                    logger.debug(f"Updated process '{process_name}' type to '{proc_type}'")
+            else:
+                process_config = ProcessConfig(
+                    name=process_name,
+                    parallel=process_data.get("parallel", True),
+                    code_function=code_function,  # Set code_function from config
+                    process_type=str(proc_type)
+                )
+                    
+                self.processes[process_name] = process_config
+            
+        # Manual-step mode: steps are executed inside process runners, not scheduled as nodes.
+    
+    def _parse_process_adjlist(self, adjlist_value: Any) -> None:
+        """Parse process-level DAG from adjacency list string or list of lines.
+        
+        Adjacency list lines follow NetworkX semantics: first token is the source node,
+        subsequent tokens are target nodes. Lines may include comments after a '#'.
+        """
+        if isinstance(adjlist_value, str):
+            lines = adjlist_value.splitlines()
+        elif isinstance(adjlist_value, list):
+            lines = [str(x) for x in adjlist_value]
+        else:
+            raise ValueError("process_adjlist must be a string or list of lines")
+        
+        for raw_line in lines:
+            line = raw_line.strip()
+            if not line:
+                continue
+            # Remove comments after '#'
+            if "#" in line:
+                line = line.split("#", 1)[0].strip()
+            if not line:
+                continue
+            parts = line.split()
+            if not parts:
+                continue
+            source = parts[0]
+            targets = parts[1:]
+            
+            # Ensure source process exists
+            if source not in self.processes:
+                self.processes[source] = ProcessConfig(name=source, depends_on=[])
+            
+            # Add or update targets with dependency on source
+            for target in targets:
+                if target not in self.processes:
+                    self.processes[target] = ProcessConfig(name=target, depends_on=[source])
+                else:
+                    deps = self.processes[target].depends_on or []
+                    if source not in deps:
+                        self.processes[target].depends_on = deps + [source]
+    
+    def _discover_steps_from_registry(self) -> None:
+        """Manual-step mode: keep for compatibility; intentionally a no-op."""
+        try:
+            step_registry = get_step_registry()
+            registered_steps = step_registry.list_steps() if step_registry else []
+            logger.debug(f"Manual-step mode enabled; ignoring {len(registered_steps)} registered steps during parsing")
+        except Exception:
+            logger.debug("Manual-step mode enabled; no step registry available")
+    
+    def _generate_networkx_config(self, pipeline_config: Dict[str, Any]) -> NetworkXGraphConfig:
+        """Generate NetworkX configuration from parsed processes and steps."""
+        
+        execution_config = pipeline_config.get("execution", {})
+        execution = {
+            "parallel": execution_config.get("parallel", True),
+            "failure_mode": execution_config.get("failure_mode", "stop"),
+            "max_workers": execution_config.get("max_workers", 4)
+        }
+        
+        return NetworkXGraphConfig(
+            processes=list(self.processes.values()),
+            steps=list(self.steps.values()),
+            execution=execution
+        )
+
+
+def parse_networkx_pipeline_from_config(pipeline_config: Dict[str, Any]) -> NetworkXGraphConfig:
+    """Parse pipeline configuration into NetworkX format."""
+    parser = NetworkXPipelineParser()
+    return parser.parse_pipeline_config(pipeline_config) 

mlops/core/payload_spill.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import io
+import os
+import uuid
+import logging
+from pathlib import Path
+from typing import Any, Dict, Iterable, Tuple, Optional
+
+try:
+    import numpy as np  # type: ignore
+except Exception:  # pragma: no cover
+    np = None  # type: ignore
+
+if np is not None:
+
+    class SpillArray(np.ndarray):  # type: ignore[misc]
+        """ndarray subclass with list-like truthiness semantics."""
+
+        def __new__(cls, input_array: "np.ndarray", origin_type: Optional[str] = None):
+            obj = np.asarray(input_array).view(cls)
+            obj._origin_type = origin_type  # type: ignore[attr-defined]
+            return obj
+
+        def __array_finalize__(self, obj):
+            if obj is None:
+                return
+            self._origin_type = getattr(obj, "_origin_type", None)
+
+        def __bool__(self) -> bool:  # pragma: no cover - trivial behaviour
+            return bool(self.size)
+
+else:  # pragma: no cover - numpy unavailable
+    SpillArray = None  # type: ignore
+
+PAYLOAD_REF_KEY = "__mlops_payload_ref__"
+PAYLOAD_META_VERSION = 2
+
+
+def _get_logger() -> logging.Logger:
+    return logging.getLogger(__name__)
+
+
+def _coerce_array(value: Any) -> Optional["np.ndarray"]:
+    """Best-effort conversion of supported payloads to a numpy array."""
+    if np is None:
+        return None
+    if isinstance(value, np.ndarray):
+        return value
+    if hasattr(value, "to_numpy"):
+        try:
+            arr = value.to_numpy()
+            return arr if isinstance(arr, np.ndarray) else None
+        except Exception:
+            return None
+    if isinstance(value, (list, tuple)):
+        try:
+            arr = np.asarray(value)
+            if arr.dtype == object:
+                return None
+            return arr
+        except Exception:
+            return None
+    return None
+
+
+def _estimate_bytes(value: Any) -> int:
+    if np is not None and isinstance(value, np.ndarray):
+        return int(value.nbytes)
+    if isinstance(value, (bytes, bytearray)):
+        return len(value)
+    if isinstance(value, (list, tuple)):
+        try:
+            if all(isinstance(v, (int, float)) for v in value):
+                return len(value) * 8
+        except Exception:
+            return 0
+    return 0
+
+
+def _should_spill(value: Any, threshold_bytes: int) -> bool:
+    arr = _coerce_array(value)
+    if arr is None:
+        return False
+    approx = _estimate_bytes(arr)
+    return approx >= threshold_bytes
+
+
+def _serialize_array(arr: "np.ndarray", origin_type: Optional[str] = None) -> Tuple[bytes, Dict[str, Any]]:
+    buf = io.BytesIO()
+    np.savez_compressed(buf, data=arr)
+    meta = {
+        "shape": list(arr.shape),
+        "dtype": str(arr.dtype),
+        "approx_bytes": int(arr.nbytes),
+        "format": "npz",
+        "origin_type": origin_type or "ndarray",
+    }
+    return buf.getvalue(), meta
+
+
+def _wrap_spilled_array(arr: "np.ndarray", origin_type: Optional[str]) -> Any:
+    if origin_type == "list":
+        return arr.tolist()
+    if origin_type == "tuple":
+        return tuple(arr.tolist())
+    if np is not None and isinstance(arr, np.ndarray) and SpillArray is not None:
+        try:
+            wrapped = arr.view(SpillArray)
+            wrapped._origin_type = origin_type  # type: ignore[attr-defined]
+            return wrapped
+        except Exception:
+            return arr
+    return arr
+
+
+def _build_payload_filename(run_id: Optional[str], process_name: Optional[str], key_path: Iterable[str]) -> str:
+    segments = ["payloads"]
+    if run_id:
+        segments.append(run_id)
+    if process_name:
+        segments.append(process_name)
+    path_tuple = tuple(key_path)
+    if path_tuple:
+        path_segment = "-".join(part or "part" for part in path_tuple)
+    else:
+        path_segment = "data"
+    segments.append(path_segment)
+    segments.append(str(uuid.uuid4()))
+    filename = "/".join(s.strip("/").replace(" ", "_") for s in segments)
+    return f"{filename}.npz"
+
+
+def _store_bytes(state_manager: Any, filename: str, payload: bytes) -> str:
+    """Persist bytes via object store when configured, otherwise on the local cache path."""
+    if getattr(state_manager, "object_store", None):
+        try:
+            build_uri = getattr(state_manager, "_build_object_uri", None)
+            if callable(build_uri):
+                uri = build_uri(filename)
+            else:
+                uri = state_manager.object_store.build_uri(filename)  # type: ignore[call-arg]
+            state_manager.object_store.put_bytes(uri, payload, content_type="application/octet-stream")  # type: ignore[attr-defined]
+            return uri
+        except Exception as e:
+            _get_logger().warning(f"[PayloadSpill] Failed to put bytes to object store ({filename}): {e}")
+    cache_dir = getattr(state_manager, "cache_dir", Path("."))
+    local_path = Path(cache_dir) / filename
+    local_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(local_path, "wb") as fout:
+        fout.write(payload)
+    return str(local_path.resolve())
+
+
+def spill_large_payloads(result: Dict[str, Any],
+                         state_manager: Any,
+                         run_id: Optional[str],
+                         process_name: Optional[str],
+                         threshold_bytes: int = 5_000_000) -> Dict[str, Any]:
+    """
+    Replace large numeric payloads inside a result dict with lightweight references.
+    """
+    if not isinstance(result, dict) or state_manager is None or np is None:
+        return result
+
+    def _process(key: str, value: Any, path: Tuple[str, ...]) -> Any:
+        if isinstance(value, dict):
+            return {k: _process(k, v, path + (k,)) for k, v in value.items()}
+        if isinstance(value, list):
+            if _should_spill(value, threshold_bytes):
+                return _spill_value(value, path)
+            try:
+                return [_process(str(idx), item, path + (str(idx),)) for idx, item in enumerate(value)]
+            except Exception:
+                pass
+        if isinstance(value, tuple):
+            if _should_spill(value, threshold_bytes):
+                return _spill_value(list(value), path)
+            return tuple(_process(str(idx), item, path + (str(idx),)) for idx, item in enumerate(value))
+        if _should_spill(value, threshold_bytes):
+            return _spill_value(value, path)
+        return value
+
+    def _spill_value(payload_value: Any, key_path: Tuple[str, ...]) -> Dict[str, Any]:
+        arr = _coerce_array(payload_value)
+        if arr is None:
+            return payload_value  # type: ignore[return-value]
+        if isinstance(payload_value, list):
+            origin_type = "list"
+        elif isinstance(payload_value, tuple):
+            origin_type = "tuple"
+        elif np is not None and isinstance(payload_value, np.ndarray):
+            origin_type = "ndarray"
+        else:
+            origin_type = type(payload_value).__name__
+        data_bytes, meta = _serialize_array(arr, origin_type=origin_type)
+        path_tuple = key_path or ("payload",)
+        filename = _build_payload_filename(run_id, process_name, path_tuple)
+        uri = _store_bytes(state_manager, filename, data_bytes)
+        ref = {
+            PAYLOAD_REF_KEY: True,
+            "uri": uri,
+            "meta": meta,
+            "version": PAYLOAD_META_VERSION,
+            "key_path": "/".join(path_tuple),
+            "process": process_name,
+            "run_id": run_id,
+        }
+        try:
+            _get_logger().info(f"[PayloadSpill] Spilled payload for {process_name}:{ref['key_path']} -> {uri} ({meta['approx_bytes']} bytes)")
+        except Exception:
+            pass
+        return ref
+
+    new_result = {}
+    for k, v in result.items():
+        key = str(k)
+        new_result[key] = _process(key, v, (key,))
+    return new_result
+
+
+def hydrate_payload_refs(data: Any, state_manager: Any) -> Any:
+    """Replace payload reference dicts with their hydrated numpy arrays."""
+    if state_manager is None or np is None:
+        return data
+    if isinstance(data, dict):
+        if data.get(PAYLOAD_REF_KEY):
+            return _load_payload(data, state_manager)
+        return {k: hydrate_payload_refs(v, state_manager) for k, v in data.items()}
+    if isinstance(data, list):
+        return [hydrate_payload_refs(v, state_manager) for v in data]
+    if isinstance(data, tuple):
+        return tuple(hydrate_payload_refs(v, state_manager) for v in data)
+    return data
+
+
+def _load_payload(ref: Dict[str, Any], state_manager: Any) -> Any:
+    uri = ref.get("uri")
+    if not uri:
+        return ref
+    try:
+        payload_bytes = None
+        if str(uri).startswith("gs://") and getattr(state_manager, "object_store", None):
+            payload_bytes = state_manager.object_store.get_bytes(uri)  # type: ignore[attr-defined]
+        else:
+            path = Path(uri)
+            if not path.is_absolute() and getattr(state_manager, "cache_dir", None):
+                path = Path(state_manager.cache_dir) / path
+            with open(path, "rb") as fin:
+                payload_bytes = fin.read()
+        if payload_bytes is None:
+            raise RuntimeError("No payload bytes resolved")
+        with np.load(io.BytesIO(payload_bytes), allow_pickle=False) as npz:
+            arr = npz["data"]
+        meta = ref.get("meta") if isinstance(ref.get("meta"), dict) else {}
+        origin_type = meta.get("origin_type") if isinstance(meta, dict) else None
+        return _wrap_spilled_array(arr, origin_type)
+    except Exception as e:
+        _get_logger().warning(f"[PayloadSpill] Failed to hydrate payload {uri}: {e}")
+        return ref
+
+
+if __name__ == "__main__":  # pragma: no cover - developer self-test hook
+    logging.basicConfig(level=logging.INFO)
+    if np is None:
+        print("NumPy is not available; skipping payload spill self-test.")
+    else:
+        class _TempStateManager:
+            def __init__(self) -> None:
+                self.cache_dir = Path(os.environ.get("PAYLOAD_SPILL_TMP", "/tmp/payload_spill_test"))
+                self.object_store = None
+
+        sm = _TempStateManager()
+        sample = {"X_test": np.random.rand(2000, 200)}  # ~3.2 MB -> forces spill
+        spilled = spill_large_payloads(sample, sm, run_id="selftest", process_name="demo", threshold_bytes=1_000_000)
+        print(spilled)
+
+