Kaio Python Client

Kaio is an MLOps platform, built to take machine learning from your IDE to production. Run your Python scripts and notebooks on cloud CPUs and GPUs: point Kaio at your code, pick an instance, and download the results when the job finishes. Follow every job on the Kaio dashboard.

Installation

pip install kaio

Quick Start

from kaio import Client

# Initialize client
client = Client("https://api.kaion5.com")

# Login with your API key
client.login("your-api-key")

# Submit a job
result = client.submit_job(
    directory="./my_code",
    job_name="training-job",
    instance_type="ml.g4dn.xlarge",
    entrypoint="train.py"
)

Data

A job's code directory is uploaded with every submit. Data reaches a job in two other ways: Kaio Datasets, uploaded once and used by any job, or your code loading it from where it's stored.

Kaio Datasets

Upload a file or folder once, then use it in any number of jobs:

client.upload_dataset("sensor-readings", "./data/readings")   # a folder or a single file

client.submit_job(directory="./my_code", entrypoint="train.py", input_data="sensor-readings")

The job finds the files in datasets/sensor-readings/, beside your code, so open("datasets/sensor-readings/day1.csv") works the same locally when your data is there (with datasets/ listed in .kaioignore).

• Kept until you delete them, unlike job files, which go after 7 days.
• Several per job: input_data=["a", "b"], up to 5. They must fit in half of volume_size_gb, packed and unpacked.
• Replacing: uploading the same name again replaces the dataset once the new upload finishes.
• Never in your outputs: datasets aren't copied to the job's output.
• On the dashboard: each dataset's size, files and the jobs that used it.

Loading Data in Your Code

Jobs have internet access, so your code loads data the same way it does on your machine, for example from the Hugging Face Hub:

my_code/
├── train.py
├── requirements.txt   # datasets, transformers
└── .env               # HF_TOKEN=hf_... (never uploaded)

# train.py
from datasets import load_dataset
from transformers import AutoModelForSequenceClassification

dataset = load_dataset("your-name/your-dataset", revision="<commit>")
model = AutoModelForSequenceClassification.from_pretrained("your-name/your-model", revision="<commit>")

• Credentials: tokens and keys go in the directory's .env or env=, for example HF_TOKEN=hf_..., which Hugging Face libraries read automatically. They're sent separately from your code, kept encrypted only while the job runs, masked as **** in logs, and deleted when the job ends. See Environment Variables and Tokens.
• Writing back: your code can push results with the same credentials, for example model.push_to_hub("your-name/your-model") with a write token.

Leaving Files Out: .kaioignore

A .kaioignore file in the code directory lists what not to upload, one pattern per line, like .gitignore:

# not uploaded
data/
*.parquet
checkpoints/
.venv/

A pattern without a / matches a name at any depth, one with a / matches from the code directory, and a trailing / matches folders only. Lines starting with # are comments. Unlike .gitignore, ! patterns aren't supported. .env files are never uploaded either way.

submit_job warns when the upload is over 100 MB or holds more than 10 MB of data files or model weights, and names the largest of them.

Features

• Scripts and notebooks: run a .py file or a whole .ipynb notebook, as a processing job or a training job
• CPUs and GPUs: pick the instance; the job's image matches your local PyTorch or TensorFlow, and defaults to PyTorch on Python 3.12
• Several instances at once: split work across instances with instance_count
• Kaio Datasets: upload data once and use it in any job
• Internet access: jobs load data and models from wherever your code reaches, such as the Hugging Face Hub
• Tokens and keys: pass them as environment variables; they're never uploaded and are deleted when the job ends
• Only your code is uploaded: .env files never are, and .kaioignore leaves out anything else
• Follow your jobs: status, logs and outputs from Python or on the dashboard
• Projects: each API key belongs to a project, and the jobs you submit with it are that project's
• Your requirements.txt: installed in the job as it is

API Reference

Client

Client(api_base, verbose=False)

Initialize the Kaio client.

Parameters:

• api_base (str): Base URL of the Kaio API endpoint
• verbose (bool): Enable verbose logging for debugging. Defaults to False.

Example:

client = Client("https://api.kaion5.com", verbose=True)

login(api_key)

Log in with your API key. The session renews itself, so call this once.

Parameters:

• api_key (str): Your Kaio API key, from the dashboard's API Keys page

Returns:

• Client: Self for method chaining

Raises:

• requests.HTTPError: If authentication fails

client.project

The project of the key you logged in with, as {"id": ..., "name": ...}. Every job you submit with this key belongs to that project. See Projects.

submit_job(**kwargs)

Run your code as a processing or training job. The SDK packages and uploads the directory, picks the image, and starts the job.

Parameters:

• directory (str): Path to your code directory. Everything in it is uploaded with each job, except .env files and what .kaioignore lists. Defaults to the current directory.
• job_name (str): Unique name for the job. Defaults to "job".
• instance_type (str): Instance type, such as ml.m5.large (CPU) or ml.g4dn.xlarge (GPU). Defaults to ml.m5.large.
• instance_count (int): Number of instances for parallel processing. Defaults to 1.
  • Multi-Instance Processing: When instance_count > 1, Kaio starts that many instances, each running your code, to process large datasets in parallel
  • Environment Variables: Each instance receives SM_CURRENT_HOST (current instance identifier) and SM_HOSTS (comma-separated list of all instance identifiers)
  • No Inter-Instance Networking: Instances run independently without network communication between them
  • Custom Parallelization Logic: Use SM_CURRENT_HOST and SM_HOSTS to implement data partitioning and parallel processing
• volume_size_gb (int): EBS volume size in GB. Defaults to 5. Maximum 50GB.
• entrypoint (str): Main script to execute (.py or .ipynb). Defaults to "train.py".
• input_data (str or list, optional): The name of a dataset from upload_dataset(), or a list of up to 5. See Kaio Datasets.
• framework (str, optional): ML framework ("pytorch", "tensorflow", "sklearn").
• framework_version (str, optional): Framework version.
• image_uri (str, optional): Custom Docker image URI. If provided, overrides framework resolution.
• workflow (str): Workflow type - "processing" or "training". Defaults to "processing".
• training_args (dict, optional): Advanced training configuration for training workflow.
• max_runtime_hours (float): Hours before the job is stopped. Defaults to 1, at most 24.
• env (dict, optional): Environment variables for your code, such as tokens and API keys. See Environment Variables and Tokens.
• env_file (str, optional): A .env file with more variables. Defaults to the directory's own .env, if it has one.

Returns:

• dict: Job submission result with status, job_name, entrypoint and env_names

Raises:

• requests.HTTPError: If the API refuses the job; the message gives the API's reason
• FileNotFoundError: If directory or entrypoint doesn't exist
• ValueError: If code package exceeds volume capacity, or an environment variable isn't allowed

get_jobs(include_archived=False, project_id=None)

Get all jobs for the authenticated user.

Parameters:

• include_archived (bool): Also list archived jobs. Defaults to False.
• project_id (str, optional): Only list this project's jobs. Defaults to every project's jobs.

Returns:

• dict: List of jobs with count. Each job has its unique_job_name, status, timestamp, project_id and workflow (processing or training).

Raises:

• requests.HTTPError: If API error

Example:

jobs = client.get_jobs()
print(f"Total jobs: {jobs['count']}")

get_job(unique_job_name)

Get specific job status and details.

Parameters:

• unique_job_name (str): Unique job name returned from submit_job

Returns:

• dict: status, status_history (each status with the time it began), timestamp, project_id, workflow and env_names (the names of the job's environment variables, never their values)

A job's status moves through QUEUED, PROVISIONING, RUNNING and ends as COMPLETED, FAILED, STOPPED or TIMED_OUT (it ran past max_runtime_hours). A stopped job passes through STOPPING.

Raises:

• requests.HTTPError: If job not found or API error

Example:

job = client.get_job("unique-job-name-123")
print(f"Status: {job['status']}")

delete_job_data(unique_job_name)

Delete a finished job's input, output and logs now, instead of after the usual 7 days. The job stays in your job list, marked data_deleted. This can't be undone.

archive_job(unique_job_name) and restore_job(unique_job_name)

Hide a finished job from get_jobs() and the dashboard's default list, or show it again. Its data is unchanged. get_jobs(include_archived=True) lists archived jobs too.

Example:

client.delete_job_data("unique-job-name-123")
client.archive_job("unique-job-name-123")
jobs = client.get_jobs(include_archived=True)

upload_dataset(name, path), list_datasets() and delete_dataset(name)

Upload a file or folder once as a dataset for jobs to use with input_data=name, list your datasets and their sizes, or delete one. A name is 1-50 letters, digits, ., _ or -. Datasets are kept until you delete them. See Kaio Datasets.

client.upload_dataset("sensor-readings", "./data/readings")
client.list_datasets()   # {"datasets": [{"name", "size_bytes", "uploaded_at"}], ...}
client.delete_dataset("sensor-readings")

download_output(unique_job_name, output_dir=".")

Download completed job output files.

Parameters:

• unique_job_name (str): Unique job name returned from submit_job
• output_dir (str): Local directory to save output. Defaults to current directory.

Returns:

• list: List of downloaded file paths

Raises:

• RuntimeError: If job is not completed
• requests.HTTPError: If download fails

Example:

files = client.download_output("unique-job-name-123", "./results")
for file_path in files:
    print(f"Downloaded: {file_path}")

Training Arguments

For training workflow jobs, use the training_args parameter to configure advanced training features:

Spot Instances

Training jobs can run on spot instances: AWS's spare capacity, which can take longer to start (up to max_wait seconds) and can be interrupted. An interrupted job carries on on a new instance, and your code can resume from the files it saved in /opt/ml/checkpoints.

training_args = {
    "spot_instances": {
        "enabled": True,
        "max_wait": 3600  # Max wait time in seconds
    }
}

Environment Variables and Tokens

Give your code tokens and keys, such as a Hugging Face token for a private dataset, without putting them in your code or uploading them:

client.submit_job(directory="./training_code", env={"HF_TOKEN": "hf_..."})

# Or from a .env file. The directory's own .env is used automatically.
client.submit_job(directory="./training_code", env_file="secrets.env")

Your code reads them as ordinary environment variables, for example os.environ["HF_TOKEN"]. They're also set while your requirements.txt installs, so a private package index token works.

• Never uploaded: .env files are left out of the code package.
• Kept only while the job runs: the variables travel over HTTPS, are stored encrypted, and only your own jobs can read them. They're deleted as soon as the job ends, however it ends.
• Masked in logs: a value of 8 or more characters is replaced with **** wherever it's printed in the job's logs. Values your code writes into output files aren't masked; outputs are kept for 7 days.
• Limits: up to 50 variables and 4 KB in total. Names may not start with SM_, SAGEMAKER_ or KAIO_, or replace PATH-like variables.

Framework Auto-Detection

The SDK picks the job's image from your local environment:

• PyTorch installed: the AWS PyTorch image closest to your version, up to 2.9 (Python 3.12)
• TensorFlow installed: the closest AWS TensorFlow image, up to 2.19 (Python 3.12)
• Neither (including scikit-learn only): PyTorch 2.9 on Python 3.12, which also has scikit-learn, pandas, NumPy and SciPy. It's also the image for Hugging Face libraries; add datasets or transformers to your requirements.txt
• framework="sklearn": AWS's scikit-learn image (1.4, Python 3.10)

Code Requirements

Data and Code Structure

The directory you submit with client.submit_job() is packaged and uploaded with every job, except .env files and what .kaioignore lists. See Data.

Security: Jobs run in Kaio's cloud account on AWS. Don't upload sensitive or confidential data. A job's code, data and outputs are deleted after 7 days; datasets you upload with upload_dataset() are kept until you delete them.

Job records: Kaio keeps a record of each job (name, instance type, status, run time, compute metrics and storage settings) for billing and to improve the service. Its code, data, logs and outputs are deleted after 7 days.

File Size Limits

Code packages must not exceed half your volume size:

• 5GB volume → 2.5GB max code package
• 10GB volume → 5GB max code package
• Maximum storage per job: 50GB
• Maximum upload: 5GB

Multi-Instance Code Structure

For multi-instance processing jobs, your code can use these environment variables to split the work:

import os

def main():
    # Which instance this is, and all of them
    current_host = os.environ.get('SM_CURRENT_HOST', 'algo-1')
    all_hosts = os.environ.get('SM_HOSTS', 'algo-1').split(',')
    
    # Determine this instance's role
    host_rank = all_hosts.index(current_host)
    total_hosts = len(all_hosts)
    
    print(f"Running on {current_host} (rank {host_rank}/{total_hosts})")
    
    # Each instance processes its own share of the data
    process_data_partition(host_rank, total_hosts)

if __name__ == "__main__":
    main()

Requirements File

A requirements.txt in the code directory is installed in the job before your code runs.

Automatic Dependencies

The job runner installs these in the job environment (your requirements.txt is not modified):

• nbconvert - For Jupyter notebook execution
• psutil - For system monitoring
• GPUtil - For GPU monitoring

Workflows

Processing Workflow

A processing job runs your script or notebook on one or more instances. With instance_count > 1, each instance runs your code, and your code splits the work between them.

Environment Variables: Each instance receives SM_CURRENT_HOST and SM_HOSTS for custom parallelization logic.

Limitations: No inter-instance communication.

Training Workflow

Training jobs add:

• Spot instances: see Spot Instances
• Checkpoints: files your code writes to /opt/ml/checkpoints are saved while the job runs, so it can resume after an interruption

Multi-Instance Processing Example

import os

# Get instance information
current_host = os.environ.get('SM_CURRENT_HOST', 'algo-1')
all_hosts = os.environ.get('SM_HOSTS', 'algo-1').split(',')
host_rank = all_hosts.index(current_host)
total_hosts = len(all_hosts)

print(f"Instance {host_rank + 1} of {total_hosts} (Host: {current_host})")

# Partition data based on host rank
data_partition_size = len(dataset) // total_hosts
start_idx = host_rank * data_partition_size
end_idx = start_idx + data_partition_size if host_rank < total_hosts - 1 else len(dataset)

# Process assigned data partition
my_data = dataset[start_idx:end_idx]
process_data_partition(my_data)

Examples

Processing Workflow Examples

from kaio import Client

client = Client("https://api.kaion5.com")
client.login("your-api-key")

# Single instance data processing
result = client.submit_job(
    directory="./data_processing",
    job_name="data-preprocessing",
    workflow="processing",
    instance_type="ml.m5.xlarge",
    entrypoint="preprocess.py"
)

# Multi-instance parallel processing
result = client.submit_job(
    directory="./parallel_processing",
    job_name="multi-instance-processing",
    workflow="processing",
    instance_type="ml.m5.2xlarge",
    instance_count=4,
    entrypoint="process_data.py",
    volume_size_gb=50
)

Training Workflow Examples

# Basic training job
result = client.submit_job(
    directory="./training_code",
    job_name="model-training",
    workflow="training",
    instance_type="ml.g4dn.xlarge",
    entrypoint="train.py"
)

# Training on spot instances
result = client.submit_job(
    directory="./training_code",
    job_name="spot-training",
    workflow="training",
    instance_type="ml.g4dn.xlarge",
    training_args={
        "spot_instances": {
            "enabled": True,
            "max_wait": 7200  # 2 hours max wait
        }
    }
)

Framework-Specific Examples

# TensorFlow with custom image
result = client.submit_job(
    directory="./tensorflow_training",
    job_name="tf-training",
    workflow="training",
    instance_type="ml.g4dn.xlarge",
    framework="tensorflow",
    framework_version="2.13.0",
    entrypoint="train.py"
)

# Jupyter notebook execution
result = client.submit_job(
    directory="./notebooks",
    job_name="notebook-analysis",
    workflow="processing",
    instance_type="ml.m5.xlarge",
    entrypoint="analysis.ipynb"
)

Job Management

Getting All Jobs

# Get list of all your jobs
jobs = client.get_jobs()
print(f"Total jobs: {jobs['count']}")
for job in jobs['jobs']:
    print(f"Job: {job['unique_job_name']} - Status: {job['status']}")

Getting Job Status

# Get specific job details (use unique_job_name from submit_job response)
job_details = client.get_job("unique-job-name-123")
print(f"Status: {job_details['status']}")
for change in job_details['status_history']:
    print(f"{change['at']}  {change['status']}")

Projects

Each API key belongs to a project, which you create and name on the dashboard's API Keys page; keys from before projects belong to Default. The key alone decides a job's project: every job you submit with it, processing and training alike, belongs to that project and belongs to that project. To file jobs under another project, log in with one of that project's keys.

client.login(api_key)
print(f"Submitting to {client.project['name']}")

# This project's jobs, processing and training together
for job in client.get_jobs(project_id=client.project["id"])["jobs"]:
    print(job["workflow"], job["unique_job_name"], job["status"])

get_jobs() without project_id still lists every project's jobs.

Downloading Job Outputs

# Download all output files from a completed job
downloaded_files = client.download_output("unique-job-name-123", "./results")
for file_path in downloaded_files:
    print(f"Downloaded: {file_path}")

Complete Job Workflow

# Submit training job
result = client.submit_job(
    directory="./my_training",
    job_name="model-training",
    workflow="training",
    instance_type="ml.g4dn.xlarge"
)

unique_job_name = result['unique_job_name']
print(f"Job submitted: {unique_job_name}")

# Monitor job status
import time
while True:
    job = client.get_job(unique_job_name)
    status = job['status']
    print(f"Job status: {status}")
    
    if status in ['COMPLETED', 'FAILED', 'STOPPED', 'TIMED_OUT']:
        break
    
    time.sleep(30)  # Check every 30 seconds

# Download results if completed
if status == 'COMPLETED':
    files = client.download_output(unique_job_name, "./results")
    print(f"Downloaded {len(files)} files")

Error Handling

A request the API refuses raises requests.HTTPError. Its message is the status code and the API's reason.

import requests

try:
    result = client.submit_job(
        directory="./code",
        job_name="my-job",
        instance_type="ml.g4dn.xlarge"
    )
except requests.HTTPError as e:
    print(e)  # the status code and the API's reason
    print(e.response.status_code)
except ValueError as e:
    print(f"Configuration error: {e}")
except FileNotFoundError as e:
    print(f"File not found: {e}")

License

MIT License