Skip to content

Repository files navigation

Hartsy SimpleTuner worker

This repository owns the worker image used by HartsyWeb's SimpleTuner backend. It turns unmodified bghira/SimpleTuner (simpletuner/train.py, invoked via accelerate launch 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 SimpleTunerCatalog.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 SimpleTuner is invoked as a subprocess, not imported

SimpleTuner is licensed AGPL-3.0. This worker's own code (handler.py, tests, Dockerfile) is MIT-licensed and never imports SimpleTuner's Python modules — it only launches accelerate launch ... simpletuner/train.py as a subprocess and reads the events SimpleTuner's own webhook system POSTs to a receiver this worker runs. Get a legal read before changing this boundary (e.g. before ever patching SimpleTuner's own source).

Build and push from this directory:

docker build --pull --build-arg SIMPLETUNER_REVISION=ebe3cb241d422427d11d3a15ee107fc929749f98 -t your-registry/hartsy-simpletuner:ebe3cb24 .
docker push your-registry/hartsy-simpletuner:ebe3cb24

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-simpletuner:latest
docker.io/kalebbroo/runpod-serverless-simpletuner: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 ST_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, and the worker binds a fixed local port (8891) for its webhook receiver. 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:

  • ST_S3_BUCKET, ST_S3_ACCESS_KEY, ST_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.
  • ST_S3_ENDPOINT for S3-compatible storage, when not using AWS
  • ST_S3_REGION (default us-east-1)
  • ST_S3_PUBLIC_BASE_URL for public objects, or omit it to return seven-day signed URLs
  • ST_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: ST_MAX_ARCHIVE_BYTES, ST_MAX_EXTRACTED_BYTES, ST_MAX_DATASET_ARCHIVES, ST_URL_TTL_SECONDS, ST_WORK_ROOT, ST_VOLUME_ROOT.

Hartsy needs RUNPOD_API_KEY and RUNPOD_SIMPLETUNER_ENDPOINT_ID. It can override SIMPLETUNER_EXECUTION_TIMEOUT_MS and SIMPLETUNER_JOB_TTL_MS within RunPod's seven-day maximum.

The worker validates its installed SimpleTuner revision and all worker-owned paths, safely downloads and extracts Hartsy's prepared dataset, writes config.json and multidatabackend.json as two separate files (SimpleTuner's data_backend_config field points at the latter — confirmed there's no inline-dataset config path), sets CONFIG_BACKEND=json/CONFIG_PATH=.../SIMPLETUNER_JOB_ID=... and launches unmodified accelerate launch ... simpletuner/train.py, receives SimpleTuner's own "raw" webhook events on a localhost HTTP server it runs in a background thread (confirmed payload shape: {"type", "job_id", "data": {"status"}, "progress": {"current", "total", "percent", "eta_seconds", "metrics": {"loss", ...}}}), republishes normalized progress through RunPod, and returns an artifact manifest for durable Hartsy ingestion once the subprocess exits.

Versioned request contract

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

{
  "input": {
    "contract_version": 1,
    "internal_job_id": 12345,
    "session_id": "st-12345-1788050000",
    "config": {
      "model_type": "lora",
      "model_family": "flux",
      "output_dir": "/workspace/output",
      "data_backend_config": "/workspace/output/multidatabackend.json",
      "webhook_config": {"webhook_type": "raw", "callback_url": "http://127.0.0.1:8891/progress"}
    },
    "data_backend": [
      {"id": "hartsy-dataset", "type": "local", "instance_data_dir": "/dataset", "default": true},
      {"id": "hartsy-text-embeds", "type": "local", "dataset_type": "text_embeds", "cache_dir": "/workspace/output/cache/text", "default": true}
    ],
    "dataset_urls": ["https://storage.example/signed-dataset.zip"],
    "simpletuner_revision": "ebe3cb241d422427d11d3a15ee107fc929749f98"
  },
  "policy": {
    "executionTimeout": 86400000,
    "ttl": 604800000
  }
}

config and data_backend are deliberately separate top-level fields, not one embedded object — that mirrors how SimpleTuner itself consumes them as two files (config.json plus whatever data_backend_config points at). The worker writes both to disk before launching training. The worker also rejects any webhook_config.callback_url that doesn't point at its own localhost receiver, so a compromised or malformed config can't redirect SimpleTuner's progress events (which may include prompt/caption text) to an external host.

Progress updates are published through runpod.serverless.progress_update() and include progress, message, step, totalSteps, loss, etaSeconds, and logsTail — sourced from SimpleTuner's own webhook events rather than a log file or database, unlike the sibling AI Toolkit/OneTrainer workers.

A completed job returns:

{
  "sessionId": "st-12345-1788050000",
  "simpleTunerRevision": "ebe3cb241d422427d11d3a15ee107fc929749f98",
  "finalLoss": 0.1842,
  "totalSteps": 2000,
  "artifacts": [
    {
      "url": "https://storage.example/training/pytorch_lora_weights.safetensors",
      "fileName": "pytorch_lora_weights.safetensors",
      "contentType": "application/octet-stream",
      "fileSize": 123456789,
      "path": "artifacts/pytorch_lora_weights.safetensors",
      "kind": "weights",
      "isPrimary": true
    }
  ],
  "samples": [],
  "sampleDetails": []
}

Safety and limits

  • Dataset downloads and every redirect must use an origin in ST_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, and unsafe session IDs are rejected. output_dir/data_backend_config must be the exact worker-owned paths, and webhook_config's callback URL must point at the worker's own localhost receiver.
  • Every data_backend entry needs a non-empty id; a text_embeds-typed backend is required alongside the primary dataset (confirmed against SimpleTuner's own dataloader validation), and instance_data_dir must stay inside /dataset.
  • The container base image, SimpleTuner 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, and the webhook receiver binds a single fixed local port.
  • The final artifact is uploaded to S3-compatible storage before its URL is 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, the multidatabackend text_embeds-required validation, the confirmed raw-webhook event shape, and both LoRA and full-fine-tune artifact discovery.

Known v1 limitations

  • No sample/validation images are uploaded yet. SimpleTuner's raw webhook payload for training.status events (the ones this worker parses) doesn't carry image URLs in what was confirmed against source at the pinned commit — a general (non-pinned-commit) pass over SimpleTuner's docs suggested some webhook message types may include an images field, but that wasn't independently re-confirmed. Wire this up once verified, rather than guessing at the field.
  • ControlNet, masked-loss, and multi-GPU/DeepSpeed training are deliberately out of scope (see SimpleTunerCatalog's doc comment on the HartsyWeb side).
  • caption_strategy and crop on the generated image data-backend entry were not independently re-confirmed against the field registry at the pinned commit — re-verify if dataset loading behaves unexpectedly.

About

Serverless endpoint to run SimpleTuner in RunPod

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages