expops 0.1.3__py3-none-any.whl

mlops/reporting/__init__.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from typing import Any
+
+from .context import ChartContext
+from .registry import chart
+
+
+def run_chart_entrypoint(*args: Any, **kwargs: Any) -> int:
+    from .entrypoint import run_chart_entrypoint as _impl
+
+    return _impl(*args, **kwargs)
+
+__all__ = ["chart", "ChartContext", "run_chart_entrypoint"]
+
+

mlops/reporting/context.py

@@ -0,0 +1,187 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Optional
+
+
+class ChartContext:
+    """Runtime context passed to chart functions.
+
+    The platform supports two chart styles:
+    - **static**: chart function receives resolved `metrics` (dict)
+    - **dynamic**: chart function receives `probe_paths` and may subscribe to updates
+    """
+
+    def __init__(
+        self,
+        output_dir: Path,
+        project_id: str,
+        run_id: str,
+        kv_path: str,
+        probe_id: Optional[str],
+        theme: str,
+        chart_name: str,
+        metrics: Optional[dict[str, Any]] = None,
+        probe_ids: Optional[dict[str, str]] = None,
+        chart_type: str = "static",
+    ):
+        self.output_dir = output_dir
+        self.project_id = project_id
+        self.run_id = run_id
+        self.kv_path = kv_path
+        self.probe_id = probe_id
+        self.theme = theme
+        self.chart_name = chart_name
+        self.metrics = metrics or {}
+        self.probe_ids = probe_ids or {}
+        self.chart_type = chart_type
+
+    def _firestore_client(self) -> Optional[Any]:
+        try:
+            from google.cloud import firestore  # type: ignore
+        except Exception:
+            return None
+        try:
+            return firestore.Client()
+        except Exception:
+            return None
+
+    def load_payload(self) -> Optional[dict[str, Any]]:
+        """Deprecated helper: load a document at `self.kv_path` from Firestore (if available)."""
+        client = self._firestore_client()
+        if client is None:
+            return None
+        try:
+            root = client.collection("mlops_projects").document(self.project_id)
+            rel_path = (self.kv_path or "").strip("/")
+            parts = rel_path.split("/") if rel_path else []
+            if not parts or len(parts) % 2 != 0:
+                return None
+            ref = root
+            for i in range(0, len(parts), 2):
+                ref = ref.collection(parts[i]).document(parts[i + 1])
+            snap = ref.get()
+            return snap.to_dict() if getattr(snap, "exists", False) else None
+        except Exception:
+            return None
+
+    def get_firestore_client(self) -> Optional[Any]:
+        """Deprecated: prefer `_firestore_client()` or KV APIs."""
+        return self._firestore_client()
+    
+    # Deprecated: id-based refs no longer used
+    def get_probe_ref(self, probe_id: str) -> Optional[Any]:
+        return None
+    
+    def get_probe_metrics_ref(self, probe_id: str) -> Optional[Any]:
+        # Deprecated
+        return None
+
+    def get_probe_metrics_ref_by_path(self, probe_path: str) -> Optional[Any]:
+        """Resolve a probe_path to its metrics document reference.
+
+        This allows dynamic charts to subscribe directly using the configured
+        probe path without dealing with probe IDs.
+
+        Args:
+            probe_path: The logical probe path as specified in chart config.
+
+        Returns:
+            Firestore document reference for the metrics document, or None if
+            it cannot be resolved or Firestore is unavailable.
+        """
+        client = self._firestore_client()
+        if client is None:
+            return None
+        try:
+            from mlops.storage.path_utils import encode_probe_path  # type: ignore
+        except Exception:
+            return None
+        try:
+            enc = encode_probe_path(probe_path)
+            return (
+                client.collection("mlops_projects")
+                .document(self.project_id)
+                .collection("metric")
+                .document(self.run_id)
+                .collection("probes_by_path")
+                .document(enc)
+            )
+        except Exception:
+            return None
+
+    def savefig(self, filename: "os.PathLike[str] | str", fig: Optional[Any] = None, **savefig_kwargs: Any) -> Path:
+        """Save a matplotlib figure under this context's output directory.
+
+        - Ensures parent directories exist
+        - If `filename` is relative, it is resolved under `self.output_dir`
+        - If `fig` is provided, calls `fig.savefig(...)`; otherwise uses pyplot.savefig on current figure
+
+        Returns the resolved output path.
+        """
+        out_path = Path(filename)
+        if not out_path.is_absolute():
+            out_path = self.output_dir / out_path
+        try:
+            out_path.parent.mkdir(parents=True, exist_ok=True)
+        except Exception:
+            # Best-effort; downstream save will raise if it truly cannot write
+            pass
+        if fig is not None:
+            fig.savefig(out_path, **savefig_kwargs)
+        else:
+            # Import locally to avoid hard dependency at module import time
+            try:
+                import matplotlib.pyplot as _plt  # type: ignore
+                _plt.savefig(out_path, **savefig_kwargs)
+            except Exception:
+                # Re-raise for caller context
+                raise
+        return out_path
+
+    def get_run_status(self) -> Optional[str]:
+        """Return lowercased run status (best-effort)."""
+        try:
+            from .kv_utils import create_kv_store
+            kv = create_kv_store(self.project_id)
+            if kv and hasattr(kv, "get_run_status"):
+                status = kv.get_run_status(self.run_id)
+                if status:
+                    return str(status).lower()
+        except Exception:
+            pass
+        client = self._firestore_client()
+        if client is None:
+            return None
+        try:
+            snap = (
+                client.collection("mlops_projects")
+                .document(self.project_id)
+                .collection("runs")
+                .document(self.run_id)
+                .get()
+            )
+            if not getattr(snap, "exists", False):
+                return None
+            data = snap.to_dict() or {}
+            status = data.get("status")
+            return str(status).lower() if status else None
+        except Exception:
+            return None
+
+    def is_run_finished(self) -> bool:
+        """Return True if the run status indicates completion/failure/cancellation.
+
+        This is a convenience predicate for chart scripts to decide when to
+        unsubscribe their listeners and exit.
+        """
+        status = self.get_run_status()
+        return status in {"completed", "failed", "cancelled"}
+
+
+__all__ = [
+    "ChartContext",
+]
+
+

mlops/reporting/entrypoint.py

@@ -0,0 +1,292 @@
+from __future__ import annotations
+
+import argparse
+import inspect
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Any, Callable, Optional
+
+from .context import ChartContext
+from .kv_utils import create_kv_store
+from .registry import CHART_FUNCS
+
+
+def _write_text(path: Path, content: str) -> None:
+    try:
+        path.write_text(content, encoding="utf-8")
+    except Exception:
+        pass
+
+
+def _append_text(path: Path, content: str) -> None:
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+        with path.open("a", encoding="utf-8") as f:
+            f.write(content)
+    except Exception:
+        pass
+
+
+def _is_falsey_env(name: str, default: str = "1") -> bool:
+    return str(os.environ.get(name, default)).strip().lower() in {"0", "false", "no"}
+
+
+def _maybe_seed_chart_subprocess() -> None:
+    """Best-effort deterministic seeding for chart subprocesses."""
+    if _is_falsey_env("MLOPS_TASK_LEVEL_SEEDING", default="1"):
+        return
+    try:
+        base_seed = int(os.environ.get("MLOPS_RANDOM_SEED") or 42)
+    except Exception:
+        base_seed = 42
+
+    try:
+        import random
+
+        random.seed(base_seed)
+    except Exception:
+        pass
+
+    # Avoid importing heavyweight optional deps just for seeding. If a chart script
+    # imports these libraries, they'll be present in sys.modules and we can seed.
+    try:
+        np = sys.modules.get("numpy")
+        if np is None:
+            import numpy as np  # type: ignore
+        np.random.seed(base_seed)  # type: ignore[attr-defined]
+    except Exception:
+        pass
+
+    torch = sys.modules.get("torch")
+    if torch is not None:
+        try:
+            torch.manual_seed(base_seed)  # type: ignore[attr-defined]
+        except Exception:
+            pass
+        try:
+            if getattr(torch, "cuda", None) and torch.cuda.is_available():  # type: ignore[attr-defined]
+                torch.cuda.manual_seed_all(base_seed)  # type: ignore[attr-defined]
+        except Exception:
+            pass
+        try:
+            torch.use_deterministic_algorithms(True)  # type: ignore[attr-defined]
+        except Exception:
+            pass
+
+    tf = sys.modules.get("tensorflow")
+    if tf is not None:
+        try:
+            tf.random.set_seed(base_seed)  # type: ignore[attr-defined]
+        except Exception:
+            pass
+
+
+def _import_chart_modules(output_dir: Path) -> None:
+    """Import user chart modules (best-effort) so @chart() registrations run."""
+    files_csv = os.environ.get("MLOPS_CHART_IMPORT_FILES", "").strip()
+    if files_csv:
+        for fpath in [p.strip() for p in files_csv.split(",") if p.strip()]:
+            try:
+                import hashlib
+                import importlib.util
+
+                f_abs = str(Path(fpath).resolve())
+                mod_name = f"mlops_chart_{hashlib.sha256(f_abs.encode()).hexdigest()[:12]}"
+                spec = importlib.util.spec_from_file_location(mod_name, f_abs)
+                if spec and spec.loader:
+                    mod = importlib.util.module_from_spec(spec)
+                    spec.loader.exec_module(mod)  # type: ignore[attr-defined]
+            except Exception as exc:
+                _append_text(output_dir / "import_error.txt", f"Failed to import chart file {fpath}: {exc}\n")
+
+def _normalize_probe_mappings(probe_paths_config: Any) -> dict[str, str]:
+    out: dict[str, str] = {}
+    if isinstance(probe_paths_config, list):
+        for item in probe_paths_config:
+            if isinstance(item, dict):
+                for key, path in item.items():
+                    out[str(key)] = str(path)
+        return out
+    if isinstance(probe_paths_config, dict):
+        for key, path in probe_paths_config.items():
+            out[str(key)] = str(path)
+    return out
+
+
+def _resolve_chart_fn(chart_name: str) -> Optional[Callable[..., Any]]:
+    if chart_name and chart_name in CHART_FUNCS:
+        return CHART_FUNCS[chart_name]
+
+    if not chart_name:
+        return None
+
+    # Allow chart scripts to call `mlops.reporting.run_chart_entrypoint()` directly
+    # without @chart() by searching parent frames for a global callable.
+    frame = inspect.currentframe()
+    try:
+        cur = frame.f_back if frame else None
+        for _ in range(12):
+            if cur is None:
+                break
+            mod_name = cur.f_globals.get("__name__")
+            if isinstance(mod_name, str) and mod_name.startswith("mlops.reporting"):
+                cur = cur.f_back
+                continue
+            maybe = cur.f_globals.get(chart_name)
+            if callable(maybe):
+                return maybe
+            cur = cur.f_back
+    except Exception:
+        return None
+    finally:
+        # Avoid reference cycles through frames.
+        del frame
+
+    return None
+
+
+def run_chart_entrypoint(argv: Optional[list[str]] = None, require_function: bool = True) -> int:
+    """Standard entrypoint for user chart scripts.
+
+    - Resolves project/run from env/args
+    - For static charts: Fetches metrics from multiple probe paths
+    - For dynamic charts: Passes probe_paths directly to the user function
+    - Invokes a registered or in-module function named by MLOPS_CHART_NAME
+      with signature (metrics, ctx) for static or (probe_paths, ctx) for dynamic
+    """
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--theme", default=os.environ.get("MLOPS_CHART_THEME", "light"))
+    parser.add_argument("--project-id", default=os.environ.get("MLOPS_PROJECT_ID", ""))
+    parser.add_argument("--run-id", default=os.environ.get("MLOPS_RUN_ID", ""))
+    parser.add_argument("--oneshot", action="store_true")
+    args = parser.parse_args(argv)
+
+    output_dir = Path(os.environ.get("MLOPS_OUTPUT_DIR", "./out"))
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    chart_name = os.environ.get("MLOPS_CHART_NAME", "").strip()
+    chart_type = os.environ.get("MLOPS_CHART_TYPE", "static").strip().lower()
+    if chart_type not in {"static", "dynamic"}:
+        chart_type = "static"
+
+    # Best-effort: import chart modules so @chart() registration runs.
+    _import_chart_modules(output_dir)
+    _maybe_seed_chart_subprocess()
+
+    if not args.project_id:
+        _write_text(output_dir / "error.txt", "Missing --project-id or MLOPS_PROJECT_ID")
+        return 1
+
+    run_id = (args.run_id or os.environ.get("MLOPS_RUN_ID", "") or "").strip()
+    if not run_id:
+        _write_text(output_dir / "error.txt", "Missing --run-id or MLOPS_RUN_ID")
+        return 1
+
+    probe_paths_json = os.environ.get("MLOPS_PROBE_PATHS", "").strip()
+    if not probe_paths_json:
+        _write_text(
+            output_dir / "error.txt",
+            "Missing MLOPS_PROBE_PATHS. All charts must define probe_paths in config.",
+        )
+        return 1
+
+    try:
+        probe_paths_config = json.loads(probe_paths_json)
+    except Exception as e:
+        _write_text(output_dir / "error.txt", f"Invalid MLOPS_PROBE_PATHS JSON: {e}")
+        return 1
+
+    probe_mappings = _normalize_probe_mappings(probe_paths_config)
+
+    metrics: dict[str, Any] = {}
+    kv_store = None
+    try:
+        kv_store = create_kv_store(args.project_id)
+        if kv_store and hasattr(kv_store, "get_probe_metrics_by_path"):
+            for user_key, probe_path in probe_mappings.items():
+                if chart_type != "static":
+                    metrics[user_key] = {}
+                    continue
+                try:
+                    probe_metrics = kv_store.get_probe_metrics_by_path(run_id, probe_path) or {}
+                    metrics[user_key] = probe_metrics
+                except Exception:
+                    metrics[user_key] = {}
+    except Exception:
+        kv_store = None
+
+    ctx = ChartContext(
+        output_dir=output_dir,
+        project_id=args.project_id,
+        run_id=run_id,
+        kv_path="",  # legacy; not used by the current runner
+        probe_id=None,  # legacy; probe IDs are not used by path-based charts
+        theme=args.theme,
+        chart_name=chart_name,
+        metrics=metrics,
+        probe_ids={},
+        chart_type=chart_type,
+    )
+
+    fn = _resolve_chart_fn(chart_name)
+    if fn is None:
+        if require_function:
+            _write_text(
+                output_dir / "error.txt",
+                f"No chart function found for '{chart_name}'. Define def {chart_name}(metrics, ctx): ... or use @chart().",
+            )
+            return 1
+        return 0
+
+    try:
+        sig = inspect.signature(fn)
+        params = list(sig.parameters.keys())
+        num_params = len(params)
+
+        if chart_type == "dynamic":
+            first_name = params[0] if num_params >= 1 else None
+            payload = probe_mappings
+            if num_params >= 2:
+                fn(payload, ctx)
+            elif num_params == 1:
+                if first_name in ("probe_paths", "paths"):
+                    fn(payload)
+                else:
+                    fn(ctx)
+            else:
+                _write_text(
+                    output_dir / "error.txt",
+                    f"Dynamic chart function '{chart_name}' must accept at least one parameter (probe_paths or ctx).",
+                )
+                return 1
+        else:
+            if num_params >= 2:
+                fn(metrics, ctx)
+            elif num_params == 1:
+                if params[0] == "metrics":
+                    fn(metrics)
+                else:
+                    fn(ctx)
+            else:
+                _write_text(
+                    output_dir / "error.txt",
+                    f"Chart function '{chart_name}' must accept at least one parameter (metrics or ctx).",
+                )
+                return 1
+    except Exception as _fe:
+        import traceback
+
+        error_msg = f"Chart function failed: {chart_name} -> {_fe}\n{traceback.format_exc()}"
+        _write_text(output_dir / "error.txt", error_msg)
+        return 1
+
+    return 0
+
+
+__all__ = ["run_chart_entrypoint"]
+
+if __name__ == "__main__":
+    import sys as _sys
+    _sys.exit(run_chart_entrypoint())

mlops/reporting/kv_utils.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Optional, Tuple
+
+from mlops.core.workspace import get_projects_root, get_workspace_root
+from mlops.storage.factory import create_kv_store as _create_kv_store
+
+
+def _load_backend_cfg_from_project_config(project_id: str) -> dict[str, Any]:
+    root = get_workspace_root()
+    cfg_path = get_projects_root(root) / project_id / "configs" / "project_config.yaml"
+    if not cfg_path.exists():
+        return {}
+    try:
+        import yaml  # type: ignore
+    except Exception:
+        return {}
+    try:
+        cfg = yaml.safe_load(cfg_path.read_text()) or {}
+        cache_cfg = (((cfg.get("model") or {}).get("parameters") or {}).get("cache") or {})
+        backend_cfg = (cache_cfg.get("backend") or {}) if isinstance(cache_cfg, dict) else {}
+        return backend_cfg if isinstance(backend_cfg, dict) else {}
+    except Exception:
+        return {}
+
+
+def _as_int(val: Any) -> Optional[int]:
+    if val is None or val == "":
+        return None
+    try:
+        return int(val)
+    except Exception:
+        return None
+
+
+def create_kv_store(project_id: str) -> Optional[Any]:
+    """Create a KV store instance for chart subprocesses.
+
+    Priority:
+    - `MLOPS_KV_BACKEND` (if set)
+    - project config `projects/<id>/configs/project_config.yaml` (cache.backend)
+    - environment-driven heuristics
+    """
+    backend_cfg = _load_backend_cfg_from_project_config(project_id)
+    root = get_workspace_root()
+    project_root = get_projects_root(root) / project_id
+    return _create_kv_store(
+        project_id,
+        backend_cfg,
+        env=os.environ,
+        workspace_root=root,
+        project_root=project_root,
+    )
+
+
+def resolve_kv_path_from_env_or_firestore(project_id: str, run_id: str, probe_path: str) -> Tuple[str, Optional[str]]:
+    """Deprecated: kept for older chart scripts (probe IDs removed)."""
+    try:
+        from mlops.storage.path_utils import encode_probe_path  # type: ignore
+    except Exception:
+        encode_probe_path = lambda s: s  # type: ignore
+    if run_id and probe_path:
+        enc = encode_probe_path(probe_path)
+        return f"metric/{run_id}/probes_by_path/{enc}", None
+    if run_id:
+        return f"runs/{run_id}", None
+    return "", None
+
+
+__all__ = [
+    "create_kv_store",
+    "resolve_kv_path_from_env_or_firestore",
+]
+
+

mlops/reporting/registry.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TypeVar, overload
+
+ChartFn = Callable[..., Any]
+CHART_FUNCS: dict[str, ChartFn] = {}
+
+_F = TypeVar("_F", bound=ChartFn)
+
+
+@overload
+def chart(fn: _F) -> _F: ...
+
+
+@overload
+def chart(name: str | None = None) -> Callable[[_F], _F]: ...
+
+
+def chart(name: str | None | _F = None):
+    """Register a chart function by name.
+
+    The registration key is either:
+    - the function name, or
+    - the explicit `name=...`
+
+    The runner (`mlops.reporting.entrypoint`) dispatches based on `MLOPS_CHART_NAME`.
+    """
+
+    def _decorator(fn: _F) -> _F:
+        key = (name if isinstance(name, str) else getattr(fn, "__name__", "")).strip()
+        if key:
+            CHART_FUNCS[key] = fn
+        return fn
+
+    # Allow `@chart` without parentheses.
+    if callable(name) and not isinstance(name, str):
+        fn2 = name
+        name = None
+        return _decorator(fn2)  # type: ignore[arg-type]
+
+    return _decorator
+
+
+__all__ = [
+    "chart",
+    "CHART_FUNCS",
+]
+
+

mlops/runtime/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+from .context import RunContext
+
+__all__ = [
+    "RunContext",
+]
+
+

mlops/runtime/context.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+
+@dataclass(frozen=True)
+class RunContext:
+    """Typed, run-scoped context passed through orchestration/execution boundaries.
+
+    This is intentionally small in Phase 2: it centralizes the identity (project/run)
+    and key resolved paths/config snapshots so components don't rely on implicit
+    global state or environment variables.
+    """
+
+    workspace_root: Path
+    project_id: str
+    project_root: Path
+    run_id: str
+
+    runtime_python: Optional[str] = None
+    reporting_python: Optional[str] = None
+
+    cache_backend: Dict[str, Any] = field(default_factory=dict)
+    cache_config: Dict[str, Any] = field(default_factory=dict)
+    reporting_config: Dict[str, Any] = field(default_factory=dict)
+
+
+__all__ = [
+    "RunContext",
+]
+
+