Skip to main content

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.

OrderTaskRun status on successDoes
1training.request_instancetraining_instance_requestedRents a pod/instance from RunPod, or Vast.ai on fallback.
2training.wait_instance_readytraining_instance_readyPolls until SSH answers a READY probe.
3training.upload_payloadtraining_payload_uploadedUploads dataset, manifest, metadata, and train_unsloth.py.
4training.start_remotetraining_startedLaunches training detached and captures its PID.
5training.monitortraining_monitoringPolls the PID, then verifies the expected outputs exist.
6training.sync_artifactstraining_artifacts_syncedDownloads the adapter bundle and publishes the avatar profile.
7training.terminate_instance— (training_instance_terminated in the manifest)Releases the GPU.
8training.completecompletedWrites 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:

  1. The run status must be review_approved or postprocess_running.
  2. A training provider must be configured — see Provider configuration.
  3. artifact_manifest.final_dataset.examples must be non-empty, review.accepted must be true, and creator_id, avatar_id, and domain_id must 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 optionsConnectTimeout=20, ServerAliveInterval=10, ServerAliveCountMax=3, the provider's private key via -i. Unless {PROVIDER}_SSH_STRICT_HOST_KEY is truthy, StrictHostKeyChecking=no and UserKnownHostsFile=/dev/null are added.
  • Key hygiene — both provider clients refuse to construct without a readable key file, and every key path resolution chmods the key to 600, raising if it cannot.
  • Uploads stream over SSH stdin (mkdir -p … && cat > path), so payloads never need a local temp file.
  • Downloads try scp first and fall back to ssh … 'test -f path && cat path' for endpoints with the SFTP subsystem disabled. allow_missing=True turns a missing file into a no-op.
  • Transient-error retries_with_vast_endpoint_retry re-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 a 3 + attempt second backoff. It only retries errors that look transient: connection refused/reset/timed out, no route to host, channel closed, endpoint not ready, or a failed ssh/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:

  1. tar -czf adapter_bundle.tar.gz -C final_adapter_dir . runs on the instance;
  2. the bundle is downloaded and extracted locally with filter="data";
  3. a missing adapter_model.safetensors after extraction is a hard failure;
  4. training_summary.json is fetched if present, and train.log only when TRAINING_DOWNLOAD_REMOTE_LOG=true — both best-effort;
  5. training_metadata.json is 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}:

SettingDefault
Base modelunsloth/Qwen3-32B-unsloth-bnb-4bit
LoRA rank / alpha / dropout16 / 32 / 0
Target modulesq_proj, k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj
Batch size × grad accumulation8 × 2
Epochs / max steps3 / 50
Learning rate / scheduler / weight decay2e-4 / cosine / 0.01
Precisionbf16 when supported, else fp16 on CUDA, else fp32
Optimizeradamw_8bit with bitsandbytes on CUDA, else adamw_torch
Checkpoint interval / loggingevery 50 steps / every 10 steps
Max sequence length / seed4096 / 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.

  1. Load the questions and the domain's base_system_prompt. The avatar's stored system_prompt overrides the domain prompt when present.
  2. Answer twice per question — once with no lora_name (base) and once with the avatar's adapter — against LLM_URL with a bearer token from LLM_API_TOKEN, temperature=0.1, max_tokens from LORA_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 by LORA_EVAL_INFERENCE_CONCURRENCY (4, clamped 2–16). An empty answer raises.
  3. Judge each pair with the OpenRouter model, in JSON mode, at temperature=0, reasoning disabled, max_tokens from LORA_EVAL_JUDGE_MAX_TOKENS (140, clamped 64–256), concurrency LORA_EVAL_JUDGE_CONCURRENCY (4, clamped 1–10). The judge returns {analysis, accuracy, hallucination, improvement_vs_base}, validated by Pydantic with accuracy bounded 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.
  4. Aggregate into avg_accuracy, hallucination_rate (%), and improvement_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:

MetricThreshold
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.

warning

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_no and retry with countdown = 2 ** retries, up to 4 attempts (3 for training.process).
  • Final failure calls _best_effort_terminate_job_instance, marks the job failed with the error text, and patches the run to failed.

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

VariableDefaultPurpose
TRAINING_LORA_ADAPTERS_ROOT / TRAINING_NAS_ARTIFACT_ROOTstate/modelsWhere adapter bundles are stored.
TRAINING_DOWNLOAD_REMOTE_LOGfalseAlso download train.log during artifact sync.
TRAINING_MAX_TASK_RETRIES, TRAINING_TASK_RETRY_BASE_SEC3, 6Present in .env.example; task retry counts are currently hard-coded per task.
LLM_URL, LLM_API_TOKENCore inference API used for LoRA evaluation answers.

RunPod

VariableDefault
RUNPOD_ENABLEDfalse
RUNPOD_API_KEY, RUNPOD_TEMPLATE_ID— (both required)
RUNPOD_GRAPHQL_URL, RUNPOD_REST_URLhttps://api.runpod.io/graphql, https://rest.runpod.io/v1
RUNPOD_GPU_TYPE_ID, RUNPOD_GPU_TYPE_IDS— (candidate list)
RUNPOD_CLOUD_TYPESECURE
RUNPOD_CONTAINER_DISK_GB, RUNPOD_VOLUME_GB, RUNPOD_VOLUME_MOUNT_PATH50, 20, /workspace
RUNPOD_GPU_COUNT, RUNPOD_SUPPORT_PUBLIC_IP, RUNPOD_DOCKER_ARGS1, 1, empty
RUNPOD_READY_TIMEOUT_SEC, RUNPOD_POLL_INTERVAL_SEC, RUNPOD_MONITOR_TIMEOUT_SEC900, 8, 7200
RUNPOD_SSH_USER, RUNPOD_SSH_KEY_PATH, RUNPOD_SSH_STRICT_HOST_KEYroot, — (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

VariableDefault
VAST_ENABLEDfalse
VAST_API_KEY, VAST_TEMPLATE_HASH_ID (or VAST_TEMPLATE_ID)— (both required)
VAST_BASE_URLhttps://console.vast.ai/api/v0
VAST_GPU_NAME, VAST_GPU_FALLBACK_NAMESH100 SXM, H100 PCIe,NVIDIA H100 80GB HBM3,H100,NVIDIA H200 SXM
VAST_RENT_TYPE, VAST_DISK_GB, VAST_NUM_GPUS, VAST_IMAGEon-demand, 50, 1, unset
VAST_READY_TIMEOUT_SEC, VAST_POLL_INTERVAL_SEC, VAST_MONITOR_TIMEOUT_SEC900, 8, 7200
VAST_SSH_USER, VAST_SSH_KEY_PATH, VAST_SSH_PUBLIC_KEY, VAST_SSH_STRICT_HOST_KEYroot, — (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_COMMANDsame shape as the RunPod template