Metadata-Version: 2.4
Name: kaio
Version: 0.3.1
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
Requires-Dist: packaging>=21.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://www.kaion5.com/home/kaio-platform.html).

Today Kaio runs your data processing and training jobs. We're building it into a complete MLOps workflow, from first experiment to production system.

## Installation

```bash
pip install kaio
```

## Quick Start

```python
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 and Models: we suggest the Hugging Face Hub

Your code is uploaded with every job; your data shouldn't have to be. We suggest keeping your datasets and model weights on the [Hugging Face Hub](https://huggingface.co) and loading them in your code. The job downloads them itself, so nothing is uploaded from your machine, and the same data is there for every run. (If your data can't go on the Hub, see [Kaio Datasets](#kaio-datasets-when-your-data-cant-go-on-the-hub).)

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

# Pin a commit or tag, so a rerun reads exactly the same data
dataset = load_dataset("your-name/your-dataset", revision="<commit>")
model = AutoModelForSequenceClassification.from_pretrained("your-name/your-model", revision="<commit>")
```

```python
client.submit_job(directory="./my_code", entrypoint="train.py")
```

- **Private datasets and models:** put your Hugging Face token in the directory's `.env` as `HF_TOKEN=hf_...`, or pass `env={"HF_TOKEN": "hf_..."}`. Hugging Face libraries read it automatically. Your token is handled securely:
  - sent over HTTPS, separately from your code, and never uploaded with it;
  - kept encrypted only while the job runs, and readable only by your own job;
  - masked as `****` if your code prints it, and deleted as soon as the job ends.

  See [Environment Variables and Tokens](#environment-variables-and-tokens).
- **Saving results to the Hub:** your code can push with the same token, for example `model.push_to_hub("your-name/your-model")`. The token needs write access.
- **Large uploads:** `submit_job` warns when the upload is over 100 MB or holds more than 10 MB of data files or model weights, and suggests loading them from the Hub.

### Kaio Datasets: when your data can't go on the Hub

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

```python
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 if you keep your data there (and list `datasets/` in `.kaioignore`).

- **Kept until you delete them**, unlike job files, which go after 7 days. Free, up to 20 GB per account (packed size).
- **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.

### Leaving Files Out: `.kaioignore`

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

```
# loaded from the Hub instead
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.

## 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`
- **Data from the Hugging Face Hub**: jobs load datasets and models themselves, so only your code is uploaded
- **Kaio Datasets**: upload data once and reuse it across jobs, when it can't go on the 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
- **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

#### `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](#kaio-datasets-when-your-data-cant-go-on-the-hub).
- `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](#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 API calls fail
- `FileNotFoundError`: If directory or entrypoint doesn't exist
- `ValueError`: If code package exceeds volume capacity, or an environment variable isn't allowed

#### `get_jobs()`

Get all jobs for the authenticated user.

**Returns:**
- `dict`: List of jobs with count

**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` 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 7 days. The job stays in your job list with its cost, 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 and cost are 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, up to 20 GB per account. See [Kaio Datasets](#kaio-datasets-when-your-data-cant-go-on-the-hub).

```python
client.upload_dataset("sensor-readings", "./data/readings")
client.list_datasets()   # {"datasets": [{"name", "size_bytes", "uploaded_at"}], "used_bytes", "quota_bytes"}
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:**
```python
files = client.download_output("unique-job-name-123", "./results")
for file_path in files:
    print(f"Downloaded: {file_path}")
```

## Supported Instance Types

Jobs run in AWS af-south-1 (Cape Town). Instance types not offered there are rejected at submission.

### CPU Instances
- `ml.t3.*`, `ml.m5.*`, `ml.m6i.*`, `ml.c5.*`, `ml.c6i.*`, `ml.r5.*`, `ml.r5d.*`

### GPU Instances
- `ml.g4dn.xlarge`, `ml.g4dn.2xlarge`, `ml.g4dn.4xlarge`, `ml.g4dn.8xlarge`, `ml.g4dn.12xlarge` (4 GPUs), `ml.g4dn.16xlarge`

## 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`. You pay only for the time AWS bills: its spot discount (up to 90% off the on-demand price) comes straight to you, and time spent waiting after an interruption isn't charged.

```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(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. Don't write values into your output files; those are kept for 7 days like any output.
- **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), if you need it specifically

## 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. Load data from the Hugging Face Hub instead of uploading it (see [Data and Models](#data-and-models-we-suggest-the-hugging-face-hub)).

**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 (keep large datasets out of the code directory)

### 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
**Important**: Include a `requirements.txt` file in your code directory listing all Python packages your job needs. This file will be automatically installed within the job environment.

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

Processing jobs suit data processing, feature engineering, model evaluation, and simple model training. When using `instance_count > 1`, you can parallelize processing of large datasets across multiple instances.

**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](#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

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

## 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(
    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

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

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

```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']}")
```

### 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(
    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

```python
import requests

try:
    result = client.submit_job(
        directory="./code",
        job_name="my-job",
        instance_type="ml.g4dn.xlarge"
    )
except requests.HTTPError as e:
    if e.response.status_code == 403:
        print("Access denied - check your API key")
    else:
        print(f"API error: {e}")
except ValueError as e:
    print(f"Configuration error: {e}")
except FileNotFoundError as e:
    print(f"File not found: {e}")
```

## Job Dashboard

**Access your jobs at**: https://www.kaion5.com/home/kaio-platform.html

After signing in to the Kaio Platform, you can:
- Monitor job status and progress
- View job logs and details
- Download job outputs and results
- Manage your job history

## License

MIT License
