Publish worker output onto the mounted network volume instead of re-uploading it

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
 

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,

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()