cortexgrid 0.2.85__py3-none-any.whl

cortexgrid/__init__.py

@@ -0,0 +1,195 @@
+"""cortexgrid — connect your code to the RoboLab compute cluster.
+
+import cortexgrid
+
+# Fire-and-forget training on the DGX; returns a job id immediately.
+# The jobs control plane picks up the submission and dispatches it to Ray.
+job_id = cortexgrid.remote(my_train, config, num_gpus=1, retry=True)
+print(f"Submitted: {job_id}")
+
+# Check on it later
+for job in cortexgrid.list_experiment_run_jobs(run_id):
+    status = cortexgrid.get_ray_job_status(job.get_ray_job_id())
+
+# Inside the training function — checkpoint after each epoch
+with cortexgrid.checkpoint() as ckpt:
+    ckpt.epoch = epoch
+    ckpt.save_training_state(model, optimizer, scheduler)
+
+# On resume — load checkpoint if it exists
+ckpt = cortexgrid.resume()
+if ckpt:
+    ckpt.restore_training_state(model, optimizer, scheduler)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable
+
+from cortexgrid.checkpoint import checkpoint, resume
+from cortexgrid.experiment import (
+    Experiment,
+    delete_experiment,
+    delete_run,
+    list_experiments,
+)
+from cortexgrid.infra import get_ray_job_server_uri
+from cortexgrid.jobs import (
+    schedule_remote_job,
+    list_experiment_run_jobs,
+    stop_experiment_run_jobs,
+    JobLifecycle,
+    LifecycleEvent,
+    Payload,
+)
+from cortexgrid.secrets import (
+    delete_secret,
+    get_secret,
+    list_secrets,
+    set_secret,
+)
+from cortexgrid.mlflow_util import (
+    log_metric,
+    log_metrics,
+    log_params,
+    log_artifact,
+    get_mlflow_client,
+    list_run_metrics,
+    get_metric_history,
+    list_run_params,
+    list_run_artifacts,
+)
+from cortexgrid.ray_util import (
+    get_ray_status,
+    get_ray_logs,
+    get_ray_job_url,
+    stop_ray_job,
+    submit_ray_job,
+    list_ray_jobs_with_submission_id,
+    get_ray_job_status,
+    ray_submission_id,
+    get_ray_job_attempt,
+    JobStatus,
+)
+from cortexgrid.s3_util import delete_prefix, download, get_s3_client, upload, upload_dir
+from cortexgrid.model_storage import (
+    SavedModel,
+    delete_model,
+    list_models,
+    load_model,
+    model_registry_status,
+)
+from cortexgrid.model_storage import save_model as _save_model_storage
+from cortexgrid.model_serving import (
+    Deployment,
+    ServingStatus,
+    deploy_model,
+    list_deployed_models,
+    model_serving_status,
+    undeploy_model,
+)
+
+
+def remote(
+    fn: Callable[..., Any],
+    *args: Any,
+    num_gpus: int = 0,
+    num_cpus: int = 1,
+    retry: bool = False,
+    **kwargs: Any,
+) -> str:
+    """Submit a function to the control plane. Returns a job ID."""
+    experiment = Experiment.get_instance()
+    return schedule_remote_job(
+        experiment.experiment_name,
+        experiment.run_id,
+        fn,
+        *args,
+        num_gpus=num_gpus,
+        num_cpus=num_cpus,
+        retry=retry,
+        **kwargs,
+    )
+
+
+def save_model(
+    weights_dir: str | Path, serve_app: type, family: str, suffix: str
+) -> SavedModel:
+    """Persist a weights directory under the current Experiment's run, paired
+    with the serve-app class that will front it at deploy time."""
+    experiment = Experiment.get_instance()
+    return _save_model_storage(
+        weights_dir,
+        serve_app,
+        suffix,
+        family,
+        run_id=experiment.run_id,
+        run_name=experiment.run_name(),
+    )
+
+
+__all__ = [
+    "Experiment",
+    "delete_experiment",
+    "delete_run",
+    # Ray / jobs
+    "remote",
+    "get_ray_job_status",
+    "list_experiment_run_jobs",
+    "stop_experiment_run_jobs",
+    "JobStatus",
+    "JobLifecycle",
+    "LifecycleEvent",
+    "Payload",
+    "get_ray_status",
+    "get_ray_logs",
+    "get_ray_job_url",
+    "stop_ray_job",
+    "submit_ray_job",
+    "get_ray_job_server_uri",
+    "list_ray_jobs_with_submission_id",
+    "ray_submission_id",
+    "get_ray_job_attempt",
+    # MLflow
+    "log_metric",
+    "log_metrics",
+    "log_params",
+    "log_artifact",
+    "get_mlflow_client",
+    "list_run_metrics",
+    "get_metric_history",
+    "list_run_params",
+    "list_run_artifacts",
+    "list_experiments",
+    # Checkpointing
+    "checkpoint",
+    "resume",
+    # S3
+    "upload",
+    "upload_dir",
+    "download",
+    "get_s3_client",
+    "delete_prefix",
+    # Secrets
+    "get_secret",
+    "set_secret",
+    "list_secrets",
+    "delete_secret",
+    # Model registry
+    "SavedModel",
+    "save_model",
+    "load_model",
+    "list_models",
+    "model_registry_status",
+    "delete_model",
+    # Model serving
+    "Deployment",
+    "ServingStatus",
+    "deploy_model",
+    "model_serving_status",
+    "undeploy_model",
+    "list_deployed_models",
+]
+
+# trigger: 315-trigger-tests-2026-05-16

cortexgrid/_bundle.py

@@ -0,0 +1,207 @@
+"""Code bundler: the files needed to run a piece of Python on another machine.
+
+`bundle(seed)` describes what is needed to run the module at `seed`: the local
+files it reaches -- following its import graph and each package's __init__
+chain, resolving imports the way the interpreter does -- and the third-party
+dependencies those files import. A file is local unless it lives in an installed
+package location (site-packages / dist-packages); imports that resolve into one
+are not followed. The standard library is excluded (it ships with the
+interpreter). Bundles of several seeds combine with `BundleDesc.merge`.
+
+`stage(files, dest)` lays a bundle out under `dest` at each file's import path,
+so `dest` on sys.path (e.g. a Ray working_dir) makes every module importable.
+"""
+
+from __future__ import annotations
+
+import ast
+from collections.abc import Iterator
+from dataclasses import dataclass
+import functools
+import importlib.machinery
+import importlib.util
+from pathlib import Path
+import shutil
+import sys
+
+
+ThirdPartyDependencyName = str
+ThirdPartyDependencyVersion = str
+
+@dataclass
+class BundleDesc:
+    local_files: set[Path]
+    tp_deps: dict[ThirdPartyDependencyName, ThirdPartyDependencyVersion]
+
+    def merge(self, other: BundleDesc) -> BundleDesc:
+        return BundleDesc(
+            local_files=self.local_files.union(other.local_files),
+            tp_deps={**self.tp_deps, **other.tp_deps},
+        )
+
+
+def bundle(seed: Path) -> BundleDesc:
+    """What is needed to run the module at `seed`: its local files, each at its
+    real path. Third-party dependency detection is not implemented yet, so
+    `tp_deps` is always empty."""
+    seed = seed.resolve()
+    files: set[Path] = set()
+    queue: list[Path] = [seed]
+    while queue:
+        file = queue.pop()
+        if file in files or not _is_local(file):
+            continue
+        files.add(file)
+        queue.extend(_init_chain(file))  # importing a module runs its __init__ chain
+        if file.suffix == ".py":
+            for name in _imports(file):
+                dep = _module_file(name)
+                if dep is not None:
+                    queue.append(dep)
+    return BundleDesc(local_files=files, tp_deps={})
+
+
+def stage(files: set[Path], dest: Path) -> None:
+    """Copy `files` under `dest`, each at its import path (relative to the
+    sys.path entry it lives under), so `dest` on sys.path imports them all."""
+    for file in files:
+        target = dest / file.relative_to(_sys_path_root(file))
+        target.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(file, target)
+
+
+# Distributions already present in the Ray worker image (k8s/docker/ray/Dockerfile).
+# They and their whole dependency trees -- torch's CUDA stack, sympy, numpy, ray,
+# mlflow, ... -- are on the worker already, so a caller subtracts them from a
+# bundle rather than shipping them again.
+_WORKER_BAKED = ("ray", "mlflow", "torch", "smart_open", "dotenv", "psutil")
+
+
+@functools.lru_cache(maxsize=1)
+def worker_provides() -> frozenset[Path]:
+    """Every file the Ray worker image already provides. Subtract from a bundle
+    before shipping: `bundle(entry).local_files - worker_provides()`."""
+    provided: set[Path] = set()
+    for name in _WORKER_BAKED:
+        origin = _module_file(name)
+        if origin is not None:
+            provided |= bundle(origin).local_files
+    return frozenset(provided)
+
+
+def _is_local(file: Path) -> bool:
+    """True unless `file` lives in an installed-package location."""
+    return not any(part in ("site-packages", "dist-packages") for part in file.parts)
+
+
+def _module_file(name: str) -> Path | None:
+    """The file the interpreter would load for import `name`; None for the
+    standard library, builtins, namespace packages, and names that are not
+    modules (e.g. a function in `from pkg import function`)."""
+    if name.split(".")[0] in sys.stdlib_module_names:
+        return None
+    spec = _find_spec(name)
+    if spec is None or spec.origin in (None, "built-in", "frozen"):
+        return None
+    return Path(spec.origin).resolve()
+
+
+def _find_spec(name: str) -> importlib.machinery.ModuleSpec | None:
+    """The spec the import system would find for `name`, without importing any
+    of its parent packages (importlib.util.find_spec imports them).
+
+    Mirrors the import system's own lookup: each meta path finder is asked for
+    `name`, with the parent's submodule_search_locations as the import path for
+    a submodule. https://docs.python.org/3.11/reference/import.html#the-meta-path
+    """
+    parent, _, child = name.rpartition(".")
+    path: list[str] | None = None
+    if parent:
+        parent_spec = _find_spec(parent)
+        if parent_spec is None or parent_spec.submodule_search_locations is None:
+            return None
+        path = list(parent_spec.submodule_search_locations)
+        portions = _namespace_portions(child, path)
+        if portions:
+            # The finders cannot build a nested namespace package's spec without
+            # its parent in sys.modules, so build it here.
+            spec = importlib.machinery.ModuleSpec(name, None, is_package=True)
+            spec.submodule_search_locations = portions
+            return spec
+    for finder in sys.meta_path:
+        spec = finder.find_spec(name, path)
+        if spec is not None:
+            return spec
+    return None
+
+
+def _namespace_portions(child: str, path: list[str]) -> list[str]:
+    """The directories that make `child` a namespace package under `path`: each
+    `<entry>/<child>` directory, provided no entry holds `child` as a module or
+    regular package. Empty if `child` is not a namespace package.
+    https://peps.python.org/pep-0420/#specification
+    """
+    suffixes = importlib.machinery.all_suffixes()
+    portions: list[str] = []
+    for entry in map(Path, path):
+        directory = entry / child
+        if any(
+            (entry / f"{child}{suffix}").is_file()
+            or (directory / f"__init__{suffix}").is_file()
+            for suffix in suffixes
+        ):
+            return []
+        if directory.is_dir():
+            portions.append(str(directory))
+    return portions
+
+
+def _imports(file: Path) -> Iterator[str]:
+    """Absolute module names imported by a .py file, relative imports resolved
+    against the file's own package."""
+    package = _package(file)
+    for node in ast.walk(ast.parse(file.read_text())):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                yield alias.name
+        elif isinstance(node, ast.ImportFrom):
+            if node.level:
+                try:
+                    base = importlib.util.resolve_name(
+                        "." * node.level + (node.module or ""), package
+                    )
+                except (ImportError, ValueError):
+                    continue
+            else:
+                base = node.module
+            if base:
+                yield base
+                for alias in node.names:
+                    yield f"{base}.{alias.name}"
+
+
+def _init_chain(file: Path) -> Iterator[Path]:
+    """The package __init__.py files above `file`, up to its sys.path root."""
+    directory = file.parent
+    while (directory / "__init__.py").exists():
+        yield (directory / "__init__.py").resolve()
+        directory = directory.parent
+
+
+def _package(file: Path) -> str:
+    """Dotted name of the package `file` lives in ('' at the sys.path root)."""
+    parts: list[str] = []
+    directory = file.parent
+    while (directory / "__init__.py").exists():
+        parts.insert(0, directory.name)
+        directory = directory.parent
+    return ".".join(parts)
+
+
+def _sys_path_root(file: Path) -> Path:
+    """The sys.path entry `file` is imported from: its first ancestor directory
+    without an __init__.py."""
+    directory = file.parent
+    while (directory / "__init__.py").exists():
+        directory = directory.parent
+    return directory

cortexgrid/_ray_job_driver.py

@@ -0,0 +1,48 @@
+"""Cluster-side entrypoint. Run by Ray as: python -m cortexgrid._ray_job_driver payload.pkl"""
+
+from __future__ import annotations
+
+import cloudpickle  # type: ignore
+import logging
+import sys
+from pathlib import Path
+
+from cortexgrid.checkpoint import set_cortexgrid_job_id
+from cortexgrid.experiment import Experiment
+from cortexgrid.jobs import Payload
+
+log = logging.getLogger("ray-job-driver")
+
+
+def main(payload_path: str) -> None:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s.%(msecs)03d %(levelname)s %(name)s: %(message)s",
+        datefmt="%H:%M:%S",
+    )
+
+    log.info("Loading payload: %s", payload_path)
+
+    if not Path(payload_path).exists():
+        raise FileNotFoundError(f"Payload not found: {payload_path}")
+
+    payload: Payload = cloudpickle.loads(Path(payload_path).read_bytes())
+    log.info("Payload loaded: %s", payload_path)
+    log.info("Loading experiment: %s/%s", payload.experiment_name, payload.run_id)
+
+    set_cortexgrid_job_id(payload.job_id)
+    Experiment.from_experiment(payload.experiment_name, payload.run_id)
+    log.info("Experiment loaded: %s/%s", payload.experiment_name, payload.run_id)
+
+    log.info(
+        "Starting job in experiment: %s/%s; args: %r; kwargs: %r",
+        payload.experiment_name,
+        payload.run_id,
+        payload.args,
+        payload.kwargs,
+    )
+    payload.fn(*payload.args, **payload.kwargs)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])

cortexgrid/_serve_entry.py

@@ -0,0 +1,39 @@
+"""Generic Ray Serve application builder used by every cortexgrid.deploy_model.
+
+Ray Serve's REST `import_path` resolves to `cortexgrid._serve_entry:build`.
+On the cluster replica, `build` imports the serve-app class bundled at
+`save_model` time (its import path was stored as an MLflow tag), reads its
+`num_gpus`/`num_replicas` class attributes for actor placement, wraps it as a
+Ray Serve deployment, and binds it with the (family, suffix, run_name)
+identifiers.
+
+The serve-app owns everything about traffic: its own routes, request schemas,
+streaming, and timeouts. cortexgrid does not interpose a request/response
+contract - it only schedules the app and hands it the identifiers it needs to
+fetch its own weights via `cortexgrid.load_model`.
+
+Design note: resource needs (`num_gpus`/`num_replicas`) are read from plain
+class attributes rather than a cortexgrid decorator or base class. This is a
+deliberate, provisional choice - kept minimal until we see how serve-apps
+declare resources in practice; revisit if plain class attributes prove too
+limited.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+from ray import serve
+from ray.serve.deployment import Application
+
+
+def build(args: dict[str, Any]) -> Application:
+    module_name, class_name = args["class_import_path"].split(":")
+    serve_app = getattr(importlib.import_module(module_name), class_name)
+    num_gpus = getattr(serve_app, "num_gpus", 0)
+    num_replicas = getattr(serve_app, "num_replicas", 1)
+    return serve.deployment(serve_app).options(
+        num_replicas=num_replicas,
+        ray_actor_options={"num_gpus": num_gpus},
+    ).bind(args["family"], args["suffix"], args["run_name"])

cortexgrid/checkpoint.py

@@ -0,0 +1,217 @@
+"""Durable checkpointing for cortexgrid jobs.
+
+Save arbitrary state (primitives, torch tensors, state_dicts) via MLflow
+artifacts and resume from the latest checkpoint on retry.
+
+Usage (save)::
+
+    with cortexgrid.checkpoint() as ckpt:
+        ckpt.epoch = epoch
+        ckpt.global_step = step
+        ckpt.save_training_state(model, optimizer, scheduler)
+
+Usage (resume)::
+
+    ckpt = cortexgrid.resume()
+    if ckpt:
+        ckpt.restore_training_state(model, optimizer, scheduler)
+        start_epoch = ckpt.epoch + 1
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import tempfile
+from pathlib import Path
+from typing import Any
+
+import cloudpickle  # type: ignore
+from mlflow.tracking import MlflowClient
+
+from cortexgrid import s3_util
+from cortexgrid.experiment import Experiment, get_mlflow_tracking_uri
+
+log = logging.getLogger(__name__)
+_CORTEXGRID_JOB_ID: str | None = None
+
+
+class Checkpoint:
+    """Attribute-based checkpoint persisted via MLflow artifacts.
+
+    Assign any cloudpickle-compatible value to an attribute and it will be
+    saved when the context manager exits::
+
+        with cortexgrid.checkpoint() as ckpt:
+            ckpt.epoch = 5
+            ckpt.model_state = model.state_dict()
+
+    Read values back after calling ``cortexgrid.resume()``::
+
+        ckpt = cortexgrid.resume()
+        if ckpt:
+            model.load_state_dict(ckpt.model_state)
+    """
+
+    _INTERNAL = frozenset(("_prefix", "_data"))
+
+    def __init__(self, prefix: str, *, _data: dict[str, Any] | None = None) -> None:
+        object.__setattr__(self, "_prefix", prefix)
+        object.__setattr__(self, "_data", _data if _data is not None else {})
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        if name in self._INTERNAL:
+            object.__setattr__(self, name, value)
+        else:
+            self._data[name] = value
+
+    def __getattr__(self, name: str) -> Any:
+        if name in ("_prefix", "_data", "_INTERNAL"):
+            return object.__getattribute__(self, name)
+        try:
+            return self._data[name]
+        except KeyError:
+            raise AttributeError(f"Checkpoint has no attribute {name!r}")
+
+    def __bool__(self) -> bool:
+        return bool(self._data)
+
+    def __repr__(self) -> str:
+        keys = ", ".join(sorted(self._data))
+        return f"Checkpoint(prefix={self._prefix!r}, attrs=[{keys}])"
+
+    def save_training_state(
+        self,
+        model: Any,
+        optimizer: Any,
+        scheduler: Any | None = None,
+    ) -> None:
+        """Save model, optimizer, and optionally scheduler state_dicts."""
+        self._data["_model_state"] = model.state_dict()
+        self._data["_optimizer_state"] = optimizer.state_dict()
+        if scheduler is not None:
+            self._data["_scheduler_state"] = scheduler.state_dict()
+
+    def restore_training_state(
+        self,
+        model: Any,
+        optimizer: Any,
+        scheduler: Any | None = None,
+    ) -> None:
+        """Load state_dicts back into existing model/optimizer/scheduler."""
+        model.load_state_dict(self._data["_model_state"])
+        optimizer.load_state_dict(self._data["_optimizer_state"])
+        if scheduler is not None and "_scheduler_state" in self._data:
+            scheduler.load_state_dict(self._data["_scheduler_state"])
+
+    def __enter__(self) -> Checkpoint:
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        if exc_type is None:
+            self._persist()
+
+    def _persist(self) -> None:
+        """Upload attr blobs to MinIO; log manifest.json via MLflow."""
+        exp = Experiment.get_instance()
+        client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
+        tmpdir = Path(tempfile.mkdtemp())
+
+        manifest: dict[str, Any] = {"attrs": {}}
+
+        for name, value in self._data.items():
+            filename = f"{name}.pkl"
+            (tmpdir / filename).write_bytes(cloudpickle.dumps(value))
+            uri = s3_util.upload(
+                str(tmpdir / filename), dest_path=f"{self._prefix}/{filename}"
+            )
+            manifest["attrs"][name] = {"uri": uri}
+
+        (tmpdir / "manifest.json").write_text(json.dumps(manifest))
+        client.log_artifact(
+            exp.run_id, str(tmpdir / "manifest.json"), artifact_path=self._prefix
+        )
+        log.info("Checkpoint saved: %s (%d attrs)", self._prefix, len(self._data))
+
+    @classmethod
+    def _load(cls, prefix: str) -> Checkpoint | None:
+        """Download manifest via MLflow; download attr blobs from MinIO."""
+        exp = Experiment.get_instance()
+        client = MlflowClient(tracking_uri=get_mlflow_tracking_uri())
+
+        manifest_rel = f"{prefix}/manifest.json"
+        if not any(
+            a.path == manifest_rel for a in client.list_artifacts(exp.run_id, prefix)
+        ):
+            return None
+
+        try:
+            manifest_path = client.download_artifacts(exp.run_id, manifest_rel)
+            manifest = json.loads(Path(manifest_path).read_text())
+        except Exception:
+            return None
+
+        data: dict[str, Any] = {}
+        tmpdir = Path(tempfile.mkdtemp())
+        for name, info in manifest["attrs"].items():
+            try:
+                _, _, src_path = info["uri"].removeprefix("s3://").partition("/")
+                file_path = s3_util.download(
+                    src_path, local_path=str(tmpdir / Path(src_path).name)
+                )
+                data[name] = cloudpickle.loads(Path(file_path).read_bytes())
+            except Exception:
+                log.warning("Failed to load checkpoint attribute %r", name)
+                return None
+
+        log.info("Checkpoint loaded: %s (%d attrs)", prefix, len(data))
+        return cls(prefix, _data=data)
+
+
+def set_cortexgrid_job_id(job_id: str) -> None:
+    global _CORTEXGRID_JOB_ID
+    _CORTEXGRID_JOB_ID = job_id
+
+
+def get_cortexgrid_job_id() -> str | None:
+    """Return the current job ID, or None if not running inside a job."""
+    global _CORTEXGRID_JOB_ID
+    return _CORTEXGRID_JOB_ID
+
+
+def _checkpoint_prefix() -> str:
+    job_id = get_cortexgrid_job_id() or "global"
+    return f"checkpoint/{job_id}"
+
+
+def checkpoint() -> Checkpoint:
+    """Create a checkpoint for the current job. Use as a context manager.
+
+    Returns a no-op checkpoint if not running inside a cortexgrid job,
+    so callers don't need to guard with ``if`` checks.
+
+    Example::
+
+        with cortexgrid.checkpoint() as ckpt:
+            ckpt.epoch = epoch
+            ckpt.save_training_state(model, optimizer, scheduler)
+    """
+    prefix = _checkpoint_prefix()
+    return Checkpoint(prefix)
+
+
+def resume() -> Checkpoint | None:
+    """Load the latest checkpoint for the current job, or None.
+
+    Returns None if not running inside a cortexgrid job or if no
+    checkpoint exists.
+
+    Example::
+
+        ckpt = cortexgrid.resume()
+        if ckpt:
+            ckpt.restore_training_state(model, optimizer, scheduler)
+            start_epoch = ckpt.epoch + 1
+    """
+    prefix = _checkpoint_prefix()
+    return Checkpoint._load(prefix)