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:

[project]
dependencies = [
    "cortexgrid",
]

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

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

Usage

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)

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.

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)

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):

• bundle(entry) traces the import graph from the function's source file, resolving each import the way the interpreter does (via sys.path), and returns every file needed to run it -- your own modules and third-party packages alike, wherever they live. The standard library is excluded (it ships with the interpreter)
• Everything ships as source: the bundle is staged at each file's import path and tarred into the Ray working_dir. Nothing is pip-installed on the worker
• Dependencies the worker image already has are subtracted rather than shipped: bundle(entry) - worker_provides(), where worker_provides() is the bundle of the packages baked into the ray image (torch and its CUDA stack, ray, mlflow, ...). See k8s/docker/ray/Dockerfile
• 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

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)

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

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:

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 for both lifecycles end to end - upload/deploy/undeploy/delete, status queries (model_registry_status, model_serving_status), and error handling.

API reference

ML compute stack

The DGX Spark runs the following services as k8s workloads managed by Argo CD (see ../k8s/argo_deployments/):