From f788de7e78aab75244c2a9b6454d59f336691ca6 Mon Sep 17 00:00:00 2001 From: kalebbroo Date: Sun, 30 Aug 2026 22:31:40 -0400 Subject: [PATCH] Publish worker output onto the mounted network volume instead of re-uploading it The endpoint already has the network volume mounted at /runpod-volume, and RunPod's S3-compatible API addresses that same volume by the identical file paths. upload_file() was PUTting output files to that volume over HTTPS from handler.py, when it could copy them onto the mount directly. The S3 client is still needed to sign the returned object URL for HartsyWeb, since the volume is private and HartsyWeb runs outside RunPod. --- README.md | 8 ++++---- handler.py | 14 ++++++++++++-- tests/test_handler.py | 29 +++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index f8bb0ea..3497f31 100644 --- a/README.md +++ b/README.md @@ -28,21 +28,21 @@ The first successful push creates `kalebbroo/runpod-serverless-aitoolkit` if it The workflow publishes BuildKit provenance and an SBOM with release images and uses the GitHub Actions cache for subsequent builds. -Create a queue-based RunPod Serverless endpoint from that immutable image. Use one concurrent request per worker because `/dataset` and `/workspace/output` are intentionally reset for every request. Allocate enough container disk for the model cache plus dataset/output and choose a GPU configuration appropriate for the largest model families you enable. Hartsy sends a 24-hour execution policy and seven-day status TTL by default; the endpoint must permit those limits. +Create a queue-based RunPod Serverless endpoint from that immutable image with a network volume attached — outputs publish onto it at `AITK_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. Allocate enough container disk for the model cache plus dataset/output and choose a GPU configuration appropriate for the largest model families you enable. Hartsy sends a 24-hour execution policy and seven-day status TTL by default; the endpoint must permit those limits. Required worker environment: -- `AITK_S3_BUCKET`, `AITK_S3_ACCESS_KEY`, `AITK_S3_SECRET_KEY` +- `AITK_S3_BUCKET`, `AITK_S3_ACCESS_KEY`, `AITK_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. - `AITK_S3_ENDPOINT` for S3-compatible storage, when not using AWS - `AITK_S3_REGION` (default `us-east-1`) - `AITK_S3_PUBLIC_BASE_URL` for public objects, or omit it to return seven-day signed URLs - `AITK_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: `AITK_MAX_ARCHIVE_BYTES`, `AITK_MAX_EXTRACTED_BYTES`, `AITK_MAX_DATASET_ARCHIVES`, `AITK_URL_TTL_SECONDS`, `AITK_WORK_ROOT`, and `AI_TOOLKIT_PYTHON`. +Optional limits: `AITK_MAX_ARCHIVE_BYTES`, `AITK_MAX_EXTRACTED_BYTES`, `AITK_MAX_DATASET_ARCHIVES`, `AITK_URL_TTL_SECONDS`, `AITK_WORK_ROOT`, `AITK_VOLUME_ROOT`, and `AI_TOOLKIT_PYTHON`. Hartsy needs `RUNPOD_API_KEY` and `RUNPOD_AI_TOOLKIT_ENDPOINT_ID`. It can override `AI_TOOLKIT_EXECUTION_TIMEOUT_MS` and `AI_TOOLKIT_JOB_TTL_MS` within RunPod's seven-day maximum. -The worker validates its installed AI Toolkit revision and all worker-owned paths, safely downloads and extracts Hartsy's prepared dataset, launches unmodified `run.py`, reads the official `loss_log.db`, uploads new samples as they appear, emits structured RunPod progress updates, and returns an artifact manifest for durable Hartsy ingestion. +The worker validates its installed AI Toolkit revision and all worker-owned paths, safely downloads and extracts Hartsy's prepared dataset, launches unmodified `run.py`, reads the official `loss_log.db`, publishes new samples onto the network volume as they appear, emits structured RunPod progress updates, and returns an artifact manifest for durable Hartsy ingestion. ## Versioned request contract diff --git a/handler.py b/handler.py index 3fd799e..5ef35bf 100644 --- a/handler.py +++ b/handler.py @@ -27,6 +27,7 @@ EXPECTED_REVISION = "be995185f598c83abb990a088e9f634c4d36eb46" TOOLKIT_ROOT = Path(os.getenv("AI_TOOLKIT_ROOT", "/app/ai-toolkit")).resolve() WORK_ROOT = Path(os.getenv("AITK_WORK_ROOT", "/workspace")).resolve() +VOLUME_ROOT = Path(os.getenv("AITK_VOLUME_ROOT", "/runpod-volume")).resolve() DATASET_ROOT = Path("/dataset").resolve() OUTPUT_ROOT = (WORK_ROOT / "output").resolve() MAX_ARCHIVE_BYTES = max(1, int(os.getenv("AITK_MAX_ARCHIVE_BYTES", str(4 * 1024**3)))) @@ -239,12 +240,21 @@ def object_url(client, bucket: str, key: str) -> str: def upload_file(client, session_id: str, path: Path, relative_name: str) -> dict[str, Any]: - """Upload one worker output and return its provider-neutral manifest entry.""" + """Publish one worker output onto the attached network volume and return its provider-neutral manifest entry. + + The endpoint's network volume is already mounted at VOLUME_ROOT, and its object keys are the same paths as + RunPod's S3-compatible API uses for that bucket — so publishing is a local copy, not a network upload. `client` + is kept only to sign the returned URL via object_url(). + """ bucket = os.environ["AITK_S3_BUCKET"] safe_name = "/".join(part for part in Path(relative_name).parts if part not in {"", ".", ".."}) key = f"training/{session_id}/{safe_name}" + if not VOLUME_ROOT.is_mount(): + raise RuntimeError(f"{VOLUME_ROOT} is not a mounted network volume; cannot publish worker output") + destination = (VOLUME_ROOT / key).resolve() + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(path, destination) content_type = mimetypes.guess_type(path.name)[0] or "application/octet-stream" - client.upload_file(str(path), bucket, key, ExtraArgs={"ContentType": content_type}) return { "url": object_url(client, bucket, key), "fileName": path.name, diff --git a/tests/test_handler.py b/tests/test_handler.py index 12e9378..e74b5ec 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -151,6 +151,35 @@ def test_storage_endpoints_and_returned_urls_require_https(self): with self.assertRaisesRegex(ValueError, "must be an absolute HTTPS URL"): handler.object_url(presigner, "bucket", "artifact.safetensors") + def test_upload_file_copies_onto_the_mounted_network_volume_instead_of_uploading(self): + """Publishing a worker output is a local copy onto the attached volume, not a network upload.""" + with tempfile.TemporaryDirectory() as temporary: + source = Path(temporary) / "model.safetensors" + source.write_bytes(b"weights") + volume_root = Path(temporary) / "volume" + volume_root.mkdir() + client = mock.Mock() + with mock.patch.object(handler, "VOLUME_ROOT", volume_root), \ + mock.patch.object(Path, "is_mount", return_value=True), \ + mock.patch.dict(os.environ, {"AITK_S3_BUCKET": "bucket"}, clear=True), \ + mock.patch.object(handler, "object_url", return_value="https://storage.example/signed") as object_url: + manifest = handler.upload_file(client, "session", source, "artifacts/model.safetensors") + client.upload_file.assert_not_called() + object_url.assert_called_once_with(client, "bucket", "training/session/artifacts/model.safetensors") + self.assertEqual(b"weights", (volume_root / "training/session/artifacts/model.safetensors").read_bytes()) + self.assertEqual("https://storage.example/signed", manifest["url"]) + self.assertEqual("artifacts/model.safetensors", manifest["path"]) + + def test_upload_file_requires_a_real_network_volume_mount(self): + """A worker without an attached volume fails loudly instead of silently losing output.""" + with tempfile.TemporaryDirectory() as temporary: + source = Path(temporary) / "model.safetensors" + source.write_bytes(b"weights") + with mock.patch.object(handler, "VOLUME_ROOT", Path(temporary) / "not-mounted"), \ + mock.patch.dict(os.environ, {"AITK_S3_BUCKET": "bucket"}, clear=True): + with self.assertRaisesRegex(RuntimeError, "not a mounted network volume"): + handler.upload_file(mock.Mock(), "session", source, "artifacts/model.safetensors") + if __name__ == "__main__": unittest.main()