Cortex Training API reference¶

This reference covers CortexTrainingClient in the customer-facing cortex_training Python namespace. Install the package, then import the client:

uv pip install git+https://github.com/snowflakedb/cortex-training.git

from cortex_training import CortexTrainingClient, SubJobConfig, JobType
from cortex_training import (
    serialize_forward_backward_args,
    build_forward_backward_kwargs,
    build_forward_backward_payload,
)

Parameter descriptions below are taken from the SDK docstrings and method signatures in CortexTrainingClient and SubJobConfig.

CortexTrainingClient¶

HTTP client for the Cortex Training API.

Recommended construction for Snowflake accounts:

client = CortexTrainingClient.from_pat(
    host="myaccount.snowflakecomputing.com",
    pat="...",
    database="MY_DB",
    schema="PUBLIC",
)

Base URL path used for requests:

https://{host}/api/v2/databases/{database}/schemas/{schema}/cortex-training

from_pat¶

Authenticate with a Programmatic Access Token (PAT).

Sends Authorization: Bearer <pat> plus the X-Snowflake-Authorization-Token-Type: PROGRAMMATIC_ACCESS_TOKEN header on every request. No login round-trip.

CortexTrainingClient.from_pat(
    host,
    pat,
    database,
    schema,
    endpoint='cortex-training',
    verify_ssl=True,
    **kwargs,
)

Constructor options¶

from_pat forwards optional constructor kwargs to __init__. You can also construct a client directly for local/mock use by passing base_url.

CortexTrainingClient(
    base_url,
    database,
    schema,
    endpoint='cortex-training',
    poll_interval=0.5,
    poll_timeout=1800.0,
    poll_backoff_multiplier=1.25,
    poll_max_interval=6.0,
    pool_maxsize=1024,
    max_retries=10,
)

Job lifecycle¶

create_job¶

Create a job from a list of sub-jobs and return its server job_id.

Each SubJobConfig is validated client-side before the request is sent (see SubJobConfig.validate). job_id is optional; when omitted the server generates one. experiment_name is optional; when omitted the server auto-creates an experiment for the job.

client.create_job(sub_jobs, job_id=None, experiment_name=None)

create_job_from_body¶

Create a job from a raw REST CreateJob request body.

This is useful for tooling that already has the REST JSON payload, while create_job remains the typed path for Python callers.

client.create_job_from_body(body)

wait_for_job¶

Poll until the job reaches RUNNING. Raises on terminal states or timeout.

client.wait_for_job(job_id)

get_job¶

client.get_job(job_id)

list_jobs¶

client.list_jobs(status=None)

cancel_job¶

client.cancel_job(job_id)

get_capacity¶

Return the calling account’s reserved GPU capacity and current usage.

This method works only for reserved GPUs. For accounts that use on-demand GPUs, has_reservation is False and the response doesn’t reflect on-demand capacity.

Backed by the account-scoped endpoint /cortex-training/capacity (not under /{job_id}). The account is resolved server-side from the caller’s session — never from a request field — so a caller can only ever read its own account’s capacity.

The returned dict includes these fields:

• has_reservation (bool): whether the account has a configured GPU reservation.
• max_total_gpus (int): the account’s GPU limit. -1 means no limit.
• reserved_gpus (int): total GPUs reserved for the account. Deprecated: use max_total_gpus with has_reservation instead.
• in_use_gpus (int): GPUs held by the account’s jobs that are placing, initializing, or running.
• pending_gpus (int): GPUs requested by jobs still waiting for capacity. These count against the limit and reduce available_gpus.
• available_gpus (int): GPUs a new job could start with now. A value of 0 doesn’t block submission: a job within the limit is queued.

client.get_capacity()

No parameters.

get_experiment_run¶

Return {experiment_name, experiment_run_name} for job_id.

Backed by the endpoint /cortex-training/{job_id}/experiment-run.

client.get_experiment_run(job_id)

Training¶

These methods submit async requests. Each returns a request_id you pass to poll_request.

forward_backward¶

Submit a forward+backward pass. Returns request_id.

client.forward_backward(job_id, data)

step¶

Submit an optimizer step. Returns request_id.

client.step(job_id, learning_rate=None)

poll_request¶

Poll until the async request completes. Returns the result dict.

client.poll_request(job_id, request_id)

get_request_status¶

client.get_request_status(
    job_id,
    request_id,
    max_events=None,
    cursor=None,
)

Checkpoints¶

save¶

Submit a checkpoint save. Returns request_id.

client.save(job_id, checkpoint_id=None, checkpoint_type=None)

load¶

Submit a checkpoint load into an existing job. Returns request_id.

client.load(
    job_id,
    checkpoint_id,
    source_job_id=None,
    target_sub_job_id=None,
)

Returns str: request_id for polling via poll_request().

list_checkpoints¶

client.list_checkpoints(job_id)

export_checkpoint¶

client.export_checkpoint(job_id, checkpoint_id)

delete_checkpoint¶

Delete a checkpoint. Returns None on success.

client.delete_checkpoint(job_id, checkpoint_id)

Inference¶

generate¶

Submit a generate request to a sampling sub-job. Returns request_id.

The typed fields are serialized as a DSSST1 safetensors frame and sent as application/octet-stream. Oversized frames are byte-chunked across multiple POSTs; the final POST returns the pollable request id. Poll with poll_request.

Prompt-length validation is asymmetric by input type:

• Pre-tokenized prompts (list[int]) are checked client-side against the sampling sub-job’s max_seq_len and fail fast with a ValueError when len(prompt) >= max_seq_len — the token count is exact, so this never rejects a prompt vLLM would accept.
• String prompts are not validated here. Counting their tokens would require replicating the server tokenizer (including special tokens), so an over-long string instead fails server-side in vLLM with a maximum model length error rather than client-side.

client.generate(
    job_id,
    prompts,
    sampling_params=None,
    routing_key=None,
    strict=None,
)

weight_sync¶

Sync model weights from a training sub-job into one or more sampling sub-jobs (used in RL training loops). Returns request_id.

Multi-sub-job sessions require the operation envelope to include a routing hint; by default the operation is routed through the source training sub-job.

client.weight_sync(
    job_id,
    source_sub_job_id,
    target_sub_job_ids,
    weight_format=None,
    sub_job_id=None,
    sub_job_type=None,
)

Model deployments¶

A model deployment serves a base model or a weights-only checkpoint as an OpenAI-compatible chat completions endpoint. These methods manage the deployment; call the endpoint itself with any OpenAI-compatible client. Only EVALUATION deployments are supported, and a deployment holds its GPUs until you drop it.

A deployment is identified by an ID of the form sf-ct-<UUID>. Pass it as model when you call the endpoint.

Deployment spec¶

create_model_deployment takes a spec with exactly three top-level fields:

{
  "deployment": {"class": "EVALUATION", "min_replicas": 1, "gpus_per_replica": 1, "gpu_type": "H200"},
  "source": {"type": "base_model", "model": "Qwen/Qwen3-1.7B"},
  "inference": {"max_seq_len": 8192}
}

A checkpoint source must name a completed weights-only checkpoint and doesn’t name a model: the client reads the model from the source training job. Which models can be deployed, and the minimum GPUs per replica each needs, is decided by the server.

The checkpoint can come from a full fine-tune or a LoRA training job. For a LoRA checkpoint, the deployment serves the base model with your trained adapter, using the LoRA settings from the training job. You can’t set LoRA settings in the spec.

Deployment statuses¶

A deployment reports one of these statuses:

Calling the endpoint¶

The endpoint is POST {base_url}/chat/completions, where base_url is https://{host}/api/v2/databases/{database}/schemas/{schema}/cortex-training/inference/v1. The client exposes it as the model_deployment_base_url property, and wait_for_model_deployment returns it as base_url. The endpoint behaves as follows:

Without "stream": true, the response is one chat.completion object.

create_model_deployment¶

Create a model deployment from a deployment spec and return its metadata, including the deployment ID. The deployment starts in PENDING; call wait_for_model_deployment to wait until it’s READY.

client.create_model_deployment(spec)

wait_for_model_deployment¶

Poll until the deployment is READY and return its metadata, including base_url. Raises ModelDeploymentWaitError if the deployment reaches FAILED, TERMINATED, or DROPPED, and ModelDeploymentWaitTimeout when the client’s poll_timeout passes. A timeout doesn’t drop the deployment, which keeps holding GPUs until you drop it.

client.wait_for_model_deployment(deployment_id)

get_model_deployment¶

Return the metadata for one deployment, including deployment, status, model_source, and the capacity fields. base_url is included unless the deployment is FAILED, DROPPING, DROPPED, or TERMINATED, and failure is included for a FAILED or TERMINATED deployment.

client.get_model_deployment(deployment_id)

describe_model_deployment¶

Return one deployment’s status, endpoint, and configuration. The configuration is grouped under spec in the same deployment, source, and inference shape as the create spec, but it’s a read view and isn’t guaranteed to be a valid create_model_deployment input. A deployment created from a checkpoint currently reports a base_model source.

client.describe_model_deployment(deployment_id)

list_model_deployments¶

Return the metadata for every deployment in the client’s database and schema. Pass status (for example, "READY" or ["PENDING", "INITIALIZING"], case-insensitive) to keep only deployments in those statuses. Entries don’t include base_url; use get_model_deployment to get it.

client.list_model_deployments(status=None)

drop_model_deployment¶

Drop a deployment and release its GPUs. Safe to retry: a deployment that’s already DROPPING or DROPPED is reported as dropped. A FAILED or TERMINATED deployment holds no GPUs and can’t be dropped, so the method returns {"dropped": False, "status": ...} instead of raising an error.

client.drop_model_deployment(deployment_id)

download_model_deployment_logs¶

Download a deployment’s logs into output_dir. log_type is "execution" (the default) or "stdout". Returns {"deployment", "logs"}, where logs lists each saved file’s filename and saved_path (plus chunk_count for stdout). Use these logs to find out why a deployment FAILED.

client.download_model_deployment_logs(deployment_id, output_dir, log_type='execution')

download_model_deployment_metrics¶

Download a deployment’s metrics into output_dir. Returns {"deployment", "metrics"}, where metrics lists each saved file’s filename, saved_path, and chunk_count.

client.download_model_deployment_metrics(deployment_id, output_dir)

Logs¶

fetch_execution_logs¶

Download every log file under the job’s experiment run stage.

Returns a list of {"sub_job_id", "filename", "s3_uri", "content"} entries, one per object found under a _logs/<sub_job_id>/ subtree in the run’s stage. Filenames such as execution.jsonl and server.log are both included; bodies are decoded as UTF-8.

Steps: experiment-run → SYSTEM$GET_VSTAGE_WRITE_CREDS → ListObjectsV2 (scoped by stage prefix) → GetObject per match.

client.fetch_execution_logs(job_id)

stream_logs¶

Yield a sub-job’s log entries, polling with backoff.

Drains the log then, with follow=True (the default), keeps polling for new entries until the caller stops iterating — suitable for a live tail. With follow=False it stops once the log reports eof.

client.stream_logs(
    job_id,
    sub_job_id=None,
    sub_job_type=None,
    cursor=None,
    max_lines=None,
    poll_interval=None,
    follow=True,
)

Returns Iterator[dict].

SubJobConfig¶

One sub-job within a CreateJob request.

training_job¶

Build a training SubJobConfig. All training fields are explicit.

SubJobConfig.training_job(
    model_name,
    optimizer,
    max_seq_len,
    train_batch_size,
    n_gpus,
    gradient_clipping=None,
    multiplex_job_id=None,
    load_optimizer_states=None,
    extra_training=None,
    global_batch_size=None,
    dtype=None,
    seed=None,
    model_post_init=None,
    source_checkpoint_info=None,
)

sampling_job¶

Build a sampling/log-probability SubJobConfig.

SubJobConfig.sampling_job(
    model_name,
    max_seq_len,
    n_gpus,
    multiplex_job_id=None,
    extra_sampling=None,
    job_type=JobType.SAMPLING,
    global_batch_size=None,
    dtype=None,
    seed=None,
    model_post_init=None,
    source_checkpoint_info=None,
)

JobType¶

class JobType(str, Enum):
    TRAINING = "training"
    SAMPLING = "sampling"
    LOG_PROBABILITY = "log_probability"

Payload helpers¶

Helpers that build the forward and backward pass payload in the client’s wire format.

serialize_forward_backward_args¶

Serialize {"args": args, "kwargs": kwargs} for Cortex Training fwd-bwd.

serialize_forward_backward_args(args=(), kwargs=None)

build_forward_backward_kwargs¶

Build tensor kwargs for a forward-backward training batch.

payload is intentionally JSON-shaped so CLI examples can stay readable: callers may provide either tokenizable texts plus a tokenizer or direct tensor data under input_ids / labels / position_ids.

build_forward_backward_kwargs(payload)

build_forward_backward_payload¶

Build a serialized fwd-bwd byte payload from readable JSON.

build_forward_backward_payload(spec)