Hartsy OneTrainer worker

This repository owns the worker image used by HartsyWeb's OneTrainer backend. It turns unmodified Nerogar/OneTrainer (scripts/train.py, invoked as a plain subprocess — never imported into this repository's own code) into Hartsy's versioned RunPod Serverless contract. Both the container and each request must match OneTrainerCatalog.UpstreamRevision; mismatches are rejected before a billable training process starts. Structured to mirror HartsyAI/Runpod-Serverless-AIToolkit — read that repo's README first if anything here is unclear, the two share almost all of their dataset-safety and storage-publishing code.

Why OneTrainer is invoked as a subprocess, not imported

OneTrainer is licensed AGPL-3.0. This worker's own code (handler.py, tests, Dockerfile) is MIT-licensed and never imports OneTrainer's Python modules — it only launches scripts/train.py as a subprocess and reads the TensorBoard event files and sample images that process writes to disk. Get a legal read before changing this boundary (e.g. before ever patching OneTrainer's own source).

Build and push from this directory:

docker build --pull --build-arg ONETRAINER_REVISION=23df3832e8f213d64c337e39f9365ba46e99fed0 -t your-registry/hartsy-onetrainer:23df3832 .
docker push your-registry/hartsy-onetrainer:23df3832

Automatic Docker Hub publishing

The Build and publish worker image GitHub Actions workflow tests and builds every pull request. A push to main, or a manual workflow dispatch, publishes the Linux AMD64 image to:

docker.io/kalebbroo/runpod-serverless-onetrainer:latest
docker.io/kalebbroo/runpod-serverless-onetrainer:sha-<full-git-commit>

Configure this GitHub Actions repository secret before publishing:

• DOCKERHUB_TOKEN: a Docker Hub personal access token with Read & Write permission

The public Docker Hub username kalebbroo is configured directly in the workflow and does not need to be stored as a secret. Do not store a Docker Hub password or token in the repository. Use the immutable sha-<full-git-commit> tag in the RunPod template; latest is a convenience tag for inspection and manual testing.

Create a queue-based RunPod Serverless endpoint from that immutable image with a network volume attached — outputs publish onto it at OT_VOLUME_ROOT (default /runpod-volume, RunPod's default Serverless mount point). Use one concurrent request per worker because /dataset and /workspace/output are intentionally reset for every request. This can reuse the same network volume and S3 API credentials already attached to the AI Toolkit endpoint — there's no need to provision separate storage.

Required worker environment:

• OT_S3_BUCKET, OT_S3_ACCESS_KEY, OT_S3_SECRET_KEY — the attached network volume's ID and an S3 API key with access to it. The worker publishes outputs with a local copy onto the volume mount, not a network upload; these credentials only sign the returned object URL.
• OT_S3_ENDPOINT for S3-compatible storage, when not using AWS
• OT_S3_REGION (default us-east-1)
• OT_S3_PUBLIC_BASE_URL for public objects, or omit it to return seven-day signed URLs
• OT_DATASET_ALLOWED_ORIGINS, a comma-separated list of exact HTTPS origins Hartsy may use for prepared archives, for example https://storage.hartsy.ai,https://example.r2.cloudflarestorage.com

Optional limits: OT_MAX_ARCHIVE_BYTES, OT_MAX_EXTRACTED_BYTES, OT_MAX_DATASET_ARCHIVES, OT_URL_TTL_SECONDS, OT_WORK_ROOT, OT_VOLUME_ROOT, and ONETRAINER_PYTHON.

Hartsy needs RUNPOD_API_KEY and RUNPOD_ONETRAINER_ENDPOINT_ID. It can override ONETRAINER_EXECUTION_TIMEOUT_MS and ONETRAINER_JOB_TTL_MS within RunPod's seven-day maximum.

The worker validates its installed OneTrainer revision and all worker-owned paths, safely downloads and extracts Hartsy's prepared dataset, writes config.json and concepts.json as two separate files (OneTrainer's concept_file_name config field points at the latter — it does not accept an embedded concepts array in the normal code path), launches unmodified scripts/train.py, tails the TensorBoard event file OneTrainer writes under workspace_dir/tensorboard/ for loss/train_step (falling back to smooth_loss/train_step), publishes new samples discovered under the nested workspace_dir/samples/{index} - {prompt}/ directories as they appear, emits structured RunPod progress updates, and returns an artifact manifest for durable Hartsy ingestion.

Versioned request contract

HartsyWeb submits contract version 1 using RunPod's /run endpoint:

{
  "input": {
    "contract_version": 1,
    "internal_job_id": 12345,
    "session_id": "ot-12345-1788050000",
    "config": {
      "model_type": "STABLE_DIFFUSION_XL_10_BASE",
      "training_method": "LORA",
      "workspace_dir": "/workspace/output",
      "concept_file_name": "/workspace/output/concepts.json",
      "output_model_destination": "/workspace/output/example_ot-12345-1788050000.safetensors",
      "output_model_format": "KOHYA_LORA"
    },
    "concepts": [
      {"name": "hartsy-dataset", "path": "/dataset", "text": {"prompt_source": "sample"}, "repeats": 1, "enabled": true}
    ],
    "dataset_urls": ["https://storage.example/signed-dataset.zip"],
    "onetrainer_revision": "23df3832e8f213d64c337e39f9365ba46e99fed0"
  },
  "policy": {
    "executionTimeout": 86400000,
    "ttl": 604800000
  }
}

config and concepts are deliberately separate top-level fields, not one embedded object — that mirrors how OneTrainer itself consumes them as two files (--config-path plus whatever concept_file_name points at). The worker writes both to disk before invoking train.py.

Progress updates are published through runpod.serverless.progress_update() and include message, step, loss, samples, and logsTail. There is no upfront totalSteps/progress percentage — OneTrainer is epoch-driven and its step count depends on dataset size, which is only known once training starts — so HartsyWeb falls back to its own execution-time-based estimate, the same fallback AIToolkitTrainingBackend uses whenever a worker's own progress figure is unavailable.

A completed job returns:

{
  "sessionId": "ot-12345-1788050000",
  "oneTrainerRevision": "23df3832e8f213d64c337e39f9365ba46e99fed0",
  "finalLoss": 0.1842,
  "totalSteps": 1250,
  "artifacts": [
    {
      "url": "https://storage.example/training/example.safetensors",
      "fileName": "example_ot-12345-1788050000.safetensors",
      "contentType": "application/octet-stream",
      "fileSize": 123456789,
      "path": "artifacts/example_ot-12345-1788050000.safetensors",
      "kind": "weights",
      "isPrimary": true
    }
  ],
  "samples": ["https://storage.example/training/sample_step_50.png"],
  "sampleDetails": []
}

Safety and limits

• Dataset downloads and every redirect must use an origin in OT_DATASET_ALLOWED_ORIGINS; the complete request is capped at 8 archives, 4 GiB compressed, and 8 GiB extracted by default.
• Custom S3 endpoints, public object bases, and every returned artifact URL must use HTTPS.
• ZIP path traversal, symbolic links, duplicate output paths, unsafe session IDs, and concept paths outside /dataset are rejected. workspace_dir/concept_file_name/output_model_destination must stay inside worker-owned roots.
• The container base image, OneTrainer commit, and worker dependencies are pinned. The handler also verifies the installed revision at runtime.
• The endpoint must use one concurrent request per worker. Each request resets /dataset and /workspace/output.
• Samples and final artifacts are uploaded to S3-compatible storage before their URLs are returned to HartsyWeb.

Test locally

python -m pip install -r requirements.txt
python -m unittest discover -s tests -v
python -m py_compile handler.py tests/test_handler.py
docker build --check .

These tests cover revision matching, path ownership, secure ZIP extraction, nested sample-directory discovery, and TensorBoard scalar-tag reading (via dependency injection, not by mocking Python's import machinery — the real tensorboard package is never imported during tests).

Known v1 limitations

• No sample-prompt configuration is sent (deliberately deferred — see OneTrainerCatalog's doc comment on the HartsyWeb side), so a run may produce zero samples until that's wired up.
• Single-concept datasets only; multi-concept/regularization-image datasets are deferred.
• Some field names in the generated config (marked in OneTrainerConfigFactory's comments) were confirmed against the pinned commit's source directly; a few narrower details (the full enumeration of valid ConceptTextConfig.prompt_source values, ConceptConfig's own top-level field names) were not independently re-confirmed — re-verify against source if training silently ignores a field.