AI Factory SDK

Python SDK for the AI Factory Compute API: submit and manage HPC jobs from Python.

Features

• Synchronous and asynchronous clients (AIFactoryClient, AsyncAIFactoryClient)
• Typed request/response models with Pydantic validation
• Job polling with configurable timeout and retry (client.wait())
• Automatic retry on transient errors (429, 5xx)
• PEP 561 compatible — full type annotation coverage
• ai-factory CLI for shell workflows (ai-factory jobs list/get/submit-container/cancel)

Installation

pip install ai-factory-sdk

Or with uv:

uv add ai-factory-sdk

Pre-release versions

Development builds published from the dev branch use PEP 440 pre-release suffixes (e.g., 0.2.0.dev1). Install them with:

pip install ai-factory-sdk --pre

Quick Start

from ai_factory.sdk import AIFactoryClient, JobRequest

# Credentials resolve from ~/.ai-factory/config.yaml, env vars, or constructor
# args (see "Configuration" below). Passed explicitly here for clarity:
with AIFactoryClient(
    api_key="dev-portal-api-key",
    slurm_token="slurm-jwt",
    slurm_user="jane",
) as client:
    # Submit a job
    resp = client.submit_job(
        JobRequest(name="hello", script="#!/bin/bash\necho Hello from SLURM")
    )
    print(f"Submitted job {resp.job_id}")

    # Wait for completion
    if resp.job_id is not None:
        detail = client.wait(str(resp.job_id), timeout=3600)
        print(f"Job finished with status: {detail.status}")

Async Usage

import asyncio
from ai_factory.sdk import AsyncAIFactoryClient, JobRequest

async def main():
    async with AsyncAIFactoryClient(
        api_key="dev-portal-api-key",
        slurm_token="slurm-jwt",
        slurm_user="jane",
    ) as client:
        resp = await client.submit_job(
            JobRequest(name="async-job", script="#!/bin/bash\nsleep 10 && echo done")
        )
        if resp.job_id is not None:
            detail = await client.wait(str(resp.job_id))
            print(detail.status)

asyncio.run(main())

Container Jobs

from ai_factory.sdk import AIFactoryClient, ContainerJobRequest

with AIFactoryClient(
    api_key="dev-portal-api-key",
    slurm_token="slurm-jwt",
    slurm_user="jane",
) as client:
    resp = client.submit_container(
        ContainerJobRequest(
            name="gpu-training",
            image="docker://nvcr.io/nvidia/pytorch:24.01-py3",
            container_command="python train.py",
            gres="gpu:a40:1",
            time_limit=120,
        )
    )

Configuration

The Compute API sits behind an APISIX gateway, so two distinct credentials are required:

• api_key: the Developer Portal API key. The SDK sends it as the apikey header so APISIX's key-auth plugin lets the request through.
• slurm_token: the Slurm JWT (scontrol token). The Compute API forwards it to the upstream Slurm REST endpoints.

To get an API key, a platform admin must first onboard your account. Then log in to the Developer Portal, generate a key, and request access to the Compute API in the AI Catalog. The gateway rejects a key that has no approved access. See Getting Started for the steps.

Credentials resolve from three sources, in priority order:

1. Explicit constructor arguments — Client(api_key=..., slurm_token=..., slurm_user=...).
2. Environment variables.
3. YAML config file at ~/.ai-factory/config.yaml.

If a required value is missing from all three sources, Client() raises ValueError with a message listing all three options.

Config file

Example ~/.ai-factory/config.yaml:

api_url: "https://aifactory.ai-factory.datalab.tuwien.ac.at/compute-api/v1"
api_key: "your-developer-portal-api-key"
slurm_token: "eyJhbGciOiJSUzI1NiIs..."   # scontrol token output
slurm_user: "jane.doe"
# slurm_cluster: "partner-si"            # optional, see "Choosing the Slurm cluster"

Secure the file so only your user can read it:

chmod 600 ~/.ai-factory/config.yaml

The SDK emits a UserWarning when a Client() is constructed if the file is group- or world-accessible. A malformed or unreadable file raises ConfigFileError (a subclass of SDKError).

Choosing the Slurm cluster

A Compute API deployment can route jobs to more than one Slurm cluster. The operator registers each extra cluster under a name. Set cluster to that name and the SDK sends it as the X-SLURM-CLUSTER header on every request:

client = AIFactoryClient(cluster="partner-si")   # or AI_FACTORY_SLURM_CLUSTER / slurm_cluster

Leave it unset to use the deployment's default cluster. Names are case-sensitive and use 1 to 32 lowercase letters, digits or hyphens; the client raises ValueError for any other name. Quote the name in the config file if YAML could read it as another type (slurm_cluster: "no"); an unquoted no, on or 42 raises ConfigFileError. A multi-cluster server rejects a name it does not know with a 400. Job responses name the job's cluster in the cluster field. A server without multi-cluster support ignores the header, runs the job on its default cluster and leaves cluster as None; submit_job and submit_container then emit a UserWarning. To override a configured cluster for one call, pass the default cluster's name (datalab on the TU Wien deployment); an empty value does not reset it.

Command-Line Interface

Installing the SDK also registers the ai-factory console script for users who prefer shell workflows or want to drive the platform from bash:

# Set credentials once (or use ~/.ai-factory/config.yaml — same resolution chain as Client())
export AI_FACTORY_API_KEY="your-developer-portal-api-key"     # -> apikey header
export AI_FACTORY_SLURM_TOKEN="eyJhbGciOi..."                 # -> X-SLURM-USER-TOKEN
export AI_FACTORY_SLURM_USER="jane.doe"

ai-factory --version                    # print SDK version and exit
ai-factory jobs list                    # table output
ai-factory jobs list --json             # machine-readable
ai-factory jobs get 459381              # single job detail
ai-factory jobs submit-container \
    --name training-run \
    --image docker://nvcr.io/nvidia/pytorch:24.01-py3 \
    --command "python train.py" \
    --partition GPU-a100 \
    --gres gpu:a40:1 \
    --time-limit 120
ai-factory jobs cancel 459381
ai-factory jobs list --cluster partner-si   # a registered non-default cluster

Every subcommand supports --help and --json. Errors map to distinct exit codes so shell scripts can branch on the failure mode (codes start at 10 so they do not collide with Click/Typer's argument-parse exit 2):

The CLI resolves credentials and the cluster the same way as AIFactoryClient, with --cluster in place of the constructor argument: --cluster beats AI_FACTORY_SLURM_CLUSTER, which beats slurm_cluster in ~/.ai-factory/config.yaml. jobs list prints a Cluster: line above the table when the server reports the cluster. SDK warnings, such as a cluster the server ignored, print as one Warning: line on stderr.

There is no ai-factory jobs wait subcommand yet. Poll with jobs get --json in a script until .status is one of completed / errored / cancelled, or use the Python client.wait() method directly.

API Reference

Clients

Methods

Request Models

Response Models

Exceptions

Requirements

• Python >= 3.11
• httpx >= 0.28
• pydantic >= 2.0

End-to-end verification

A SDK-driven Path 2 test lives at test/e2e/test_sdk_path2.py in the monorepo. It submits a real container job through the published-shape AIFactoryClient, polls until terminal state, and validates the JobDetail schema.

Run it locally against staging:

export COMPUTE_API_URL="https://aifactory-dev.ai-factory.datalab.tuwien.ac.at/compute-api/v1"
export APISIX_CI_API_KEY="your-developer-portal-api-key"    # sent as the apikey header
export SLURM_USERNAME="$(whoami)"
export SLURM_USER_TOKEN="$(scontrol token | cut -d= -f2)"   # Slurm JWT, rotates often

uv run pytest test/e2e/test_sdk_path2.py -v -s -m sdk_e2e

In CI, trigger the manual sdk-e2e job in the post-deploy-verify stage. The job is intentionally manual because each run queues a real Slurm job.

License

MIT