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 the job stops running
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 the job stops running, then return self
run.cancel()               # delete the job (idempotent)

run.phase                  # "Running", "Succeeded", ... (the six below)
run.mlflow_run_id          # MLflow run id once tracking starts
run.ray_job_name           # underlying RayJob name
run.start_time, run.end_time

run.phase is one of six values, in the two sets wait() keys on. Both are importable from privatemind:

Set	Phases	What it means
ACTIVE_PHASES	Pending, Queued, Running	the job is on its way through the platform; wait() polls while the phase is in here
TERMINAL_PHASES	Succeeded, Failed, Terminating	the job is done or being deleted; wait() returns

A phase in neither set raises UnknownPhaseError (carrying .run_name and .phase). It means the platform has grown a phase this SDK release cannot place, so the SDK stops rather than report a job as finished when it cannot tell: upgrade privatemind, or read run.phase and decide for yourself. wait(on_unknown_phase="return") and watch(on_unknown_phase="return") downgrade it to a UserWarning and return, which is worth asking for only deliberately. A notebook that opens with warnings.filterwarnings("ignore") turns the lenient path into a silent early return, and metrics(), params() and list_artifacts() will go on serving half a run's numbers.

wait() can narrate phase transitions as they happen. progress=True prints one timestamped line per transition plus a final line with the total elapsed time; on_change is a callback invoked with the new phase string on every transition (including the terminal one):

run.wait(progress=True)
# [12:00:01] run tj-abc12: Pending
# [12:00:06] run tj-abc12: Pending -> Running
# [12:11:43] run tj-abc12: Running -> Succeeded
# [12:11:43] run tj-abc12: Succeeded after 702.0s (terminal)

run.wait(on_change=lambda phase: print("now:", phase))

In a Jupyter notebook, run.watch() renders a live HTML status table that updates in place on every phase transition, then stops at the terminal phase. When no live Jupyter display is available (IPython missing, or no shell running) it degrades to wait(progress=True). And a run — or the whole result of list_jobs() — renders as an HTML table in Jupyter on its own, with no waiting involved:

run.watch()                # live-updating display; returns the run when done

from privatemind import list_jobs, get_job

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

list_jobs(phase="Running")  # case-insensitive client-side phase filter

r = get_job("tj-abc12")

list_jobs() returns a RunList. Slicing it, or adding two together, gives another RunList and keeps the table rendering; sorted() and the other builtins that copy a list hand back a plain list, which does not.

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.

Fine-tuning

One-call LoRA/QLoRA SFT: pm.finetune(...) validates the config and dataset client-side, stages a bundled TRL/PEFT recipe onto the home Volume, and submits it as a normal TrainingJob. The adapter lands on your Volume at output_path. Paths under your home Volume are translated to their job-side /workspace/… location automatically; datasets and weights mounted on other Volumes must be referenced exactly as the job will see them (pass validate_data=False if they aren't readable from the notebook).

import privatemind as pm

run = pm.finetune(
    model="Qwen/Qwen3-8B-Instruct",        # HF id or Volume path
    dataset="sft_data.jsonl",              # readable locally for the pre-flight check
    output_path="/workspace/adapters/support-v1",
    method="qlora",                        # 4-bit nf4 base loading
    experiment="support-bot",
)

run.wait()
print(run.phase, run.mlflow_run_id)
print(run.metrics())                       # trainer metrics, e.g. {"loss": 0.42}
print(run.params())                        # the resolved config fields

pm.finetune_dry_run(...) takes the same keyword arguments, runs the same validation, and prints the resolved configuration summary without submitting anything — use it as the pre-flight check:

pm.finetune_dry_run(
    model="Qwen/Qwen3-8B-Instruct",
    dataset="sft_data.jsonl",
    output_path="/workspace/adapters/support-v1",
    method="qlora",
)

• Dataset formats: chat_jsonl (one {"messages": [{"role": ..., "content": ...}, ...]} conversation per line) and text (one {"text": ...} row per line) — auto-detected from the first row unless you pass dataset_format=; pm.dataset_formats() lists the registry. The first 50 rows are validated before submit, so a mis-shaped file raises before any GPU time is spent (sniff_rows= to change it); pass validate_data=False to skip, e.g. when the dataset is only readable on the cluster.
• Single-node in this release: workers must be 1; use gpus_per_worker for multi-GPU.
• The training image carries the dependencies: trl, peft, transformers, accelerate and datasets must be in the image; QLoRA also needs bitsandbytes, and packing (default on) needs flash-attn (or the kernels package). The recipe fails fast with a version report when something is missing, before GPU memory is allocated.
• Output: output_path receives the adapter (adapter_config.json + adapter_model.safetensors) and the tokenizer files, also logged as an MLflow artifact when mlflow=True.

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() polls while the status is in ACTIVE_BATCH_STATUSES (validating, in_progress, finalizing, cancelling) and stops on TERMINAL_BATCH_STATUSES (completed, failed, expired, cancelled), refreshing the object in place and returning it. timeout=None waits forever; the poll interval doubles, capped at 30s or the initial interval, whichever is larger. A status in neither set means the SDK is older than the platform, and raises UnknownStatusError rather than reporting a batch that may still be running as finished; pass on_unknown_status="return" to warn and return instead. 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, UnknownPhaseError, UnknownStatusError. The last two share an UnknownStateError base, so one except covers a run and a batch. Client-side validation fails fast before any request.

License

Apache 2.0.