privatemind

Official Python SDK for the PrivateMind ML platform. Submit distributed training jobs to managed Ray + KubeRay + MLflow GPU clusters from any Python environment, and promote a notebook function to the fleet without leaving Python.

pip install privatemind

In a PrivateMind notebook the SDK is pre-installed and pre-configured, so the quickstart below runs with nothing to set up.

Full documentation: docs.privatemind.com/sdk.html

Quickstart

from privatemind import submit_training

run = submit_training(
    entrypoint="python train.py --epochs 10",
    image="registry.example.com/trainer@sha256:...",
    workers=4,
    target_cluster="<your-gpu-cluster>",
    gpus_per_worker=1,
)

print(run)            # Run(name='tj-abc12', phase='Pending')
print(run.url)        # link to the platform UI

run.wait()            # block until terminal (Succeeded / Failed)
print(run.phase, run.mlflow_run_id)

Zero-config inside a notebook

A PrivateMind workspace injects everything the SDK needs as environment, so inside a notebook you do not pass a token, URL, cluster, or image:

Variable	What it provides
PRIVATEMIND_TOKEN	per-workspace bearer token
PRIVATEMIND_URL	gateway API base URL (batch, image generation)
PRIVATEMIND_APP_URL	app backend base URL (training jobs)
PRIVATEMIND_TARGET_CLUSTER	the GPU cluster the workspace runs on
PRIVATEMIND_TRAINING_IMAGE	default worker image (falls back to the platform's list when unset)
PRIVATEMIND_HOME_VOLUME / PRIVATEMIND_HOME_PATH	the home Volume used to stage promoted code

Off-cluster (laptop, CI) you supply a token (see Configuration); the URLs default to the production platform, and everything else is an explicit argument.

GPU jobs

Set gpus_per_worker and the platform pins real GPUs for you. You do not have to know which physical GPUs your org owns: when you omit gpu_placements, the platform auto-derives a placement from the GPUs your org owns that are free right now.

run = submit_training(
    entrypoint="python train.py",
    image="...",
    workers=1,
    target_cluster="<your-gpu-cluster>",
    gpus_per_worker=1,        # placement auto-derived
)

To pin exact GPUs, pass gpu_placements (one placement per GPU worker, each with exactly gpus_per_worker indices):

gpu_placements=[{"host": "<gpu-host>", "indices": [0]}]

For multi-node, set workers greater than 1: each worker gets gpus_per_worker GPUs. Omit gpu_placements and the platform places the workers for you; explicit placements support single-worker jobs. The platform is the authority on ownership, conflicts, and quota, and rejects a job that asks for GPUs it cannot have.

Promote a notebook function

@pm.train turns a function you just validated in the notebook into a distributed TrainingJob. Calling it runs locally (validate in seconds); .promote() ships it to the fleet.

import privatemind as pm

@pm.train(workers=1, gpus_per_worker=1)
def train(lr=3e-4, epochs=10):
    import torch, mlflow                  # imported on the worker
    assert torch.cuda.is_available()
    ...                                    # your training loop

train(lr=1e-3)                            # runs locally in the notebook
run = train.promote(lr=1e-3)              # runs on the fleet -> Run
run.wait()
print(run.phase, run.mlflow_run_id)

Inside a notebook you usually pass only workers + gpus_per_worker; target_cluster, image, and the home Volume come from the notebook context. .with_options(...) returns a copy with overrides (e.g. a different image) without re-decorating.

Every promote is tracked in MLflow with no setup: the run is named after the job (not a random name), CPU/GPU/memory system metrics are captured automatically, and anything you log inside the function — metrics, artifacts, mlflow.pytorch.log_model(...) — lands on that run. Pass experiment="..." to group runs under a named experiment.

Three rules make a function promotable:

1. Take all inputs as arguments. The function and its bound arguments are cloudpickled; closing over a notebook global (a DataFrame, a loaded model) either bloats the payload or fails to pickle. Pass data via a mounted Volume, not a closure.
2. Import inside the function. Imports re-resolve on the worker, so import only what the worker image provides.
3. For GPU work, the body runs on a GPU worker. A GPU promote dispatches your function onto the GPU worker automatically, so plain torch code that uses CUDA works. For multi-GPU or distributed training, use Ray Train or ray.remote inside the function as you would in any Ray program.

The Run handle

run.refresh()              # re-fetch status
run.wait(timeout=3600)     # block until terminal, then return self
run.cancel()               # delete the job (idempotent)

run.phase                  # "Pending" | "Running" | "Succeeded" | "Failed" | ...
run.mlflow_run_id          # MLflow run id once tracking starts
run.ray_job_name           # underlying RayJob name
run.start_time, run.end_time

from privatemind import list_jobs, get_job

for r in list_jobs():
    print(r.name, r.phase)

r = get_job("tj-abc12")

Reading MLflow results

Once a run has an mlflow_run_id (check run.mlflow_run_id), read its logged metrics, params, and artifacts back through the same Run handle. This is a read-only path: it never writes to the tracker, so it's safe to call from anywhere.

run = get_job("tj-abc12")

run.metrics()                # {"loss": 0.42, "acc": 0.91}  (latest values)
run.params()                 # {"lr": "0.001", "epochs": "3"} (values are str)
run.list_artifacts()         # [RunArtifact(path=..., is_dir=..., size=...)]

The tracking URI resolves in order: the tracking_uri= keyword argument, then the PRIVATEMIND_MLFLOW_URI environment variable, then a ConfigError. Pass it explicitly when you're off-platform:

run.metrics(tracking_uri="https://mlflow.internal")
run.list_artifacts(path="model")  # root the listing at a subpath

The read path needs the optional mlflow extra:

pip install 'privatemind[mlflow]'

For ad-hoc reads against any tracking server, build a client directly:

import privatemind as pm

client = pm.mlflow_client(tracking_uri="https://mlflow.internal")
mrun = client.get_run("abc123")
print(mrun.data.metrics)

Error mapping: a missing mlflow_run_id raises PrivatemindError (the job is still pending, or was submitted with mlflow=False); a run the tracker can't find raises NotFoundError; other tracker failures raise PrivatemindError.

The SDK deliberately does not auto-log your bound function arguments to MLflow. Bound args can carry secrets, and default-on logging to an org-readable tracker would invert that protection. Auto-logging will land later as an explicit, opt-in step.

Image generation

Generate images straight from Python. A prompt goes in, base64 images come back through the same gateway choke point as training (auth, audit, billing, rate-limiting).

from privatemind import generate_image

resp = generate_image(
    model="cosmos3-super-text2image",
    prompt="a tropical beach at sunset, dramatic clouds",
    size="1024x1024",
    n=4,
)

for i, img in enumerate(resp.data):
    # b64_json is a base64-encoded PNG; the SDK does not decode it for you.
    import base64
    with open(f"out_{i}.png", "wb") as f:
        f.write(base64.b64decode(img.b64_json))

print(resp.warnings)  # non-fatal notices, e.g. off-allowlist size

size is required and must be WIDTHxHEIGHT (e.g. 1024x1024). It is a required argument, so omitting it is a TypeError, and a malformed size raises ValidationError before the request leaves.

Inference knobs are forwarded verbatim; extra_args is a blind passthrough:

resp = generate_image(
    model="cosmos3-super-text2image",
    prompt="...",
    size="1024x1024",
    num_inference_steps=50,
    guidance_scale=4.0,
    flow_shift=3.0,
    negative_prompt="blurry, low quality",
    seed=1143,
    extra_args={"new_backend_param": "value"},
)

generate_image() returns a typed ImageGenerationResponse: created, data (each item exposing b64_json and revised_prompt), and warnings. The SDK returns base64 strings and leaves decoding to you.

Prompt enrichment (upsampling)

The chat UI enriches sparse prompts into dense, structured text-to-image JSON before generation. The SDK does not bundle that step — it would couple the client to a specific LLM and a schema that evolves on the backend. Wire it in yourself with any model you like:

# 1. Ask any strong LLM to expand your description into structured T2I JSON,
#    using a text-to-image template you control.
enriched_json = my_llm.complete(t2i_template.format(description="a cat"))

# 2. Pass that JSON string straight through as the prompt.
resp = generate_image(
    model="cosmos3-super-text2image", prompt=enriched_json, size="1024x1024"
)

Batch jobs

Run bulk chat-completion work asynchronously: upload a JSONL file of requests, create a batch from it, wait, and iterate the results. The flow mirrors the OpenAI batch workflow, so an author who knows that API can carry their mental model over — same parameters, same object fields, flat methods in the SDK's own style.

from privatemind import Client

# Batch + image generation run against the gateway (PRIVATEMIND_URL).
# Training jobs run against the app backend (PRIVATEMIND_APP_URL); set it
# only if you also call submit_training / list_jobs / get_job.
client = Client()  # token + base_url from env, zero-config in notebooks

f = client.create_file(file=open("in.jsonl", "rb"), purpose="batch")
batch = client.create_batch(
    input_file_id=f.id,
    endpoint="/v1/chat/completions",
    completion_window="24h",
)
batch.wait()                       # poll until terminal; pass timeout= to bound it
for line in batch.results():       # one raw dict per JSONL output line
    print(line["custom_id"], line["response"]["status_code"])

The six calls, one-to-one with the gateway's batch endpoints: create_file(file=, purpose="batch"), get_file_content(file_id), create_batch(input_file_id=, endpoint=, completion_window=, metadata=), get_batch(batch_id), list_batches(limit=, after=), cancel_batch(batch_id). Each also exists as a module-level function (from privatemind import create_batch) sharing the lazy process-wide client.

Batch.wait() blocks until the batch reaches a terminal status (completed, failed, expired, cancelled), refreshes the object in place, and returns it. timeout=None waits forever; the poll interval doubles up to a 30s cap. Batch.results() downloads the output file and yields each parsed JSON object (custom_id, response, error); Batch.error_results() does the same for the error file — the per-request failures. Before the batch is terminal both raise and tell you to wait(); on a terminal batch whose requests all failed there is no output file, so results() raises and points you at error_results().

list_batches returns one page as a plain list. Advance with after=batches[-1].id until a page comes back empty; the gateway caps limit server-side, so don't rely on len(batches) == limit to detect more pages. The gateway has no status filter — filter client-side: [b for b in client.list_batches() if b.status == "in_progress"].

Server-side limits (request count, file size, rate) are enforced by the gateway, not the SDK, so they can change without an SDK upgrade; a rejected request surfaces as a normal SDK exception. max_tokens is defaulted server-side; set it per request line if you need a different value.

Data retention: batch input and output files are content-bearing and retained server-side for up to 29 days. The batch feature is not eligible for zero-data-retention (ZDR) arrangements.

Configuration

The client resolves auth from, in order:

1. token=... kwarg on Client(...)
2. PRIVATEMIND_TOKEN environment variable
3. ~/.privatemind/auth file (must be mode 0600, owned by you, not a symlink)

The base URL comes from base_url=... or PRIVATEMIND_URL and defaults to https://api.privatemind.com, the production platform API, which serves batch and image generation. It must be https:// unless you set allow_insecure_http=True (or PRIVATEMIND_ALLOW_INSECURE_HTTP=1).

Training jobs (submit_training, list_jobs, get_job) run against the app backend, configured via app_url=... or PRIVATEMIND_APP_URL and defaulting to https://privatemind.com (same https:// rule).

Power use: explicit Client

from privatemind import Client

with Client(token="...") as pm:  # URLs default to the production platform
    for run in pm.list_jobs():
        print(run.name, run.phase)

The module-level functions share one lazily-created Client; reconfigure(...) swaps it (useful in notebooks when credentials change).

Errors

All raise subclasses of PrivatemindError: AuthError, ConfigError, ValidationError, ForbiddenError, NotFoundError, ConflictError, RateLimitError, ServerError. Client-side validation fails fast before any request.

License

Apache 2.0.