Training Pipeline
Training rents a GPU instance, drives it entirely over SSH, pulls the resulting LoRA adapter back, and grades it. training/ implements it as eight idempotent Celery stages on the training queue, with a manifest that records provider state between stages.
The chain
smart.finalize creates the train job and calls enqueue_training_job, which builds a Celery chain and applies it to the training queue.
| Order | Task | Run status on success | Does |
|---|---|---|---|
| 1 | training.request_instance | training_instance_requested | Rents a pod/instance from RunPod, or Vast.ai on fallback. |
| 2 | training.wait_instance_ready | training_instance_ready | Polls until SSH answers a READY probe. |
| 3 | training.upload_payload | training_payload_uploaded | Uploads dataset, manifest, metadata, and train_unsloth.py. |
| 4 | training.start_remote | training_started | Launches training detached and captures its PID. |
| 5 | training.monitor | training_monitoring | Polls the PID, then verifies the expected outputs exist. |
| 6 | training.sync_artifacts | training_artifacts_synced | Downloads the adapter bundle and publishes the avatar profile. |
| 7 | training.terminate_instance | — (training_instance_terminated in the manifest) | Releases the GPU. |
| 8 | training.complete | completed | Writes the quality gate, evaluates the adapter, closes the job. |
training.process is a ninth task that re-dispatches an interrupted chain rather than running a stage itself.
Idempotence and resume
The manifest's pipeline_status is rank-ordered:
queued 0 · instance_requested 10 · instance_ready 20 · payload_uploaded 30
started 40 · monitoring 50 · artifacts_synced 60 · instance_terminated 70 · completed 80
Every stage first compares its own rank against the recorded one and returns immediately if the work is already done. Run-status patches follow the same rule — _patch_run_status never moves a run backwards, and refuses to touch a terminal run at all.
dispatch_training_chain uses this to resume: it reads the current pipeline_status, computes the remaining stages, and dispatches only those. An empty remainder means the chain finishes with run_stage_complete.
Two stages additionally enforce ordering: terminate refuses to run before training_artifacts_synced, and complete calls terminate itself if it has not happened yet — then refuses to proceed if the instance still is not terminated.
Preconditions
start_training runs three checks before creating the job:
- The run status must be
review_approvedorpostprocess_running. - A training provider must be configured — see Provider configuration.
artifact_manifest.final_dataset.examplesmust be non-empty,review.acceptedmust be true, andcreator_id,avatar_id, anddomain_idmust all be present.
If manifest construction or the run patch then fails, the just-created job is retired before the exception propagates, so a failed start does not leave a claimable job behind.
GPU providers
RunPod is primary; Vast.ai is a fallback. _request_instance tries RunPod when configured, and falls through to Vast.ai only if the RunPod error looks like a capacity problem — the message matches one of unavailable, no pod(s), no instances, insufficient, not enough, not available, capacity, gpu type, gpu_type, gpu not found, stock. Any other RunPod error is raised immediately rather than silently paying for a second provider.
Both clients pass four environment variables into the instance: TRAINING_JOB_ID, TRAINING_RUN_ID, TRAINING_AVATAR_ID, TRAINING_DOMAIN_ID. Instances are labelled onboarding-v4-{job_id}.
RunPod
podFindAndDeployOnDemand is called over GraphQL, walking a deduplicated GPU candidate list built from RUNPOD_GPU_TYPE_ID, then RUNPOD_GPU_TYPE_IDS, then a built-in list (H100 80GB HBM3, H100 PCIe, H100 NVL, H100 SXM, RTX PRO 6000 Blackwell, RTX 6000 Ada, RTX A6000). A capacity error moves to the next candidate; anything else aborts.
Readiness polls pod.runtime.ports for the public mapping of private port 22, falling back to the REST portMappings["22"] plus publicIp. Once a host and port appear, the client checks the TCP socket and then runs echo READY over SSH; only a real READY counts.
Vast.ai
Offers are searched with POST /bundles/, filtered to verified, rentable, un-rented machines with at least one direct port and the requested GPU count, ordered by hourly price. The GPU list is VAST_GPU_NAME plus VAST_GPU_FALLBACK_NAMES. For each offer, PUT /asks/{offer_id}/ creates an ssh_direct instance from VAST_TEMPLATE_HASH_ID with cancel_unavail: true; the public key derived from the private key is injected via an onstart script that appends it to /root/.ssh/authorized_keys, and then attached through the SSH API as well. HTTP 429 is retried with exponential backoff up to six times.
Readiness additionally requires actual_status == "running" before the TCP and READY probes.
Remote execution
There is no SFTP daemon requirement and no cloud storage: everything goes through ssh, scp, and ssh-keygen subprocesses.
- SSH options —
ConnectTimeout=20,ServerAliveInterval=10,ServerAliveCountMax=3, the provider's private key via-i. Unless{PROVIDER}_SSH_STRICT_HOST_KEYis truthy,StrictHostKeyChecking=noandUserKnownHostsFile=/dev/nullare added. - Key hygiene — both provider clients refuse to construct without a readable key file, and every key path resolution
chmods the key to600, raising if it cannot. - Uploads stream over SSH stdin (
mkdir -p … && cat > path), so payloads never need a local temp file. - Downloads try
scpfirst and fall back tossh … 'test -f path && cat path'for endpoints with the SFTP subsystem disabled.allow_missing=Trueturns a missing file into a no-op. - Transient-error retries —
_with_vast_endpoint_retryre-resolves the SSH endpoint (refreshing it from the provider API, then falling back to the last known host/port), probes it, and retries up to four times with a3 + attemptsecond backoff. It only retries errors that look transient: connection refused/reset/timed out, no route to host, channel closed, endpoint not ready, or a failedssh/scp.
Remote layout
{PROVIDER}_REMOTE_BASE_DIR/{job_id}/
├── input/
│ ├── dataset.json # {"examples": [...], "metadata": {...}}
│ ├── training_manifest.json # the full training manifest
│ ├── metadata.json # just the dataset metadata block
│ └── train_unsloth.py # uploaded from the repository
├── output/
│ ├── train.pid
│ ├── checkpoints/step_N.safetensors
│ ├── final_adapter_dir/ # adapter_model.safetensors, tokenizer files
│ ├── final_adapter.safetensors
│ ├── training_summary.json
│ └── adapter_bundle.tar.gz # created during artifact sync
└── logs/train.log
Launching and monitoring
The command template (RUNPOD_TRAIN_COMMAND / VAST_TRAIN_COMMAND) is formatted with the remote paths and launched fully detached:
mkdir -p <output_dir> || exit 1
rm -f <output_dir>/train.pid
nohup setsid -f bash -lc 'echo $$ > <pid_file>; exec <command>' > <log_path> 2>&1 < /dev/null
# then poll for up to 5 s until the pid file is non-empty, and print it
If no numeric PID comes back, the stage fails. Monitoring then polls kill -0 <pid> every {PROVIDER}_POLL_INTERVAL_SEC (8, minimum 5) until the process exits or {PROVIDER}_MONITOR_TIMEOUT_SEC (7200) elapses, calling touch_job on every tick so the job is not requeued as stale.
When the process is gone, the stage verifies the adapter exists — final_adapter_dir/adapter_model.safetensors, or the legacy final_adapter.safetensors. If neither is there, the error includes the last 80 lines of the remote log, which is usually enough to diagnose an OOM or a missing dependency without opening an SSH session.
Artifact sync
state/models/{normalized_avatar_id}/ is wiped and recreated, then:
tar -czf adapter_bundle.tar.gz -C final_adapter_dir .runs on the instance;- the bundle is downloaded and extracted locally with
filter="data"; - a missing
adapter_model.safetensorsafter extraction is a hard failure; training_summary.jsonis fetched if present, andtrain.logonly whenTRAINING_DOWNLOAD_REMOTE_LOG=true— both best-effort;training_metadata.jsonis written locally with the run id, job id, avatar id, domain id, the full copied-file list, and a timestamp.
The avatar id is normalized for the directory name by lowercasing and collapsing anything outside [a-z0-9_-] into single underscores.
Then _publish_avatar_profile upserts the avatar row with the refined runtime system prompt (falling back to a short generated prompt when the rewrite failed and left nothing usable), the adapter directory, and the ElevenLabs voice id — setting lora_quality_passed=false and lora_quality_score=null until evaluation runs. Adapters stay on local disk: the profile records status: "stored_locally" and published_to_remote: false.
Remote training script
training/scripts/train_unsloth.py runs on the instance. It prefers the unsloth 4-bit backend and falls back to transformers + peft when unsloth is unavailable; --backend can force either.
Examples are rendered to a single text field:
### User
Instruction:
<instruction>
Input:
<input> # only when present
Context:
<context> # only when present
### Assistant
<output>
Rows without both an instruction and an output are dropped, as are token sequences shorter than two tokens; an empty result raises rather than training on nothing.
Defaults, overridable per-flag or from manifest.training.{lora,hyperparameters,callbacks}:
| Setting | Default |
|---|---|
| Base model | unsloth/Qwen3-32B-unsloth-bnb-4bit |
| LoRA rank / alpha / dropout | 16 / 32 / 0 |
| Target modules | q_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj |
| Batch size × grad accumulation | 8 × 2 |
| Epochs / max steps | 3 / 50 |
| Learning rate / scheduler / weight decay | 2e-4 / cosine / 0.01 |
| Precision | bf16 when supported, else fp16 on CUDA, else fp32 |
| Optimizer | adamw_8bit with bitsandbytes on CUDA, else adamw_torch |
| Checkpoint interval / logging | every 50 steps / every 10 steps |
| Max sequence length / seed | 4096 / 42 |
Three callbacks run: loss-based early stopping (patience 20 logging windows, min_delta 1e-4); a domain-aware variant that adds 6 windows of patience for legal_consulting, cybersecurity, and real_estate_consulting; and a checkpoint streamer that copies each saved adapter to checkpoints/step_N.safetensors.
On completion the adapter and tokenizer are saved to final_adapter_dir/, adapter_model.safetensors is copied to final_adapter.safetensors, and training_summary.json records the backend, runtime, dataset size, every hyperparameter, the final global step, the train loss, the adapter path, and the GPU name.
Completion
run_stage_complete writes state/runs/{run_id}/quality_gate.json — a handoff record asserting that review was approved, the run has an avatar and a domain, and artifacts were synced — then runs the LoRA evaluation, merges training, quality_gate, and lora_quality into the artifact manifest, marks the job completed, moves the run to completed, and appends training_completed.
LoRA quality evaluation
training/lora_quality.py grades the adapter against the base model on the domain's 20 held-out questions.
- Load the questions and the domain's
base_system_prompt. The avatar's storedsystem_promptoverrides the domain prompt when present. - Answer twice per question — once with no
lora_name(base) and once with the avatar's adapter — againstLLM_URLwith a bearer token fromLLM_API_TOKEN,temperature=0.1,max_tokensfromLORA_EVAL_MAX_TOKENS(256, clamped 64–512),stream: false. Both sides get an evaluation-mode instruction to answer concisely and stay under 120 words. All 40 requests run through one thread pool sized byLORA_EVAL_INFERENCE_CONCURRENCY(4, clamped 2–16). An empty answer raises. - Judge each pair with the OpenRouter model, in JSON mode, at
temperature=0, reasoning disabled,max_tokensfromLORA_EVAL_JUDGE_MAX_TOKENS(140, clamped 64–256), concurrencyLORA_EVAL_JUDGE_CONCURRENCY(4, clamped 1–10). The judge returns{analysis, accuracy, hallucination, improvement_vs_base}, validated by Pydantic withaccuracybounded 0–10. Each pair gets up to three attempts: JSON mode, then non-JSON mode, then a retry with an explicit "reply again with one raw JSON object" message appended. - Aggregate into
avg_accuracy,hallucination_rate(%), andimprovement_rate(%).
The rubric is deliberately forgiving about phrasing and detail: 7–8 is "clearly useful, mostly correct, practical"; hallucination is only for invented facts, unsupported claims, fake procedures, or unsafe guesses — not for a generic or incomplete answer; and improvement_vs_base requires a practically meaningful gain in correctness, specificity, domain precision, actionability, or reasoning, never merely being longer or reworded.
Pass criteria
All three must hold:
| Metric | Threshold |
|---|---|
avg_accuracy | >= 6.5 |
hallucination_rate | < 20 % |
improvement_rate | >= 50 % |
The verdict is written to ml.avatars as lora_quality_passed, lora_quality_score (the average accuracy), and lora_pass_report (the report path), and the full per-question report — including both answers and the judge's analysis — is written to state/models/{avatar_id}/lora_quality_report.json.
A failing verdict does not fail the run. The run still reaches completed; lora_quality_passed stays false. Since GET /v1/avatars?trained_only=true filters on that flag, a failed adapter simply never reaches the trained-avatar catalog.
Failure handling and cleanup
Every training.* task shares the same wrapper: classify, retry, then fail.
- Stale errors — a missing job or run, an invalid transition, a terminal run, or "cannot start training from run status" — retire the job and raise Celery's
Ignore. No retry. - Retryable errors increment
attempt_noand retry withcountdown = 2 ** retries, up to 4 attempts (3 fortraining.process). - Final failure calls
_best_effort_terminate_job_instance, marks the jobfailedwith the error text, and patches the run tofailed.
Termination is genuinely best-effort and idempotent. It reconstructs provider state from the run's artifact manifest, then the job payload, then — as a last resort — the most recent training_started / training_payload_uploaded / training_instance_ready / training_instance_requested event payload. A termination error is recorded as terminate_error in the manifest rather than raised, and an already-terminated instance is left alone. Each attempt appends a training_instance_cleanup event.
training.terminate_instance is the one stage that attempts termination before retrying, on the reasoning that a stuck terminate should still release the GPU.
Provider configuration
A provider counts as configured only when its enable flag is truthy, its API key and template are set, and its private key file exists. With neither configured, every stage fails fast with a message listing exactly what is missing.
Shared
| Variable | Default | Purpose |
|---|---|---|
TRAINING_LORA_ADAPTERS_ROOT / TRAINING_NAS_ARTIFACT_ROOT | state/models | Where adapter bundles are stored. |
TRAINING_DOWNLOAD_REMOTE_LOG | false | Also download train.log during artifact sync. |
TRAINING_MAX_TASK_RETRIES, TRAINING_TASK_RETRY_BASE_SEC | 3, 6 | Present in .env.example; task retry counts are currently hard-coded per task. |
LLM_URL, LLM_API_TOKEN | — | Core inference API used for LoRA evaluation answers. |
RunPod
| Variable | Default |
|---|---|
RUNPOD_ENABLED | false |
RUNPOD_API_KEY, RUNPOD_TEMPLATE_ID | — (both required) |
RUNPOD_GRAPHQL_URL, RUNPOD_REST_URL | https://api.runpod.io/graphql, https://rest.runpod.io/v1 |
RUNPOD_GPU_TYPE_ID, RUNPOD_GPU_TYPE_IDS | — (candidate list) |
RUNPOD_CLOUD_TYPE | SECURE |
RUNPOD_CONTAINER_DISK_GB, RUNPOD_VOLUME_GB, RUNPOD_VOLUME_MOUNT_PATH | 50, 20, /workspace |
RUNPOD_GPU_COUNT, RUNPOD_SUPPORT_PUBLIC_IP, RUNPOD_DOCKER_ARGS | 1, 1, empty |
RUNPOD_READY_TIMEOUT_SEC, RUNPOD_POLL_INTERVAL_SEC, RUNPOD_MONITOR_TIMEOUT_SEC | 900, 8, 7200 |
RUNPOD_SSH_USER, RUNPOD_SSH_KEY_PATH, RUNPOD_SSH_STRICT_HOST_KEY | root, — (required), 0 |
RUNPOD_REMOTE_BASE_DIR, RUNPOD_REMOTE_PYTHON_BIN, RUNPOD_REMOTE_SCRIPT_NAME | /workspace/training_job_runs, python, train_unsloth.py |
RUNPOD_TRAIN_COMMAND | {python_bin} {train_script_path} --dataset … --manifest … --metadata … --output … |
Vast.ai
| Variable | Default |
|---|---|
VAST_ENABLED | false |
VAST_API_KEY, VAST_TEMPLATE_HASH_ID (or VAST_TEMPLATE_ID) | — (both required) |
VAST_BASE_URL | https://console.vast.ai/api/v0 |
VAST_GPU_NAME, VAST_GPU_FALLBACK_NAMES | H100 SXM, H100 PCIe,NVIDIA H100 80GB HBM3,H100,NVIDIA H200 SXM |
VAST_RENT_TYPE, VAST_DISK_GB, VAST_NUM_GPUS, VAST_IMAGE | on-demand, 50, 1, unset |
VAST_READY_TIMEOUT_SEC, VAST_POLL_INTERVAL_SEC, VAST_MONITOR_TIMEOUT_SEC | 900, 8, 7200 |
VAST_SSH_USER, VAST_SSH_KEY_PATH, VAST_SSH_PUBLIC_KEY, VAST_SSH_STRICT_HOST_KEY | root, — (required), derived via ssh-keygen -y, 0 |
VAST_REMOTE_BASE_DIR, VAST_REMOTE_PYTHON_BIN, VAST_REMOTE_SCRIPT_NAME | /workspace/training_job_runs, python, train_unsloth.py |
VAST_TRAIN_COMMAND | same shape as the RunPod template |