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