Metadata-Version: 2.5
Name: cortexgrid
Version: 0.2.88
Summary: Connect your ML code to the RoboLab compute cluster — Ray, MLflow, and S3
Project-URL: Homepage, https://github.com/robodatalab/cortexgrid
Project-URL: Repository, https://github.com/robodatalab/cortexgrid
Project-URL: Documentation, https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/README.md
License-Expression: Apache-2.0
License-File: LICENSE
Requires-Python: >=3.11
Requires-Dist: boto3>=1.34
Requires-Dist: cloudpickle>=3.0
Requires-Dist: fabric>=3.2.3
Requires-Dist: haikunator>=2.1.0
Requires-Dist: mlflow<4,>=3.11
Requires-Dist: packaging>=24
Requires-Dist: pip>=23.0
Requires-Dist: pydantic-settings>=2.13.1
Requires-Dist: pydantic>=2.13.3
Requires-Dist: pydotenv>=0.0.7
Requires-Dist: python-dotenv>=1.2.2
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: ray[default]<3,>=2.9
Requires-Dist: requests>=2.31
Requires-Dist: setuptools>=82.0.1
Requires-Dist: tqdm>=4.60
Description-Content-Type: text/markdown

# cortexgrid

`cortexgrid` is a Python library that connects your ML code to the deployed infrastructure. It wraps Ray, MLflow, and S3/MinIO so your training scripts don't need to know about URLs, credentials, or service endpoints.

### Installation

Add `cortexgrid` as a dependency in your project's `pyproject.toml`:

```toml
[project]
dependencies = [
    "cortexgrid",
]
```

Then `uv sync` to install it, and point it at the head's secrets server (reachable over the tailnet):

```bash
export CORTEXGRID_HEAD_URL=http://robolab-head:7700
```

### Usage

```python
import cortexgrid

cortexgrid.init(experiment="weather-forecast")
```

That single call reads the service URLs from the head's secrets server at `$CORTEXGRID_HEAD_URL` and connects to all services through them. It also creates (or finds) the named MLflow experiment and starts a new run inside it. Omit `experiment=` to auto-generate a unique name like `funky-koval-12`.

**One experiment per binary run.** `cortexgrid.init()` may only be called once per process. Every subsequent `cortexgrid.log_metric`, `cortexgrid.log_artifact`, checkpoint, and `cortexgrid.remote()` submission is scoped to that experiment+run. Remote jobs dispatched by the control plane inherit the experiment+run via the pickled payload, so their logging flows into the same MLflow run as the parent binary.

#### Experiment tracking (MLflow)

```python
cortexgrid.init(experiment="weather-forecast")

cortexgrid.log_params({"lr": 1e-3, "epochs": 20, "batch_size": 64})

for epoch in range(20):
    loss = train_one_epoch(model, dataloader)
    cortexgrid.log_metric("loss", loss, step=epoch)

    if epoch % 5 == 0:
        with cortexgrid.checkpoint() as ckpt:
            ckpt.epoch = epoch
            ckpt.save_training_state(model, optimizer)
```

No run-scoping context manager — `init()` starts the run, and every subsequent logging call flows into it. Metrics and artifacts are logged to the MLflow server on the DGX. View them at `http://<DGX_IP>:5000`.

#### Checkpointing and resuming

Inside a cortexgrid job, `cortexgrid.checkpoint()` returns an attribute-based checkpoint object that persists to MLflow artifacts when its `with` block exits. On job restart (either manual retry or `retry=True`), `cortexgrid.resume()` returns the last checkpoint for the same job ID, or `None` if there isn't one.

```python
ckpt = cortexgrid.resume()
if ckpt:
    ckpt.restore_training_state(model, optimizer)
    start_epoch = ckpt.epoch + 1
else:
    start_epoch = 0

for epoch in range(start_epoch, 20):
    train_one_epoch(model, dataloader)
    with cortexgrid.checkpoint() as ckpt:
        ckpt.epoch = epoch
        ckpt.save_training_state(model, optimizer)
```

You can assign any cloudpickle-compatible or torch-serializable value as an attribute on the checkpoint (`ckpt.metric = 0.93`, `ckpt.weights = model.state_dict()`); the `save_training_state`/`restore_training_state` helpers are a shortcut for the common model+optimizer pair.

#### Distributed compute (jobs control plane)

```python
def train_step(batch):
    # runs on the DGX GPU
    # MLflow and S3 env vars are injected automatically
    return loss

job_id = cortexgrid.remote(train_step, batch, num_gpus=1, retry=True)
print(f"Submitted: {job_id}")
```

`cortexgrid.remote` submits a job *request* (a pickled payload plus a `JobLifecycle` record) to MLflow and returns a job ID string immediately. It does not wait for the job to run or finish — use the UI at `http://<DGX_IP>:8000`, or poll `cortexgrid.list_experiment_run_jobs(run_id)`, to observe status.

A separate service — the **jobs control plane** — polls MLflow for pending job requests, matches them against the set of Ray submissions the cluster already has, and submits anything missing. It is also responsible for retrying failed jobs and honouring user-requested stops.

Each submission captures the code and dependencies the entry function needs automatically ([_bundle.py](https://github.com/robodatalab/cortexgrid/blob/main/cortexgrid/_bundle.py)):
- `bundle(entry)` traces the import graph from the function's source file, resolving each import the way the interpreter does (via `sys.path`). The standard library is excluded (it ships with the interpreter)
- Your own modules -- anything outside site-packages / dist-packages -- ship **as source**: they are staged at their import paths and tarred into the Ray `working_dir`
- Third-party packages are recorded as the installed distribution that owns the imported file, pinned to its installed version (`tqdm==4.67.3`), and Ray **pip-installs** them on the worker into a per-node cached virtualenv layered on the image (`runtime_env["pip"]`). An import into site-packages that no installed distribution owns fails `cortexgrid.remote` with `UnownedDependencyError`
- Distributions the worker image already has are not installed again: `worker_provides()` is the dependency closure of the packages baked into the ray image (torch and its CUDA stack, ray, mlflow, ...), and is subtracted from the pip list. See [k8s/docker/ray/Dockerfile](https://github.com/robodatalab/cortexgrid/blob/main/k8s/docker/ray/Dockerfile) and `_WORKER_BAKED` in `_bundle.py`, which must list the same packages
- Injects MLflow/S3 credentials so task code running on the DGX can reach all services

##### Retries

Pass `retry=True` and the control plane will resubmit the job whenever Ray reports the most recent attempt as `FAILED`. Retries are **unbounded by design**: the intended way to end a retry loop is to stop the job manually from the UI (which flips the `stop_requested` latch on the lifecycle, and the control plane stops the current Ray attempt on its next poll). This keeps the retry policy simple — you don't have to predict a good `max_retries` up front — and puts the human in the loop for anything that's failing persistently.

##### Stopping a job

```python
cortexgrid.stop_experiment_run_jobs(run_id)   # stops every job in the run
```

`stop_experiment_run_jobs` never touches Ray directly. It only flips `stop_requested` on each job's lifecycle record in MLflow. The control plane observes the flag on its next poll and calls `ray.stop_job` for any attempt that has reached Ray. For jobs that have not yet been submitted, the same flag short-circuits the submission path inside the worker.

#### Object storage (S3/MinIO)

```python
cortexgrid.upload("data/output.parquet", bucket="ray-checkpoints", key="run-42/output.parquet")
cortexgrid.download("ray-checkpoints", "run-42/output.parquet", local_path="./output.parquet")

# or get the raw boto3 client
s3 = cortexgrid.get_s3_client()
```

Works with MinIO on the DGX today, real S3 on AWS tomorrow — same code.

#### Getting raw clients

```python
mlflow_client = cortexgrid.get_mlflow_client()   # mlflow.tracking.MlflowClient
s3_client = cortexgrid.get_s3_client()           # boto3 S3 client
```

#### Model registry and serving

Save a trained model's weights together with the serve-app that fronts it, then deploy it as a Ray Serve application. A serve-app is a class fronted by a FastAPI app, marked with cortexgrid's `serve.ingress` (not Ray's):

```python
from cortexgrid import serve
from fastapi import FastAPI

app = FastAPI()

@serve.ingress(app)
class MyServeApp:
    num_gpus = 1

    def __init__(self, family: str, suffix: str, run_name: str) -> None:
        self._weights_dir = cortexgrid.load_model(family, suffix, run_name)

    @app.post("/complete")
    async def complete(self, body: dict): ...

saved = cortexgrid.save_model(weights_dir, MyServeApp, family="qwen", suffix="instruct")
deployed = cortexgrid.deploy_model("qwen", "instruct", saved.run_name, wait=True)
print(deployed.url)
```

`save_model` is synchronous (registry lifecycle: `uploading` -> `ready`); `deploy_model` schedules the serving lifecycle (`deploying` -> `running`). See [model-serving.md](https://github.com/robodatalab/cortexgrid/blob/main/docs/cortexgrid/model-serving.md) for both lifecycles end to end - upload/deploy/undeploy/delete, status queries (`model_registry_status`, `model_serving_status`), and error handling.

### API reference

| Function | Description |
|----------|-------------|
| `cortexgrid.init(experiment=None)` | Configure connections + start a new MLflow run inside the named experiment. One call per binary. |
| `cortexgrid.log_metric(key, value, step)` | Log a metric |
| `cortexgrid.log_metrics(metrics, step)` | Log multiple metrics |
| `cortexgrid.log_params(params)` | Log parameters |
| `cortexgrid.log_artifact(path, artifact_path)` | Log a file as an artifact |
| `cortexgrid.checkpoint()` | Context manager returning an attribute-based checkpoint saved to MLflow on exit |
| `cortexgrid.resume()` | Load the latest checkpoint for the current job, or `None` |
| `cortexgrid.remote(fn, *args, num_gpus=0, num_cpus=1, retry=False, **kwargs)` | Submit a function to the jobs control plane; returns a job ID |
| `cortexgrid.list_experiment_run_jobs(run_id)` | List `JobLifecycle` records for every cortexgrid job in a run |
| `cortexgrid.stop_experiment_run_jobs(run_id)` | Request every job in a run to stop (flips the `stop_requested` latch) |
| `cortexgrid.get_ray_job_status(ray_job_id)` | Live Ray status for a submission id |
| `cortexgrid.get_ray_logs(ray_job_id)` | Tail the stdout/stderr of a Ray submission |
| `cortexgrid.upload(path, bucket, key)` | Upload a file to S3/MinIO |
| `cortexgrid.download(bucket, key, path)` | Download a file from S3/MinIO |
| `cortexgrid.get_mlflow_client()` | Raw configured MLflow client |
| `cortexgrid.get_s3_client()` | Raw configured boto3 S3 client |

## ML compute stack

The DGX Spark runs the following services as k8s workloads managed by Argo CD (see [../k8s/argo_deployments/](https://github.com/robodatalab/cortexgrid/tree/main/k8s/argo_deployments/)):

| Service | Port | Purpose |
|---------|------|---------|
| Ray | 8265 | Dashboard + job submission (NodePort 30265) |
| MLflow | 5000 | Experiment tracking, model registry |
| MinIO | 9000/9001 | S3-compatible artifact storage |
| PostgreSQL | 5432 | MLflow metadata backend |
| Prometheus | 9090 | Metrics collection |
| Grafana | 3000 | Dashboards (GPU, jobs, system) |

