Metadata-Version: 2.4
Name: kaio
Version: 0.5.2
Summary: MLOps from your IDE: run Python code on cloud CPUs and GPUs
Author-email: Kaio Compute <danielb@kaion5.com>
License-Expression: MIT
Project-URL: Homepage, https://www.kaion5.com/home/kaio-platform.html
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests>=2.25.0
Dynamic: license-file

# 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](https://kaion5.com).

## Quick Start

```bash
pip install kaio
```

```python
import time
from kaio import Client

client = Client("https://api.kaion5.com")
client.login("your-api-key")  # from the dashboard's API Keys page

# Run train.py from its job folder on a GPU
job = client.submit_job(
    job_dir="./my_code",  # train.py and requirements.txt at its top
    entrypoint="train.py",
    instance_type="ml.g4dn.xlarge",
    job_name="training-job",
)
name = job["unique_job_name"]

# Wait for it to end, then download what it wrote and its logs
while (status := client.get_job(name)["status"]) not in ("COMPLETED", "FAILED", "STOPPED", "TIMED_OUT"):
    time.sleep(30)
if status == "COMPLETED":
    client.download_output(name, "./results")
```

Or from a terminal:

```bash
kaio login                                               # asks for your API key once
kaio run ./my_code -e train.py -i ml.g4dn.xlarge --wait  # prints the job's name, then its status as it changes
kaio download <job name> ./results
```

## Data

A job's folder, `job_dir`, 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:

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

client.submit_job(job_dir="./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 14 days.
- **Several per job:** `input_data=["a", "b"]`, up to 5. They must fit on the job's disk, 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.

#### From a Job's Outputs

A job's outputs can become a dataset too, so the next job takes them as its input. One processing job prepares the data, a training job learns from it:

```python
prep = client.submit_job(job_dir="./prep", entrypoint="prep.py", workflow="processing")
client.create_dataset_from_job("features", prep["unique_job_name"])   # waits for prep, then keeps what it wrote

client.submit_job(job_dir="./train", entrypoint="train.py", workflow="training", input_data="features")
```

- **What's in it:** the files the job wrote, as `download_output()` would give you, without Kaio's logs. `path="features"` takes one folder of them instead, its contents at the dataset's root.
- **Waiting:** by default it waits for the job to finish and the dataset to be ready, so the next line can use it. `wait=False` returns at once; the job must have completed, and `list_datasets()` shows the dataset under `making` until it's ready.
- **Within 14 days:** a job's own outputs are deleted 14 days after it ends; the dataset made from them is kept until you delete it.
- **Training jobs too:** a training job's model comes as its `model.tar.gz`.
- **From a terminal:** `kaio datasets from-job features <unique_job_name>`, with `--path` and `--no-wait`.

### 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](https://huggingface.co):

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

```python
# 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 job folder'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](#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 job folder lists what not to upload, one pattern per line, like `.gitignore`:

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

A pattern without a `/` matches a name at any depth, one with a `/` matches from the job folder, and a trailing `/` matches folders only. Lines starting with `#` are comments. Unlike `.gitignore`, `!` patterns aren't supported. `.env` files, `.git`, virtual environments (`.venv/`, `venv/`), `__pycache__`, `.ipynb_checkpoints`, `*.pem` and `*.key` 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; your code runs on [Kaio's runtime](#kaios-runtime), PyTorch or TensorFlow on Python 3.11–3.13
- **Several instances at once**: split work across instances with `instance_count`
- **Kaio Datasets**: upload data once and use it in any job, or keep a job's outputs as a dataset for the next 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:**
```python
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](#projects).

#### `submit_job(**kwargs)`

Run your code as a processing or training job. The SDK packages and uploads the job folder, and starts the job on [Kaio's runtime](#kaios-runtime).

**Parameters:**
- `job_dir` (str): The job's folder: your code, its `requirements.txt` and the entrypoint at its top. The job runs inside it, and the files it adds or changes there are the job's outputs. Everything in it is uploaded with each job, except `.env` files and what `.kaioignore` leaves out. Defaults to the current directory. `directory`, its old name, still works with a `DeprecationWarning`.
- `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): Extra disk per instance in GB, added to the about 100 GB every job has. Defaults to 5, at most 50.
- `entrypoint` (str): The script or notebook to run (.py or .ipynb), a file at the top of `job_dir`. 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](#kaio-datasets).
- `framework` (str, optional): `"pytorch"` (the default) or `"tensorflow"`. See [Kaio's Runtime](#kaios-runtime).
- `python_version` (str, optional): `"3.11"`, `"3.12"` (the default) or `"3.13"`.
- `framework_version` (str, optional): Ignored. Pin another version in `requirements.txt`.
- `image_uri` (str, optional): [Your own container](#your-own-container).
- `workflow` (str): Workflow type - "processing" or "training". Defaults to "processing".
- `training_args` (dict, optional): [Spot instances](#spot-instances), for any job, and distributed training.
- `max_runtime_hours` (float): Hours before the job is stopped. Defaults to 1, at most 48.
- `env` (dict, optional): Environment variables for your code, such as tokens and API keys. See [Environment Variables and Tokens](#environment-variables-and-tokens).
- `env_file` (str, optional): A `.env` file with more variables. Defaults to `job_dir`'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 the entrypoint isn't a file at the top of `job_dir`
- `ValueError`: If the code package is over 5 GB, 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:**
```python
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:**
```python
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 14 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:**
```python
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](#kaio-datasets).

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

#### `create_dataset_from_job(name, job, path=None, wait=True)`

Keep a job's outputs as a dataset for later jobs to use with `input_data=name`. With `wait=True` it waits for the job to finish, then returns the dataset once it's ready; it raises `RuntimeError` if the job didn't complete. `path` takes one folder of the outputs. See [From a Job's Outputs](#from-a-jobs-outputs).

```python
client.create_dataset_from_job("features", "20261014-101500-123-abc", path="features")
```

#### `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:**
```python
files = client.download_output("unique-job-name-123", "./results")
for file_path in files:
    print(f"Downloaded: {file_path}")
```

## Training Arguments

`training_args` sets up spot instances, for any job, and distributed training.

### Spot Instances
Any job, processing or training, can run on spot instances: AWS's spare capacity, which costs less, 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`.

```python
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:

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

# Or from a .env file. The job folder's own .env is used automatically.
client.submit_job(job_dir="./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 14 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.

## Kaio's Runtime

Your code runs on Kaio's runtime, on CPUs and GPUs:

- **PyTorch** 2.14 (the default), or **TensorFlow** 2.21 with `framework="tensorflow"`
- **Python** 3.12, or `python_version="3.11"` or `"3.13"`
- Anything else goes in your `requirements.txt`
- Distributed training runs PyTorch 2.9 on Python 3.12

### Your Own Container

`image_uri` runs a public container image with `python3` instead, such as `python:3.13-slim`.

## Code Requirements

### Data and Code Structure
The job folder you submit with `client.submit_job()` is packaged and uploaded with every job, except `.env` files and what `.kaioignore` lists. See [Data](#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 14 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 14 days.

### File Size Limits
- **Code package**: up to 5 GB per job. Put large data in a [Kaio dataset](#data) or on the Hugging Face Hub
- **Disk**: every job has about 100 GB, and at least 50 GB on any instance; `volume_size_gb` adds up to 50 GB more

### Multi-Instance Code Structure
For multi-instance processing jobs, your code can use these environment variables to split the work:

```python
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 job folder 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:

- **Checkpoints**: files your code writes to `/opt/ml/checkpoints` are saved while the job runs, so it can resume after an interruption
- **The model**: `/opt/ml/model` is saved as `model.tar.gz`

### Multi-Instance Processing Example

```python
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)
```

### Passing Outputs to the Next Job

Workflows chain through datasets: keep one job's outputs as a dataset, and give it to the next job as `input_data`. A processing job that cleans and splits data can feed a training job, or another processing job, without the data leaving Kaio:

```python
prep = client.submit_job(job_dir="./prep", entrypoint="build_features.py", workflow="processing",
                         instance_type="ml.m5.xlarge")
client.create_dataset_from_job("features", prep["unique_job_name"])

train = client.submit_job(job_dir="./train", entrypoint="train.py", workflow="training",
                          instance_type="ml.g4dn.xlarge", input_data="features")   # in datasets/features/
```

`build_features.py` writes its results beside itself, as any job's outputs; `train.py` reads them from `datasets/features/`. See [From a Job's Outputs](#from-a-jobs-outputs).

## Examples

### Processing Workflow Examples

```python
from kaio import Client

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

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

# Multi-instance parallel processing
result = client.submit_job(
    job_dir="./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

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

# Training on spot instances
result = client.submit_job(
    job_dir="./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

```python
# Kaio's TensorFlow runtime on Python 3.11
result = client.submit_job(
    job_dir="./tensorflow_training",
    job_name="tf-training",
    workflow="training",
    instance_type="ml.g4dn.xlarge",
    framework="tensorflow",
    python_version="3.11",
    entrypoint="train.py"
)

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

## Job Management

### Getting All Jobs

```python
# 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

```python
# 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.

```python
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

```python
# 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

```python
# Submit training job
result = client.submit_job(
    job_dir="./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.

```python
import requests

try:
    result = client.submit_job(
        job_dir="./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
