Metadata-Version: 2.4
Name: afrilink-sdk
Version: 0.9.4
Summary: AfriLink SDK — Train, finetune, and generate on GPUs from your notebook
Home-page: https://github.com/dataspires/afrilink-sdk
Author: DataSpires
Author-email: DataSpires <info@dataspires.com>
License: MIT
Project-URL: Homepage, https://dataspires.com
Project-URL: Documentation, https://www.dataspires.com/docs
Project-URL: Repository, https://github.com/DataSpires/afrilink-sdk
Project-URL: Bug Tracker, https://github.com/DataSpires/afrilink-sdk/issues
Keywords: hpc,high-performance-computing,finetuning,pretrain,generate,vlm,llm,lora,notebook,gpu,k3s,slurm,afrilink
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Intended Audience :: Developers
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.8
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: System :: Distributed Computing
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: requests>=2.28.0
Provides-Extra: full
Requires-Dist: psutil>=5.9.0; extra == "full"
Provides-Extra: build
Requires-Dist: cryptography>=41.0.0; extra == "build"
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: black>=23.0.0; extra == "dev"
Requires-Dist: mypy>=1.0.0; extra == "dev"
Dynamic: author
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-python

# AfriLink SDK

One-line access to GPUs for training, finetuning, and inference — from any notebook.

```bash
pip install afrilink-sdk
```

**Works in:** Google Colab · Kaggle · Jupyter · VS Code · any Python 3.8+ environment

---

## Contents

1. [60-second quickstart](#60-second-quickstart)
2. [Which method should I use?](#which-method-should-i-use)
3. [Backends: k3s vs. opentoken-a100](#backends)
4. [Authentication](#authentication)
5. [Guides by task](#guides-by-task)
   — [Pretrain](#pretrain) · [Finetune](#finetune) · [Generate](#generate) · [Custom containers](#custom-containers)
6. [Working with your output model](#working-with-your-model)
7. [API reference](#api-reference)
8. [Hardware & billing](#hardware--billing)
9. [Model & dataset registry](#model--dataset-registry)
10. [Troubleshooting](#troubleshooting)
11. [Built-in help](#built-in-help)
12. [Architecture (advanced)](#architecture-advanced)

---

## 60-second quickstart

This example uses the **default `k3s` backend** — no environment variable needed.

```python
from afrilink import AfriLinkClient

client = AfriLinkClient()
client.authenticate(api_key="afk_live_…")  # or set AFRILINK_API_KEY

job = client.pretrain(
    kind="image-classify",
    model="resnet18",
    data="./images/",
    params={"epochs": 5, "num_classes": 2},
    gpus=1,
)
result = job.run(wait=True)
print(result["status"])
```

Every other capability (finetune, generate, custom containers) follows this same `client → job → job.run()` pattern. See [Guides by task](#guides-by-task).

---

## Which method should I use?

| Your goal                                         | Call                        | You provide                                       |
| -------------------------------------------------- | ---------------------------- | -------------------------------------------------- |
| Train a **YOLO / CNN / transformer from scratch**   | `client.pretrain(kind=…)`    | `kind`, `data`, optional `weights=` / `scratch=`    |
| Fine-tune an **LLM / VLM** foundation model         | `client.finetune()`          | `model`, `training_mode`, `data`, optional `task=`  |
| Run a model **without** changing weights            | `client.generate()`          | `model`, prompts / image+question manifest         |
| Need a **custom framework/version** not curated     | `client.build_and_train()`   | image spec + script                                 |

> **Rule of thumb:** train from zero → `pretrain()`. Adapt existing weights → `finetune()`. Frozen answers → `generate()`. Need a bespoke environment → `build_and_train()`.

---

## Backends

AfriLink can run your job on two different backends. This choice affects auth, output retrieval, and billing, so it's worth understanding up front rather than discovering it mid-job.

| | `k3s` **(default)** | `opentoken-a100` |
| --- | --- | --- |
| **How to select** | Default — nothing to set | `AFRILINK_BACKEND=opentoken-a100` |
| **What it targets** | DataSpires K3s GPU cluster (multi-node, multi-GPU capable) | Single dedicated OpenToken node — gives you access to a NVIDIA Tesla V100 32GB, over SSH |
| **Job model** | Async by default — submit, then poll or list jobs, even from a new session | Synchronous container run on a single shared node |
| **File I/O** | S3-compatible presigned URLs | SCP upload; runs from `/workspace/job/...` |
| **Output retrieval** | Must sync from `/workspace/job/output/` — see [K3s output retrieval](#k3s-output-retrieval) | `client.download_model()` pulls the whole `output/` dir directly |
| **Billing** | Not yet wired (`billing: None` in result) | `$0.60/GPU-hour`, billed per completed GPU-minute |
| **Multi-GPU** | Supported by the cluster | Clamped to 1 (single-GPU node) |

If you don't set `AFRILINK_BACKEND`, you are on `k3s`. Every quickstart below states which backend it's using.

---

## Authentication

As of v0.8.x, auth is a single stateless API key — no passwords, no certificate refresh, no SSH key management on your side.

### Get a key
1. Sign up at [dataspires.com](https://dataspires.com).
2. **Profile → AfriLink SDK keys → Create new key.** Copy the `afk_live_…` value — it's shown once.
3. Store it as `AFRILINK_API_KEY`.

### Set the key

| Environment | How |
| --- | --- |
| Google Colab | 🔑 sidebar → **Add secret** → name `AFRILINK_API_KEY` → enable for notebook |
| Kaggle | Add-ons → **Secrets** → name `AFRILINK_API_KEY` → attach to notebook |
| Local Jupyter / VS Code | `os.environ["AFRILINK_API_KEY"] = "afk_live_…"` before `authenticate()` |
| Anywhere | `client.authenticate(api_key="afk_live_…")` |

```python
client = AfriLinkClient()
client.authenticate()   # resolves from secret / env / argument, in that order
```

### What happens at auth time (~1–2 sec total)

| Phase | What runs |
| --- | --- |
| 1. DataSpires session | Exchange API key at `api.dataspires.com` for a short-lived Supabase JWT (used for billing writes) |
| 2. Node reachability *(opentoken-a100 only)* | Silent SSH probe confirming your slot is live |

The JWT lives in memory for the kernel lifetime — nothing is written to disk. Rotate a key by revoking it on the dashboard and minting a new one.

---

## Guides by task

### Pretrain
*(train from scratch — YOLO, CNN, or transformer)*

**On the default `k3s` backend:**
```python
job = client.pretrain(
    kind="image-classify",
    model="resnet18",
    data="./images/",
    params={"epochs": 5, "num_classes": 2},
    gpus=1,
)
result = job.run(wait=True)
```

**YOLO example, explicitly on `opentoken-a100`:**
```python
import os
os.environ["AFRILINK_BACKEND"] = "opentoken-a100"

job = client.pretrain(
    kind="yolo-detect",
    weights="yolo11n.pt",       # or scratch=True for random init
    data="./dataset/",
    data_config="data.yaml",
    gpus=1,
    time_limit="02:00:00",
)
result = job.run(wait=True)
client.download_model(result["job_id"], "./yolo-out")
```

Curated containers resolved from `kind=`:

| `kind` prefix | Container | Frameworks |
| --- | --- | --- |
| `yolo-*` | `afrilink-yolo` | Ultralytics, PyTorch, torchvision |
| `image-classify` | `afrilink-vision` | PyTorch, torchvision |
| transformer/ViT | `afrilink-pretrain` | Transformers, accelerate |

Need a stack outside these? → [Custom containers](#custom-containers).

#### K3s output retrieval
On `k3s`, your script writes to `/workspace/job/output/`; the orchestrator syncs it to S3 on completion.

```python
from pathlib import Path
job.k8s_runner.download_output(job.job_id, Path("./my-outputs"))
```

Requires a deployed orchestrator exposing `GET /api/v1/jobs`. Full spec: [dataspires.com/docs](https://www.dataspires.com/docs).

### Finetune
*(LoRA/QLoRA adaptation of an LLM/VLM — runs in the `afrilink-finetune` container)*

```python
import pandas as pd

data = pd.DataFrame({"text": ["Below is an instruction...\n\n### Response:\n..."]})

job = client.finetune(
    model="qwen2.5-0.5b",
    training_mode="low",     # low | medium | high — see table below
    data=data,
    gpus=1,
    time_limit="01:00:00",
)
result = job.run(wait=True)
if result["status"] == "completed":
    client.download_model(result["job_id"], "./my-model")
```

| Mode | Strategy | Quantization |
| --- | --- | --- |
| `low` | QLoRA (rank 8) | 4-bit |
| `medium` | LoRA (rank 16) | 8-bit / none |
| `high` | Full LoRA (rank 64) | none |

### Generate
*(frozen inference — no weight updates)*

```python
job = client.generate(
    model="smolvlm-256m",
    data="./manifest_dir/",
    outputs=["text", "confidence"],   # add "hidden_states" for analysis
    max_new_tokens=64,
    temperature=0.0,
)
result = job.run(wait=True)
```

### Custom containers

For frameworks or model versions the curated containers don't cover. **`build_and_train()` requires `AFRILINK_BACKEND=opentoken-a100`** (not available on the default `k3s` backend).


```python
spec = dict(
    base_image="pytorch",       # preset — see table below
    pip_packages=["transformers>=4.45", "accelerate>=0.34", "peft>=0.13"],
    apt_packages=["git"],
    model_source={"kind": "huggingface", "id": "Qwen/Qwen2.5-0.5B-Instruct"},
)

# check cache first — avoids a ~5 min rebuild for a spec you've already built
hit = client.find_existing_image(**spec)

result = client.build_and_train(
    **spec,
    script="my_train.py",
    gpus=1,
    time_limit_hours=0.5,
    reuse_existing_image=True,   # default
)
client.download_model(result["run"]["job_id"], "./output")
```

**`base_image` presets:**

| Preset | Resolves to | Notes |
| --- | --- | --- |
| `pytorch` | `pytorch/pytorch:2.5.0-cuda12.4-cudnn9-runtime` | GPU default |
| `pytorch-2.4` | `pytorch/pytorch:2.4.0-cuda12.4-cudnn9-runtime` | |
| `pytorch-cpu` | `pytorch/pytorch:2.5.0-cpu-runtime` | CPU-only, smaller |
| `cuda-12.4` | `nvidia/cuda:12.4.0-runtime-ubuntu22.04` | bring-your-own-Python |
| `ultralytics` | `ultralytics/ultralytics:latest` | YOLOv8 ready |

**`model_source` kinds:** `huggingface` (`id`, optional `revision`/`subfolder`), `url`, `git`, `gs`, `s3`, or omit to load the model yourself in-script. Models are fetched at **runtime**, not baked into the image — keeps images ~2GB instead of 7+GB. Gated HF models (Llama, Gemma) work automatically if you add `HUGGINGFACE_TOKEN` as a notebook secret.

**Cache key** includes `base_image`, `pip_packages`, `apt_packages`, index URLs, and `model_source`. It excludes `script`, `env`, `extra_files`, and job/user IDs — so two runs with an identical environment spec but different training scripts share a cached image.

**Lifecycle:** the built image lives permanently in Artifact Registry (that's what cache hits read from); the local copy on the compute node is deleted at the end of each `build_and_train()` call unless `cleanup_image_after=False`.

---

## Working with your model

### Convert to GGUF for Ollama / llama.cpp
```python
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B")
model = PeftModel.from_pretrained(base, "./my-model")
merged = model.merge_and_unload()
merged.save_pretrained("./my-model-merged")
AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B").save_pretrained("./my-model-merged")

# python convert_hf_to_gguf.py ./my-model-merged --outfile my-model.gguf --outtype f16
# ./llama-quantize my-model.gguf my-model-q4.gguf Q4_K_M
# ollama create my-model -f Modelfile && ollama run my-model
```

### Publish to Hugging Face Hub
```python
from huggingface_hub import HfApi

api = HfApi(token="hf_...")
api.create_repo("your-username/my-finetuned-model", exist_ok=True)
api.upload_folder(folder_path="./my-model", repo_id="your-username/my-finetuned-model")  # adapter only
```

---

## API reference

### `AfriLinkClient`

| Method | Description |
| --- | --- |
| `authenticate(api_key=None)` | Resolve key (arg / env / Colab / Kaggle secrets), exchange for session |
| `pretrain(kind, data, gpus, ...)` | From-scratch training job |
| `finetune(model, training_mode, data, gpus, ...)` | LoRA/QLoRA job in `afrilink-finetune` |
| `generate(model, data, outputs, ...)` | Frozen inference job |
| `find_existing_image(...)` | Check cache before a custom build |
| `build_image(...)` | Build a custom image only |
| `build_and_train(...)` | Cache-check → build (or skip) → run → cleanup |
| `delete_built_image(job_id_or_image)` | Remove image from the compute node's local cache |
| `download_model(job_id, local_dir)` | Download `output/` directory |
| `upload_dataset(local_path, dataset_name)` | Upload to job-scoped staging |
| `list_containers()` / `list_available_models()` / `list_available_datasets()` | Registry lookups |
| `get_model_requirements(model, training_mode)` | GPU/memory recommendation |
| `cancel_job(job_id)` | Stop and remove a job |
| `list_jobs(include_completed=False)` | List your jobs |
| `get_job_status(job_id)` | Poll by ID — works across sessions on `k3s` |
| `get_job(job_id)` | In-session job object, if still tracked |
| `run_command(command)` | Arbitrary shell command *(opentoken-a100 only)* |

### `TrainJob` / `FinetuneJob`

Returned by `pretrain()` / `finetune()` / `build_and_train()`.

| Member | Description |
| --- | --- |
| `run(wait=True, poll_interval=30)` | Submit; `wait=True` polls to completion |
| `cancel()` | Stop the job |
| `get_logs(tail=100)` | Recent log lines |
| `estimated_cost_usd()` | Estimate from GPUs × time limit |
| `status`, `job_id` | Current state / 8-char job ID |
| `container_id` | *(opentoken-a100 only, set after `run()`)* |
| `k8s_runner` | *(k3s only)* — use for `download_output()` |

**Result shape — `opentoken-a100`:**
```json
{
  "job_id": "a1b2c3d4",
  "status": "completed",
  "output_dir": "/mnt/data/sdk-jobs/a1b2c3d4/output",
  "billing": {"total_gpu_minutes": 5.0, "total_cost_usd": 0.1667, "rate_per_gpu_hour": 0.60}
}
```

**Result shape — `k3s`:**
```json
{"job_id": "a1b2c3d4", "status": "completed", "billing": null}
```
On failure, also includes `"error": "<last 200 log lines>"`.

**Async job discovery (k3s):**
```python
result = job.run(wait=False)
for entry in client.list_jobs():
    print(entry["job_id"], entry["state"], entry.get("reason"))

status = client.get_job_status(result["job_id"])  # works even after a notebook restart
```
`list_jobs(include_completed=True)` includes completed jobs for up to 3 days.

---

## Hardware & billing

**`opentoken-a100` node** — this backend name refers to the OpenToken node; the GPU it actually gives you access to is a 1× NVIDIA Tesla V100 32GB · 12 CPU cores · 84GB RAM · 774GB storage · CUDA 13.0 pre-installed, containerized via `--gpus all`.

| Model size | Mode | Fits on 1 GPU? |
| --- | --- | --- |
| 0.5B – 3B | low / medium | yes |
| 3B – 7B | low (QLoRA 4-bit) | yes |
| 7B | medium (LoRA 8-bit) | yes |
| 13B | low (QLoRA 4-bit) | tight |
| 30B+ | low (QLoRA 4-bit) | unlikely (OOM risk) |

**Billing:** `$0.60`/GPU-hour on `opentoken-a100`, billed per completed GPU-minute (1-minute minimum), deducted automatically from your DataSpires balance. Invoices: [dataspires.com/dashboard/billing](https://dataspires.com/dashboard/billing). Cloud Build time for custom images is absorbed by the platform — you only pay for GPU time. `k3s` billing is not yet wired (tracked on the roadmap).

---

## Model & dataset registry

```python
client.list_available_models()                # all
client.list_available_models(size="tiny")      # tiny | small | medium | large
client.list_available_datasets()
client.get_model_requirements("qwen2.5-0.5b", "low")
```

| ID | Type | Params | Min VRAM |
| --- | --- | --- | --- |
| `qwen2.5-0.5b` | text | 0.5B | 4 GB |
| `gemma-3-270m` | text | 0.27B | 2 GB |
| `llama-3.2-1b` | text | 1.0B | 4 GB |
| `deepseek-r1-1.5b` | text | 1.5B | 6 GB |
| `ministral-3b` | text | 3.3B | 8 GB |
| `florence-2-base` | vision | 0.23B | 4 GB |
| `smolvlm-256m` | vision | 0.26B | 2 GB |
| `moondream2` | vision | 1.9B | 8 GB |
| `internvl2-1b` | vision | 1.0B | 4 GB |
| `llava-1.5-7b` | vision | 7.0B | 16 GB |

Anything else → [custom containers](#custom-containers) with `model_source=`.

---

## Troubleshooting

| Symptom | Likely cause | Fix |
| --- | --- | --- |
| `authenticate()` hangs or times out | Node reachability probe failing on `opentoken-a100` | Retry; if persistent, check [status page] or switch to `k3s` |
| Job stuck in `queued` | Cluster at capacity (k3s) | `client.list_jobs()` to check state; jobs auto-clear after 3 days if abandoned |
| `gpus=2` silently becomes `gpus=1` | `opentoken-a100` is single-GPU | Use `k3s` for multi-GPU jobs |
| Fresh build every time despite unchanged deps | A field outside the cache key changed (e.g. `script`) — this is expected | Only `base_image`, `pip_packages`, `apt_packages`, index URLs, and `model_source` are hashed |
| Gated HF model fails to download | Missing `HUGGINGFACE_TOKEN` | Add it as a notebook secret |
| `download_model()` returns nothing on k3s | Script didn't write to `/workspace/job/output/` | Confirm output path in your training script |

---

## Built-in help

Query the inline reference manual from any cell — no internet required:

```python
import afrilink

afrilink.docs("help")       # recommended
afrilink.docs("quickstart")
afrilink.docs("auth")
afrilink.docs("finetune")
afrilink.docs("pretrain")
afrilink.docs("generate")
afrilink.docs("specs")
afrilink.docs("datasets")
afrilink.docs("billing")

# Quoted slash form also works:
afrilink / "help"
```

---

## Architecture (advanced)

High-level flow: your notebook authenticates against `api.dataspires.com` (a Cloudflare Worker), which exchanges your API key for a short-lived Supabase JWT, a compute-node SSH key, a GCP service-account key (for custom builds), and a registry pull token. Job execution then proceeds either over SSH to the `opentoken-a100` node (Docker, running on the V100) or via the K3s orchestrator REST API (Kubernetes), with custom images built out-of-band on Cloud Build and pushed to a private Artifact Registry.

```
Notebook (Colab / Kaggle / Local)         api.dataspires.com (Cloudflare Worker)
+---------------------+                   +---------------------------+
| AfriLink SDK        | --- POST -----→   | exchange afk_live_… for:  |
|  client.authenticate()                  |  - Supabase JWT (billing) |
|                     | ←-- response ---  |  - V100 SSH key (in-mem)  |
+---------------------+                   |  - GCP SA key (build)     |
     |        ↓                           |  - GHCR PAT (image pulls) |
     |   (in-memory state)                +---------------------------+
     |
     ↓
+---------------------+        SSH        +---------------------+
| docker_runner.py    | ----------------→ | OpenToken V100 32GB |
|  - prepare_job_dir  |   /mnt/data/      |  Docker daemon      |
|  - upload via SCP   |    sdk-jobs/      |  (containerd at     |
|  - docker run --gpus=all                |   /mnt/data/)       |
|  - docker inspect (poll)                +---------------------+
+---------------------+
     |
     ↓ build path
+---------------------+   Cloud Build    +---------------------+
| build.py            | --→ submit job → | europe-west4-       |
|  - generate Docker- |     (anadrome)   | docker.pkg.dev/...  |
|    file from spec   |                  |  afrilink-user-     |
|  - tar build context|                  |  images/<user>/<job>|
|  - upload to GCS    |                  +---------------------+
+---------------------+                            |
                                                   ↓ docker pull
                                              (V100 fetches image,
                                               runs it, deletes
                                               local layer at end)
```

Full diagram and internals: [dataspires.com/docs](https://www.dataspires.com/docs).

---

## Links

- **Docs:** [dataspires.com/docs](https://www.dataspires.com/docs)
- **Repository / Issues:** [github.com/DataSpires/afrilink-sdk](https://github.com/DataSpires/afrilink-sdk)
- **License:** MIT
