From 8829c45f3d433f55426fc9dee550d0c632b3663b Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 16:17:26 -0400 Subject: [PATCH 01/23] feat: resumable network volume upload and download Adds runpod.network_volume_transfer (upload/download over the S3-compatible API) and `rp network-volume upload|download`. Multipart uploads resume from a local state file plus ListParts; downloads resume from a chunk sidecar; folder transfers skip unchanged files. --- runpod/__init__.py | 4 + runpod/network_volume_transfer.py | 722 ++++++++++++++++++++++++++ runpod/rp_cli/main.py | 117 +++++ tests/test_network_volume_transfer.py | 561 ++++++++++++++++++++ 4 files changed, 1404 insertions(+) create mode 100644 runpod/network_volume_transfer.py create mode 100644 tests/test_network_volume_transfer.py diff --git a/runpod/__init__.py b/runpod/__init__.py index 427023537..e9ddd406d 100644 --- a/runpod/__init__.py +++ b/runpod/__init__.py @@ -68,6 +68,10 @@ "patch": "runpod.apps.markers", "post": "runpod.apps.markers", "put": "runpod.apps.markers", + # network volume file transfers + "download_from_network_volume": "runpod.network_volume_transfer", + "network_volume_transfer": "runpod.network_volume_transfer", + "upload_to_network_volume": "runpod.network_volume_transfer", # logger "RunPodLogger": "runpod.serverless.modules.rp_logger", # submodule diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py new file mode 100644 index 000000000..d6b435b70 --- /dev/null +++ b/runpod/network_volume_transfer.py @@ -0,0 +1,722 @@ +"""resumable file and folder transfers to and from network volumes. + +transfers go through the runpod s3-compatible api. the bucket is the +network volume id and the endpoint is per datacenter. + +uploads above 64 MiB use multipart. the upload id is kept in a small +state file under ~/.runpod/transfers, so re-running an interrupted +upload sends only the parts the server does not already hold. downloads +fetch 64 MiB ranges in parallel into `.rpdownload`, recording +finished ranges in a sidecar json, so a re-run fetches only what is +missing. folder transfers skip files that are already present and +unchanged. +""" + +import hashlib +import json +import math +import os +import random +import threading +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from pathlib import Path, PurePosixPath +from typing import Any, Callable, Dict, List, Optional, Tuple + +MIB = 1024 * 1024 +SINGLE_PUT_MAX = 64 * MIB +DEFAULT_PART_SIZE = 64 * MIB +MAX_PART_SIZE = 500 * MIB # runpod's maximum multipart part size +MAX_PARTS = 10000 +DOWNLOAD_CHUNK_SIZE = 64 * MIB +RETRY_ATTEMPTS = 8 +UPLOADS_DIR = ".s3compat_uploads" + +ACCESS_KEY_ENV = "RUNPOD_S3_ACCESS_KEY" +SECRET_KEY_ENV = "RUNPOD_S3_SECRET_KEY" +S3_KEYS_URL = "https://console.runpod.io/user/settings" + +# progress(event, path, nbytes). events: "start" (nbytes = file size), +# "advance" (nbytes = delta), "skip" (nbytes = file size), "done". +Progress = Callable[[str, str, int], None] + +_sleep = time.sleep + + +class NetworkVolumeTransferError(RuntimeError): + """a transfer failed. re-running the same call resumes it.""" + + +class _Retryable(Exception): + """a transient failure raised by this module (e.g. a short read).""" + + +def _state_dir() -> Path: + return Path(os.path.expanduser("~/.runpod/transfers")) + + +# --------------------------------------------------------------- helpers + + +def part_size_for(size: int) -> int: + """part size for a multipart upload of `size` bytes.""" + part_size = DEFAULT_PART_SIZE + if math.ceil(size / part_size) > MAX_PARTS: + part_size = math.ceil(math.ceil(size / MAX_PARTS) / MIB) * MIB + if part_size > MAX_PART_SIZE: + raise NetworkVolumeTransferError( + f"file of {size} bytes needs parts above the 500 MiB limit" + ) + return part_size + + +def _error_code(exc: Exception) -> Tuple[str, int]: + response = getattr(exc, "response", None) or {} + code = str(response.get("Error", {}).get("Code", "")) + status = response.get("ResponseMetadata", {}).get("HTTPStatusCode") or 0 + return code, int(status) + + +def _is_not_found(exc: Exception) -> bool: + code, status = _error_code(exc) + return status == 404 or code in ("404", "NoSuchKey", "NotFound") + + +def _is_retryable(exc: Exception) -> bool: + from botocore.exceptions import ( # pylint: disable=import-outside-toplevel + ClientError, + ConnectionError as BotoConnectionError, + HTTPClientError, + IncompleteReadError, + ResponseStreamingError, + ) + + if isinstance(exc, _Retryable): + return True + if isinstance( + exc, + ( + BotoConnectionError, + HTTPClientError, + IncompleteReadError, + ResponseStreamingError, + ConnectionError, + TimeoutError, + ), + ): + return True + if isinstance(exc, ClientError): + code, status = _error_code(exc) + if status == 429 or status >= 500: + return True + return code in ("SlowDown", "RequestTimeout", "InternalError", "524") + return False + + +def _with_retries(func: Callable[[], Any], attempts: int = RETRY_ATTEMPTS) -> Any: + """call func, retrying transient failures with exponential backoff.""" + for attempt in range(attempts): + try: + return func() + except Exception as exc: # noqa: BLE001 - filtered by _is_retryable + if attempt == attempts - 1 or not _is_retryable(exc): + raise + _sleep(min(2**attempt, 30) + random.uniform(0, 1)) + raise AssertionError("unreachable") + + +def _write_json_atomic(path: Path, data: Dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "w", encoding="utf-8") as handle: + json.dump(data, handle) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + + +def _read_json(path: Path) -> Optional[Dict[str, Any]]: + try: + with open(path, encoding="utf-8") as handle: + data = json.load(handle) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None + + +def _remove(path: Path) -> None: + try: + path.unlink() + except FileNotFoundError: + pass + + +def _read_range(path: Path, offset: int, length: int) -> bytes: + with open(path, "rb") as handle: + handle.seek(offset) + return handle.read(length) + + +def _md5_hex(path: Path, offset: int, length: int) -> str: + digest = hashlib.md5() # noqa: S324 - etag comparison, not security + with open(path, "rb") as handle: + handle.seek(offset) + remaining = length + while remaining > 0: + block = handle.read(min(8 * MIB, remaining)) + if not block: + break + digest.update(block) + remaining -= len(block) + return digest.hexdigest() + + +def _is_md5_etag(etag: str) -> bool: + return len(etag) == 32 and all(c in "0123456789abcdef" for c in etag.lower()) + + +def _run_pool(tasks: List[Callable[[], Any]], concurrency: int) -> List[Any]: + """run tasks on a thread pool. the first failure cancels the rest.""" + results = [] + with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool: + futures = [pool.submit(task) for task in tasks] + try: + for future in as_completed(futures): + results.append(future.result()) + except BaseException: + for future in futures: + future.cancel() + raise + return results + + +def _strip_etag(etag: Any) -> str: + return str(etag or "").strip('"') + + +# ------------------------------------------------------------- transfers + + +class _Transfer: + """one bucket on one endpoint. the public functions build this.""" + + def __init__( + self, + client: Any, + bucket: str, + endpoint: str, + concurrency: int = 4, + progress: Optional[Progress] = None, + state_dir: Optional[Path] = None, + ): + self.client = client + self.bucket = bucket + self.endpoint = endpoint + self.concurrency = max(1, int(concurrency)) + self._progress = progress + self.state_dir = state_dir or _state_dir() + self._lock = threading.Lock() + self.summary = {"files_transferred": 0, "files_skipped": 0, "bytes_transferred": 0} + + # -- bookkeeping + + def _emit(self, event: str, path: str, nbytes: int = 0) -> None: + with self._lock: + if event == "advance": + self.summary["bytes_transferred"] += nbytes + elif event == "done": + self.summary["files_transferred"] += 1 + elif event == "skip": + self.summary["files_skipped"] += 1 + if self._progress is not None: + self._progress(event, path, nbytes) + + def _state_path(self, key: str, local: Path) -> Path: + ident = f"{self.endpoint}|{self.bucket}|{key}|{local.resolve()}" + return self.state_dir / (hashlib.sha256(ident.encode()).hexdigest() + ".json") + + # -- remote metadata + + def head(self, key: str) -> Optional[Dict[str, Any]]: + try: + return _with_retries( + lambda: self.client.head_object(Bucket=self.bucket, Key=key) + ) + except Exception as exc: # noqa: BLE001 + if _is_not_found(exc): + return None + raise + + def list_prefix(self, prefix: str) -> List[Dict[str, Any]]: + objects: List[Dict[str, Any]] = [] + kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Prefix": prefix} + while True: + try: + page = _with_retries(lambda: self.client.list_objects_v2(**kwargs)) + except Exception as exc: # noqa: BLE001 + if _is_not_found(exc): + return objects + raise + objects.extend(page.get("Contents") or []) + token = page.get("NextContinuationToken") + if not page.get("IsTruncated") or not token: + return objects + kwargs["ContinuationToken"] = token + + def _list_parts(self, key: str, upload_id: str) -> List[Dict[str, Any]]: + parts: List[Dict[str, Any]] = [] + kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Key": key, "UploadId": upload_id} + while True: + page = _with_retries(lambda: self.client.list_parts(**kwargs)) + parts.extend(page.get("Parts") or []) + marker = page.get("NextPartNumberMarker") + if not page.get("IsTruncated") or not marker: + return parts + kwargs["PartNumberMarker"] = marker + + # -- upload + + def upload_file(self, local: Path, key: str) -> None: + stat = local.stat() + size = stat.st_size + label = str(local) + self._emit("start", label, size) + if size <= SINGLE_PUT_MAX: + data = local.read_bytes() + _with_retries( + lambda: self.client.put_object(Bucket=self.bucket, Key=key, Body=data) + ) + self._emit("advance", label, size) + else: + self._multipart_upload(local, key, size, stat.st_mtime_ns, label) + self._emit("done", label) + + def _resume_parts( + self, local: Path, key: str, state: Dict[str, Any], size: int, part_size: int + ) -> Dict[int, str]: + """parts of the saved upload that match the local file.""" + upload_id = state["upload_id"] + try: + listed = self._list_parts(key, upload_id) + except Exception as exc: # noqa: BLE001 + code, _ = _error_code(exc) + if code == "NoSuchUpload" or _is_not_found(exc): + raise LookupError(upload_id) from exc + raise + done: Dict[int, str] = {} + for part in listed: + number = int(part["PartNumber"]) + offset = (number - 1) * part_size + expected = min(part_size, size - offset) + if expected <= 0 or int(part.get("Size", -1)) != expected: + continue + etag = _strip_etag(part.get("ETag")) + if _is_md5_etag(etag) and _md5_hex(local, offset, expected) != etag.lower(): + continue + done[number] = part.get("ETag") + return done + + def _multipart_upload( + self, local: Path, key: str, size: int, mtime_ns: int, label: str + ) -> None: + part_size = part_size_for(size) + count = math.ceil(size / part_size) + state_path = self._state_path(key, local) + state = _read_json(state_path) + upload_id: Optional[str] = None + done: Dict[int, str] = {} + + if state and state.get("upload_id"): + same = ( + state.get("file_size") == size + and state.get("mtime_ns") == mtime_ns + and state.get("part_size") == part_size + ) + if same: + try: + done = self._resume_parts(local, key, state, size, part_size) + upload_id = state["upload_id"] + except LookupError: + pass + else: + try: + self.client.abort_multipart_upload( + Bucket=self.bucket, Key=key, UploadId=state["upload_id"] + ) + except Exception: # noqa: BLE001 - best effort + pass + if upload_id is None: + _remove(state_path) + + if upload_id is None: + created = _with_retries( + lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=key) + ) + upload_id = created["UploadId"] + _write_json_atomic( + state_path, + { + "upload_id": upload_id, + "part_size": part_size, + "file_size": size, + "mtime_ns": mtime_ns, + "bucket": self.bucket, + "key": key, + }, + ) + + already = sum(min(part_size, size - (n - 1) * part_size) for n in done) + if already: + self._emit("advance", label, already) + + def upload_part(number: int) -> Tuple[int, str]: + offset = (number - 1) * part_size + length = min(part_size, size - offset) + + def send(): + body = _read_range(local, offset, length) + return self.client.upload_part( + Bucket=self.bucket, + Key=key, + UploadId=upload_id, + PartNumber=number, + Body=body, + ) + + response = _with_retries(send) + self._emit("advance", label, length) + return number, response["ETag"] + + todo = [n for n in range(1, count + 1) if n not in done] + for number, etag in _run_pool( + [lambda n=n: upload_part(n) for n in todo], self.concurrency + ): + done[number] = etag + + parts = [{"PartNumber": n, "ETag": done[n]} for n in sorted(done)] + try: + _with_retries( + lambda: self.client.complete_multipart_upload( + Bucket=self.bucket, + Key=key, + UploadId=upload_id, + MultipartUpload={"Parts": parts}, + ) + ) + except Exception: + # a slow complete can time out after the server finished it + head = self.head(key) + if head is None or int(head.get("ContentLength", -1)) != size: + raise + _remove(state_path) + + def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bool: + if not remote: + return False + stat = local.stat() + size = remote.get("Size", remote.get("ContentLength")) + modified = remote.get("LastModified") + if size is None or int(size) != stat.st_size or modified is None: + return False + return modified.timestamp() >= stat.st_mtime + + def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: + if not local.exists(): + raise FileNotFoundError(str(local)) + if local.is_dir(): + prefix = (local.resolve().name if remote_path is None else remote_path).strip("/") + remote = { + obj["Key"]: obj + for obj in self.list_prefix(f"{prefix}/" if prefix else "") + } + jobs = [] + for path in sorted(p for p in local.rglob("*") if p.is_file()): + rel = path.relative_to(local).as_posix() + key = f"{prefix}/{rel}" if prefix else rel + if self._unchanged_remote(remote.get(key), path): + self._emit("skip", str(path), path.stat().st_size) + continue + jobs.append((path, key)) + self._run_files([lambda p=p, k=k: self.upload_file(p, k) for p, k in jobs], jobs) + return self.summary + + if remote_path is None: + key = local.name + elif remote_path.endswith("/"): + key = remote_path.lstrip("/") + local.name + else: + key = remote_path.lstrip("/") + if self._unchanged_remote(self.head(key), local): + self._emit("skip", str(local), local.stat().st_size) + return self.summary + self.upload_file(local, key) + return self.summary + + def _run_files(self, tasks: List[Callable[[], Any]], jobs: List[Tuple[Any, Any]]) -> None: + """run per-file tasks. every file is attempted; failures are reported together.""" + failures = [] + + def guarded(index: int) -> None: + try: + tasks[index]() + except Exception as exc: # noqa: BLE001 - collected and re-raised + with self._lock: + failures.append(f"{jobs[index][0]}: {exc}") + + _run_pool([lambda i=i: guarded(i) for i in range(len(tasks))], self.concurrency) + if failures: + raise NetworkVolumeTransferError( + f"{len(failures)} file(s) failed:\n " + "\n ".join(sorted(failures)) + ) + + # -- download + + def download_file(self, key: str, size: int, etag: str, dest: Path) -> None: + label = str(dest) + if dest.is_file() and dest.stat().st_size == size: + self._emit("skip", label, size) + return + self._emit("start", label, size) + dest.parent.mkdir(parents=True, exist_ok=True) + tmp = dest.with_name(dest.name + ".rpdownload") + sidecar = dest.with_name(dest.name + ".rpdownload.json") + state = _read_json(sidecar) + done: set = set() + if ( + state + and state.get("size") == size + and state.get("etag") == etag + and state.get("chunk_size") == DOWNLOAD_CHUNK_SIZE + and tmp.is_file() + and tmp.stat().st_size == size + ): + done = {int(i) for i in state.get("done", [])} + else: + with open(tmp, "wb") as handle: + handle.truncate(size) + _write_json_atomic( + sidecar, + {"size": size, "etag": etag, "chunk_size": DOWNLOAD_CHUNK_SIZE, "done": []}, + ) + + count = math.ceil(size / DOWNLOAD_CHUNK_SIZE) + already = sum( + min(DOWNLOAD_CHUNK_SIZE, size - i * DOWNLOAD_CHUNK_SIZE) for i in done if i < count + ) + if already: + self._emit("advance", label, already) + + def fetch(index: int) -> None: + start = index * DOWNLOAD_CHUNK_SIZE + end = min(start + DOWNLOAD_CHUNK_SIZE, size) - 1 + + def get(): + response = self.client.get_object( + Bucket=self.bucket, Key=key, Range=f"bytes={start}-{end}" + ) + data = response["Body"].read() + if len(data) != end - start + 1: + raise _Retryable(f"short read for {key} bytes {start}-{end}") + return data + + data = _with_retries(get) + with open(tmp, "r+b") as handle: + handle.seek(start) + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) + with self._lock: + done.add(index) + _write_json_atomic( + sidecar, + { + "size": size, + "etag": etag, + "chunk_size": DOWNLOAD_CHUNK_SIZE, + "done": sorted(done), + }, + ) + self._emit("advance", label, end - start + 1) + + todo = [i for i in range(count) if i not in done] + _run_pool([lambda i=i: fetch(i) for i in todo], self.concurrency) + + if tmp.stat().st_size != size: + raise NetworkVolumeTransferError(f"{dest}: size mismatch after download") + os.replace(tmp, dest) + _remove(sidecar) + self._emit("done", label) + + def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, int]: + key = remote_path.lstrip("/") + head = None if key == "" or key.endswith("/") else self.head(key) + if head is not None: + name = PurePosixPath(key).name + dest = Path(local_path) if local_path is not None else Path(name) + if dest.is_dir(): + dest = dest / name + self.download_file( + key, int(head["ContentLength"]), _strip_etag(head.get("ETag")), dest + ) + return self.summary + + prefix = key.rstrip("/") + list_prefix = f"{prefix}/" if prefix else "" + objects = [ + obj + for obj in self.list_prefix(list_prefix) + if not obj["Key"].endswith("/") + and UPLOADS_DIR not in PurePosixPath(obj["Key"]).parts + ] + if not objects: + raise FileNotFoundError(f"no such file or directory on the volume: {remote_path}") + root = Path(local_path) if local_path is not None else Path( + PurePosixPath(prefix).name or self.bucket + ) + jobs = [(_safe_target(root, obj["Key"][len(list_prefix):]), obj) for obj in objects] + self._run_files( + [ + lambda d=d, o=o: self.download_file( + o["Key"], int(o["Size"]), _strip_etag(o.get("ETag")), d + ) + for d, o in jobs + ], + jobs, + ) + return self.summary + + +def _safe_target(root: Path, rel: str) -> Path: + """local path for a remote relative key, refusing anything outside root.""" + parts = PurePosixPath(rel).parts + if ( + not rel + or rel.startswith("/") + or "\\" in rel + or any(part in ("..", "") for part in parts) + or (parts and ":" in parts[0]) + ): + raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") + target = root.joinpath(*parts) + resolved_root = root.resolve() + if resolved_root not in target.resolve().parents: + raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") + return target + + +# ---------------------------------------------------------------- client + + +def _lookup_datacenter(volume_id: str) -> str: + from runpod.api.rest import run_rest_request # pylint: disable=import-outside-toplevel + + data = run_rest_request("GET", f"/v2/network-volumes/{volume_id}") or {} + datacenter = data.get("dataCenter") or data.get("dataCenterId") + if not datacenter: + raise NetworkVolumeTransferError( + f"could not find the datacenter of network volume {volume_id}" + ) + return str(datacenter) + + +def _make_transfer( + volume_id: str, + datacenter: Optional[str], + access_key: Optional[str], + secret_key: Optional[str], + endpoint_url: Optional[str], + concurrency: int, + progress: Optional[Progress], +) -> _Transfer: + import boto3 # pylint: disable=import-outside-toplevel + from botocore.config import Config # pylint: disable=import-outside-toplevel + + access_key = access_key or os.environ.get(ACCESS_KEY_ENV) + secret_key = secret_key or os.environ.get(SECRET_KEY_ENV) + if not access_key or not secret_key: + raise NetworkVolumeTransferError( + f"s3 api keys required: set {ACCESS_KEY_ENV} and {SECRET_KEY_ENV} " + f"(create keys at {S3_KEYS_URL})" + ) + datacenter = datacenter or _lookup_datacenter(volume_id) + endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" + concurrency = max(1, int(concurrency)) + config = Config( + s3={"addressing_style": "path"}, + request_checksum_calculation="when_required", + response_checksum_validation="when_required", + retries={"max_attempts": 10, "mode": "standard"}, + connect_timeout=60, + read_timeout=3600, + max_pool_connections=max(10, concurrency * concurrency), + ) + client = boto3.session.Session().client( + "s3", + region_name=datacenter, + endpoint_url=endpoint, + aws_access_key_id=access_key, + aws_secret_access_key=secret_key, + config=config, + ) + return _Transfer(client, volume_id, endpoint, concurrency, progress) + + +def upload( + volume_id: str, + local_path, + remote_path: Optional[str] = None, + *, + datacenter: Optional[str] = None, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + endpoint_url: Optional[str] = None, + concurrency: int = 4, + progress: Optional[Progress] = None, +) -> Dict[str, int]: + """upload a file or folder to a network volume. + + `remote_path` defaults to the local basename. a folder's contents go + under `remote_path/`. for a file, a `remote_path` ending in "/" keeps + the local file name. unchanged files (same size, remote not older) + are skipped, and an interrupted multipart upload resumes on re-run. + + s3 keys come from the arguments or RUNPOD_S3_ACCESS_KEY and + RUNPOD_S3_SECRET_KEY. without `datacenter`, the volume's datacenter is + looked up with the runpod api key. + + returns {files_transferred, files_skipped, bytes_transferred}. + """ + transfer = _make_transfer( + volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress + ) + return transfer.upload(Path(local_path), remote_path) + + +def download( + volume_id: str, + remote_path: str, + local_path=None, + *, + datacenter: Optional[str] = None, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + endpoint_url: Optional[str] = None, + concurrency: int = 4, + progress: Optional[Progress] = None, +) -> Dict[str, int]: + """download a file or folder from a network volume. + + a `remote_path` ending in "/" (or one that only exists as a prefix) is + a folder. `local_path` defaults to the remote basename. local files of + the same size are skipped, and partial downloads resume on re-run. + + returns {files_transferred, files_skipped, bytes_transferred}. + """ + transfer = _make_transfer( + volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress + ) + return transfer.download(remote_path, None if local_path is None else Path(local_path)) + + +upload_to_network_volume = upload +download_from_network_volume = download diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index 5be6aaef0..ac92a42c6 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1086,6 +1086,123 @@ def undeploy(app_name, env_name, yes): ) +# -------------------------------------------------------- network volumes + + +@cli.group(name="network-volume") +def network_volume(): + """Upload and download files on network volumes (resumable). + + Uses the S3-compatible API. Set RUNPOD_S3_ACCESS_KEY and + RUNPOD_S3_SECRET_KEY (create keys at console.runpod.io/user/settings). + Re-run an interrupted command to resume it. + """ + + +def _transfer_options(func): + options = [ + click.option( + "--datacenter", + default=None, + help="Volume datacenter, e.g. EU-RO-1 (looked up with your API key if omitted).", + ), + click.option("--s3-access-key", default=None, help="S3 API access key."), + click.option("--s3-secret-key", default=None, help="S3 API secret key."), + click.option("--endpoint-url", default=None, hidden=True), + click.option( + "--concurrency", default=4, show_default=True, type=click.IntRange(1, 64), + help="Parallel parts, chunks, or files.", + ), + ] + for option in reversed(options): + func = option(func) + return func + + +def _run_transfer(direction, volume_id, source, target, opts) -> None: + from rich.console import Console + from rich.progress import ( + BarColumn, + DownloadColumn, + Progress, + TextColumn, + TransferSpeedColumn, + ) + + from runpod import network_volume_transfer as nvt + + tasks = {} + with Progress( + TextColumn("{task.description}"), + BarColumn(), + DownloadColumn(), + TransferSpeedColumn(), + console=Console(stderr=True), + transient=True, + ) as bar: + + def progress(event, path, nbytes): + if event == "start": + tasks[path] = bar.add_task(Path(path).name, total=nbytes) + elif event == "advance" and path in tasks: + bar.advance(tasks[path], nbytes) + elif event == "done" and path in tasks: + bar.remove_task(tasks.pop(path)) + bar.console.print(f" {direction} {path}") + + func = nvt.upload if direction == "uploaded" else nvt.download + try: + summary = func( + volume_id, + source, + target, + datacenter=opts["datacenter"], + access_key=opts["s3_access_key"], + secret_key=opts["s3_secret_key"], + endpoint_url=opts["endpoint_url"], + concurrency=opts["concurrency"], + progress=progress, + ) + except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly + raise click.ClickException( + f"{exc}\nre-run the same command to resume" + ) from exc + + click.echo( + f"{direction} {summary['files_transferred']} file(s), " + f"{summary['bytes_transferred']} bytes; " + f"skipped {summary['files_skipped']} unchanged" + ) + + +@network_volume.command(name="upload") +@click.argument("volume_id") +@click.argument("local_path", type=click.Path(exists=True)) +@click.argument("remote_path", required=False) +@_transfer_options +def network_volume_upload(volume_id, local_path, remote_path, **opts): + """Upload a file or folder to a network volume. + + REMOTE_PATH defaults to the local name. A folder's contents go under + REMOTE_PATH/. Unchanged files are skipped. + """ + _run_transfer("uploaded", volume_id, local_path, remote_path, opts) + + +@network_volume.command(name="download") +@click.argument("volume_id") +@click.argument("remote_path") +@click.argument("local_path", required=False) +@_transfer_options +def network_volume_download(volume_id, remote_path, local_path, **opts): + """Download a file or folder from a network volume. + + A REMOTE_PATH ending in / is a folder. LOCAL_PATH defaults to the + remote name. Local files of the same size are skipped. + """ + _run_transfer("downloaded", volume_id, remote_path, local_path, opts) + + # -------------------------------------------------------- legacy groups diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py new file mode 100644 index 000000000..2efc76c4a --- /dev/null +++ b/tests/test_network_volume_transfer.py @@ -0,0 +1,561 @@ +"""tests for resumable network volume transfers against an in-process fake s3.""" + +import hashlib +import io +import json +import os +from datetime import datetime, timedelta, timezone +from unittest.mock import patch + +import pytest +from botocore.exceptions import ClientError, ReadTimeoutError +from click.testing import CliRunner + +from runpod import network_volume_transfer as nvt + + +def _client_error(code, status, op="Op"): + return ClientError( + {"Error": {"Code": code}, "ResponseMetadata": {"HTTPStatusCode": status}}, op + ) + + +class FakeS3: + """the subset of the s3 api the transfer module uses.""" + + def __init__(self, md5_etags=True): + self.objects = {} # key -> (data, last_modified) + self.uploads = {} # upload_id -> {"key": str, "parts": {n: bytes}} + self.calls = [] + self.fail = None # fail(op, kwargs) -> exception or None + self.md5_etags = md5_etags + self._next = 0 + + def _etag(self, data): + if self.md5_etags: + return '"' + hashlib.md5(data).hexdigest() + '"' + return '"opaque-' + str(len(data)) + '"' + + def _call(self, op, **kwargs): + self.calls.append((op, kwargs)) + if self.fail is not None: + exc = self.fail(op, kwargs) + if exc is not None: + raise exc + + def put(self, key, data, modified=None): + self.objects[key] = (data, modified or datetime.now(timezone.utc)) + + # -- objects + + def put_object(self, Bucket, Key, Body): + self._call("put_object", Key=Key) + self.put(Key, bytes(Body)) + return {"ETag": self._etag(bytes(Body))} + + def head_object(self, Bucket, Key): + self._call("head_object", Key=Key) + if Key not in self.objects: + raise _client_error("404", 404, "HeadObject") + data, modified = self.objects[Key] + return {"ContentLength": len(data), "ETag": self._etag(data), "LastModified": modified} + + def get_object(self, Bucket, Key, Range): + self._call("get_object", Key=Key, Range=Range) + data, _ = self.objects[Key] + start, end = (int(x) for x in Range.split("=")[1].split("-")) + return {"Body": io.BytesIO(data[start : end + 1])} + + def list_objects_v2(self, Bucket, Prefix, ContinuationToken=None): + self._call("list_objects_v2", Prefix=Prefix, ContinuationToken=ContinuationToken) + keys = sorted(k for k in self.objects if k.startswith(Prefix)) + start = int(ContinuationToken or 0) + page = keys[start : start + 2] + truncated = start + 2 < len(keys) + return { + "Contents": [ + { + "Key": k, + "Size": len(self.objects[k][0]), + "ETag": self._etag(self.objects[k][0]), + "LastModified": self.objects[k][1], + } + for k in page + ], + "IsTruncated": truncated, + "NextContinuationToken": str(start + 2) if truncated else None, + } + + # -- multipart + + def create_multipart_upload(self, Bucket, Key): + self._call("create_multipart_upload", Key=Key) + self._next += 1 + upload_id = f"up-{self._next}" + self.uploads[upload_id] = {"key": Key, "parts": {}} + return {"UploadId": upload_id} + + def _upload(self, upload_id): + if upload_id not in self.uploads: + raise _client_error("NoSuchUpload", 404, "Multipart") + return self.uploads[upload_id] + + def upload_part(self, Bucket, Key, UploadId, PartNumber, Body): + self._call("upload_part", PartNumber=PartNumber) + self._upload(UploadId)["parts"][PartNumber] = bytes(Body) + return {"ETag": self._etag(bytes(Body))} + + def list_parts(self, Bucket, Key, UploadId, PartNumberMarker=0): + self._call("list_parts", UploadId=UploadId, PartNumberMarker=PartNumberMarker) + parts = sorted(self._upload(UploadId)["parts"].items()) + parts = [p for p in parts if p[0] > int(PartNumberMarker)] + page, rest = parts[:2], parts[2:] + return { + "Parts": [ + {"PartNumber": n, "Size": len(d), "ETag": self._etag(d)} for n, d in page + ], + "IsTruncated": bool(rest), + "NextPartNumberMarker": page[-1][0] if rest else None, + } + + def complete_multipart_upload(self, Bucket, Key, UploadId, MultipartUpload): + self._call("complete_multipart_upload", UploadId=UploadId) + upload = self._upload(UploadId) + data = b"".join(upload["parts"][p["PartNumber"]] for p in MultipartUpload["Parts"]) + self.put(Key, data) + del self.uploads[UploadId] + return {} + + def abort_multipart_upload(self, Bucket, Key, UploadId): + self._call("abort_multipart_upload", UploadId=UploadId) + self.uploads.pop(UploadId, None) + + def ops(self, op): + return [kwargs for name, kwargs in self.calls if name == op] + + +@pytest.fixture(autouse=True) +def small_sizes(monkeypatch): + monkeypatch.setattr(nvt, "_sleep", lambda _: None) + monkeypatch.setattr(nvt, "SINGLE_PUT_MAX", 10) + monkeypatch.setattr(nvt, "DEFAULT_PART_SIZE", 8) + monkeypatch.setattr(nvt, "DOWNLOAD_CHUNK_SIZE", 8) + + +def _transfer(fake, tmp_path, concurrency=1, events=None): + progress = None if events is None else (lambda e, p, n: events.append((e, p, n))) + return nvt._Transfer( + fake, "vol", "https://s3", concurrency, progress, state_dir=tmp_path / "state" + ) + + +DATA = bytes(range(40)) # 5 parts / chunks of 8 bytes + + +# ---------------------------------------------------------------- upload + + +def test_part_size_for_limits(): + assert nvt.part_size_for(10) == 8 + size = 4 * 1024**4 + part = nvt.part_size_for(size) + assert part <= nvt.MAX_PART_SIZE and size / part <= nvt.MAX_PARTS + with pytest.raises(nvt.NetworkVolumeTransferError): + nvt.part_size_for(nvt.MAX_PART_SIZE * nvt.MAX_PARTS + 1) + + +def test_small_file_uses_put(tmp_path): + fake = FakeS3() + src = tmp_path / "small.bin" + src.write_bytes(b"hello") + summary = _transfer(fake, tmp_path).upload(src, "dir/") + assert fake.objects["dir/small.bin"][0] == b"hello" + assert summary == {"files_transferred": 1, "files_skipped": 0, "bytes_transferred": 5} + + +def test_multipart_resume_uploads_only_missing_and_bad_parts(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) + if op == "upload_part" and kw["PartNumber"] == 5 + else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(list((tmp_path / "state").glob("*.json"))) == 1 + upload = next(iter(fake.uploads.values())) + assert sorted(upload["parts"]) == [1, 2, 3, 4] + upload["parts"][2] = b"X" * 8 # same size, wrong content + + fake.fail = None + fake.calls.clear() + events = [] + summary = _transfer(fake, tmp_path, events=events).upload(src, "big.bin") + + assert sorted(kw["PartNumber"] for kw in fake.ops("upload_part")) == [2, 5] + assert not fake.ops("create_multipart_upload") + assert fake.objects["big.bin"][0] == DATA + assert summary["bytes_transferred"] == len(DATA) + assert not list((tmp_path / "state").glob("*.json")) + + +def test_multipart_resume_accepts_non_md5_etag_on_size(tmp_path): + fake = FakeS3(md5_etags=False) + src = tmp_path / "big.bin" + src.write_bytes(DATA) + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) + if op == "upload_part" and kw["PartNumber"] == 5 + else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + fake.fail = None + fake.calls.clear() + _transfer(fake, tmp_path).upload(src, "big.bin") + assert sorted(kw["PartNumber"] for kw in fake.ops("upload_part")) == [5] + assert fake.objects["big.bin"][0] == DATA + + +def test_changed_file_aborts_stale_upload(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) if op == "upload_part" and kw["PartNumber"] == 2 else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + fake.fail = None + src.write_bytes(DATA[::-1]) + os.utime(src, ns=(1, 1)) + fake.calls.clear() + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(fake.ops("abort_multipart_upload")) == 1 + assert len(fake.ops("upload_part")) == 5 + assert fake.objects["big.bin"][0] == DATA[::-1] + + +def test_missing_server_upload_starts_fresh(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) if op == "upload_part" and kw["PartNumber"] == 2 else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + fake.fail = None + fake.uploads.clear() # server-side cleanup + fake.calls.clear() + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(fake.ops("create_multipart_upload")) == 1 + assert fake.objects["big.bin"][0] == DATA + + +def test_complete_timeout_then_head_counts_as_success(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + original = fake.complete_multipart_upload + + def slow_complete(**kwargs): + if kwargs["UploadId"] in fake.uploads: + original(**kwargs) # the server finishes, the client times out + raise ReadTimeoutError(endpoint_url="https://s3") + raise _client_error("NoSuchUpload", 404) + + fake.complete_multipart_upload = slow_complete + summary = _transfer(fake, tmp_path).upload(src, "big.bin") + assert fake.objects["big.bin"][0] == DATA + assert summary["files_transferred"] == 1 + assert not list((tmp_path / "state").glob("*.json")) + + +def test_complete_failure_without_object_keeps_state(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) if op == "complete_multipart_upload" else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(list((tmp_path / "state").glob("*.json"))) == 1 + + +def test_transient_errors_are_retried(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + failed = set() + + def flaky(op, kw): + if op == "upload_part" and kw["PartNumber"] not in failed: + failed.add(kw["PartNumber"]) + return _client_error("524", 524) + return None + + fake.fail = flaky + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(fake.ops("upload_part")) == 10 + assert fake.objects["big.bin"][0] == DATA + + +def test_retry_gives_up_after_attempts(): + calls = [] + + def always_503(): + calls.append(1) + raise _client_error("ServiceUnavailable", 503) + + with pytest.raises(ClientError): + nvt._with_retries(always_503, attempts=3) + assert len(calls) == 3 + + +def test_folder_upload_skips_unchanged(tmp_path): + fake = FakeS3() + root = tmp_path / "data" + (root / "b").mkdir(parents=True) + (root / "a.txt").write_bytes(b"aaa") + (root / "b" / "c.txt").write_bytes(b"cccc") + (root / "d.txt").write_bytes(b"dd") + future = datetime.now(timezone.utc) + timedelta(hours=1) + past = datetime.now(timezone.utc) - timedelta(days=1) + fake.put("models/a.txt", b"xyz", future) # same size, newer: skip + fake.put("models/d.txt", b"zz", past) # same size, older: upload + + summary = _transfer(fake, tmp_path, concurrency=2).upload(root, "models") + assert sorted(kw["Key"] for kw in fake.ops("put_object")) == ["models/b/c.txt", "models/d.txt"] + assert fake.objects["models/a.txt"][0] == b"xyz" + assert summary["files_skipped"] == 1 and summary["files_transferred"] == 2 + + +def test_folder_upload_default_prefix_and_failure_report(tmp_path): + fake = FakeS3() + root = tmp_path / "data" + root.mkdir() + (root / "ok.txt").write_bytes(b"1") + (root / "bad.txt").write_bytes(b"2") + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) if kw.get("Key") == "data/bad.txt" else None + ) + with pytest.raises(nvt.NetworkVolumeTransferError, match="1 file"): + _transfer(fake, tmp_path).upload(root, None) + assert "data/ok.txt" in fake.objects + + +def test_single_file_upload_skips_unchanged(tmp_path): + fake = FakeS3() + src = tmp_path / "f.txt" + src.write_bytes(b"abc") + fake.put("f.txt", b"abc", datetime.now(timezone.utc) + timedelta(hours=1)) + summary = _transfer(fake, tmp_path).upload(src, None) + assert summary["files_skipped"] == 1 + assert not fake.ops("put_object") + + +def test_upload_missing_local_path(tmp_path): + with pytest.raises(FileNotFoundError): + _transfer(FakeS3(), tmp_path).upload(tmp_path / "nope", None) + + +# -------------------------------------------------------------- download + + +def test_download_resumes_from_sidecar(tmp_path): + fake = FakeS3() + fake.put("big.bin", DATA) + dest = tmp_path / "out" / "big.bin" + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) + if op == "get_object" and kw["Range"] == "bytes=32-39" + else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).download("big.bin", dest) + sidecar = json.loads((tmp_path / "out" / "big.bin.rpdownload.json").read_text()) + assert sidecar["done"] == [0, 1, 2, 3] + assert not dest.exists() + + fake.fail = None + fake.calls.clear() + summary = _transfer(fake, tmp_path).download("big.bin", dest) + assert [kw["Range"] for kw in fake.ops("get_object")] == ["bytes=32-39"] + assert dest.read_bytes() == DATA + assert not (tmp_path / "out" / "big.bin.rpdownload.json").exists() + assert not (tmp_path / "out" / "big.bin.rpdownload").exists() + assert summary["bytes_transferred"] == len(DATA) + + +def test_download_restarts_when_remote_changed(tmp_path): + fake = FakeS3() + fake.put("big.bin", DATA) + dest = tmp_path / "big.bin" + fake.fail = lambda op, kw: ( + _client_error("AccessDenied", 403) if kw.get("Range") == "bytes=16-23" else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).download("big.bin", dest) + fake.fail = None + fake.put("big.bin", DATA[::-1]) + fake.calls.clear() + _transfer(fake, tmp_path).download("big.bin", dest) + assert len(fake.ops("get_object")) == 5 + assert dest.read_bytes() == DATA[::-1] + + +def test_download_short_read_is_retried(tmp_path): + fake = FakeS3() + fake.put("f.bin", DATA) + original = fake.get_object + shorted = [] + + def short_once(**kwargs): + response = original(**kwargs) + if not shorted: + shorted.append(1) + return {"Body": io.BytesIO(response["Body"].read()[:-1])} + return response + + fake.get_object = short_once + _transfer(fake, tmp_path).download("f.bin", tmp_path / "f.bin") + assert (tmp_path / "f.bin").read_bytes() == DATA + + +def test_download_file_into_existing_dir_and_skip(tmp_path): + fake = FakeS3() + fake.put("models/w.bin", b"weights") + fake.put("empty.bin", b"") + _transfer(fake, tmp_path).download("models/w.bin", tmp_path) + assert (tmp_path / "w.bin").read_bytes() == b"weights" + summary = _transfer(fake, tmp_path).download("/models/w.bin", tmp_path) + assert summary["files_skipped"] == 1 + _transfer(fake, tmp_path).download("empty.bin", tmp_path / "e.bin") + assert (tmp_path / "e.bin").read_bytes() == b"" + + +def test_folder_download_filters_and_defaults(tmp_path, monkeypatch): + fake = FakeS3() + fake.put("models/a.bin", b"a") + fake.put("models/sub/", b"") + fake.put("models/sub/b.bin", b"bb") + fake.put("models/.s3compat_uploads/x", b"junk") + fake.put("other/c.bin", b"c") + monkeypatch.chdir(tmp_path) + summary = _transfer(fake, tmp_path, concurrency=2).download("models", None) + assert (tmp_path / "models" / "a.bin").read_bytes() == b"a" + assert (tmp_path / "models" / "sub" / "b.bin").read_bytes() == b"bb" + assert not (tmp_path / "models" / ".s3compat_uploads").exists() + assert summary["files_transferred"] == 2 + + +def test_folder_download_rejects_path_traversal(tmp_path): + fake = FakeS3() + fake.put("models/ok.bin", b"a") + fake.put("models/../evil.bin", b"x") + with pytest.raises(nvt.NetworkVolumeTransferError, match="unsafe"): + _transfer(fake, tmp_path).download("models/", tmp_path / "dl") + assert not (tmp_path / "evil.bin").exists() + assert not fake.ops("get_object") + + +@pytest.mark.parametrize("rel", ["/abs", "C:/x", "a\\..\\b", ""]) +def test_safe_target_rejects(tmp_path, rel): + with pytest.raises(nvt.NetworkVolumeTransferError): + nvt._safe_target(tmp_path, rel) + + +def test_download_missing_remote(tmp_path): + with pytest.raises(FileNotFoundError): + _transfer(FakeS3(), tmp_path).download("nope", tmp_path / "x") + + +# ---------------------------------------------------------------- client + + +def test_make_transfer_requires_s3_keys(monkeypatch): + monkeypatch.delenv(nvt.ACCESS_KEY_ENV, raising=False) + monkeypatch.delenv(nvt.SECRET_KEY_ENV, raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "aws") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "aws") + with pytest.raises(nvt.NetworkVolumeTransferError, match="RUNPOD_S3_ACCESS_KEY"): + nvt.upload("vol", ".", datacenter="EU-RO-1") + + +def test_make_transfer_builds_datacenter_client(monkeypatch): + monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") + monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") + with patch( + "runpod.api.rest.run_rest_request", return_value={"dataCenter": "EU-RO-1"} + ) as rest: + transfer = nvt._make_transfer("vol123", None, None, None, None, 4, None) + rest.assert_called_once_with("GET", "/v2/network-volumes/vol123") + assert transfer.client.meta.endpoint_url == "https://s3api-eu-ro-1.runpod.io/" + assert transfer.client.meta.region_name == "EU-RO-1" + assert transfer.client.meta.config.request_checksum_calculation == "when_required" + credentials = transfer.client._request_signer._credentials + assert credentials.access_key == "user_x" + assert transfer.bucket == "vol123" + + +def test_lookup_datacenter_missing(): + with patch("runpod.api.rest.run_rest_request", return_value={}): + with pytest.raises(nvt.NetworkVolumeTransferError): + nvt._lookup_datacenter("vol") + + +def test_public_functions_delegate(tmp_path): + fake_transfer = type( + "T", (), {"upload": lambda self, l, r: ("up", l, r), "download": lambda self, r, l: ("down", r, l)} + )() + with patch.object(nvt, "_make_transfer", return_value=fake_transfer): + assert nvt.upload("vol", "x", "y")[0] == "up" + assert nvt.download("vol", "y", "x")[0] == "down" + + +def test_lazy_exports(): + import runpod + + assert runpod.upload_to_network_volume is nvt.upload + assert runpod.download_from_network_volume is nvt.download + + +# ------------------------------------------------------------------- cli + + +def test_cli_upload_prints_summary(tmp_path): + from runpod.rp_cli.main import cli + + src = tmp_path / "f.txt" + src.write_text("x") + + def fake_upload(volume_id, source, target, **kwargs): + progress = kwargs["progress"] + progress("start", source, 1) + progress("advance", source, 1) + progress("done", source, 0) + assert kwargs["datacenter"] == "EU-RO-1" and kwargs["concurrency"] == 2 + return {"files_transferred": 1, "files_skipped": 0, "bytes_transferred": 1} + + with patch.object(nvt, "upload", side_effect=fake_upload): + result = CliRunner().invoke( + cli, + ["network-volume", "upload", "vol", str(src), "--datacenter", "EU-RO-1", "--concurrency", "2"], + ) + assert result.exit_code == 0, result.output + assert "uploaded 1 file(s), 1 bytes; skipped 0 unchanged" in result.output + + +def test_cli_download_failure_says_resume(): + from runpod.rp_cli.main import cli + + with patch.object(nvt, "download", side_effect=RuntimeError("boom")): + result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b"]) + assert result.exit_code != 0 + assert "boom" in result.output + assert "re-run the same command to resume" in result.output From 86c9ddf9c929ca421d47ebd7753a5dff19e6d0dd Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 16:27:00 -0400 Subject: [PATCH 02/23] fix: directory HEAD 400 and resume hint only for transient failures Runpod answers HeadObject on a directory with 400, so download falls through to listing the prefix. The CLI only suggests re-running when the failure is interrupted or transient, not for missing paths or 401/403. --- runpod/network_volume_transfer.py | 56 +++++++++++++++++--- runpod/rp_cli/main.py | 5 +- tests/test_network_volume_transfer.py | 75 ++++++++++++++++++++++++++- 3 files changed, 126 insertions(+), 10 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index d6b435b70..a6635fc7d 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -44,7 +44,11 @@ class NetworkVolumeTransferError(RuntimeError): - """a transfer failed. re-running the same call resumes it.""" + """a transfer failed. `resumable` says whether re-running can help.""" + + def __init__(self, message: str, resumable: bool = False): + super().__init__(message) + self.resumable = resumable class _Retryable(Exception): @@ -113,6 +117,32 @@ def _is_retryable(exc: Exception) -> bool: return False +def is_resumable(exc: BaseException) -> bool: + """whether re-running the same transfer can get past `exc`. + + missing paths, bad credentials, and refused requests fail the same way + on every run. interrupted and transient failures resume. + """ + if isinstance(exc, NetworkVolumeTransferError): + return exc.resumable + if isinstance(exc, (FileNotFoundError, PermissionError, ValueError)): + return False + from runpod import error # pylint: disable=import-outside-toplevel + + if isinstance(exc, (error.AuthenticationError, error.QueryError)): + return False + code, status = _error_code(exc) + if status in (400, 401, 403, 404) or code in ( + "AccessDenied", + "InvalidAccessKeyId", + "SignatureDoesNotMatch", + "NoSuchBucket", + "NoSuchKey", + ): + return False + return True + + def _with_retries(func: Callable[[], Any], attempts: int = RETRY_ATTEMPTS) -> Any: """call func, retrying transient failures with exponential backoff.""" for attempt in range(attempts): @@ -237,13 +267,18 @@ def _state_path(self, key: str, local: Path) -> Path: # -- remote metadata - def head(self, key: str) -> Optional[Dict[str, Any]]: + def head(self, key: str, directory_ok: bool = False) -> Optional[Dict[str, Any]]: + """object metadata, or None when there is no file at `key`. + + with `directory_ok`, a 400 also means "not a file": runpod answers + HeadObject on a directory with 400 rather than 404. + """ try: return _with_retries( lambda: self.client.head_object(Bucket=self.bucket, Key=key) ) except Exception as exc: # noqa: BLE001 - if _is_not_found(exc): + if _is_not_found(exc) or (directory_ok and _error_code(exc)[1] == 400): return None raise @@ -455,6 +490,7 @@ def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: def _run_files(self, tasks: List[Callable[[], Any]], jobs: List[Tuple[Any, Any]]) -> None: """run per-file tasks. every file is attempted; failures are reported together.""" failures = [] + resumable = [] def guarded(index: int) -> None: try: @@ -462,11 +498,13 @@ def guarded(index: int) -> None: except Exception as exc: # noqa: BLE001 - collected and re-raised with self._lock: failures.append(f"{jobs[index][0]}: {exc}") + resumable.append(is_resumable(exc)) _run_pool([lambda i=i: guarded(i) for i in range(len(tasks))], self.concurrency) if failures: raise NetworkVolumeTransferError( - f"{len(failures)} file(s) failed:\n " + "\n ".join(sorted(failures)) + f"{len(failures)} file(s) failed:\n " + "\n ".join(sorted(failures)), + resumable=any(resumable), ) # -- download @@ -482,6 +520,8 @@ def download_file(self, key: str, size: int, etag: str, dest: Path) -> None: sidecar = dest.with_name(dest.name + ".rpdownload.json") state = _read_json(sidecar) done: set = set() + # the etag is only compared for equality. a multipart object's etag + # is "-", never the md5 of the file. if ( state and state.get("size") == size @@ -542,14 +582,18 @@ def get(): _run_pool([lambda i=i: fetch(i) for i in todo], self.concurrency) if tmp.stat().st_size != size: - raise NetworkVolumeTransferError(f"{dest}: size mismatch after download") + raise NetworkVolumeTransferError( + f"{dest}: size mismatch after download", resumable=True + ) os.replace(tmp, dest) _remove(sidecar) self._emit("done", label) def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, int]: key = remote_path.lstrip("/") - head = None if key == "" or key.endswith("/") else self.head(key) + head = ( + None if key == "" or key.endswith("/") else self.head(key, directory_ok=True) + ) if head is not None: name = PurePosixPath(key).name dest = Path(local_path) if local_path is not None else Path(name) diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index ac92a42c6..f9b662c2b 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1164,9 +1164,8 @@ def progress(event, path, nbytes): progress=progress, ) except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly - raise click.ClickException( - f"{exc}\nre-run the same command to resume" - ) from exc + hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" + raise click.ClickException(f"{exc}{hint}") from exc click.echo( f"{direction} {summary['files_transferred']} file(s), " diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 2efc76c4a..cc5dd73f5 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -454,6 +454,69 @@ def test_folder_download_filters_and_defaults(tmp_path, monkeypatch): assert summary["files_transferred"] == 2 +def test_folder_download_when_head_on_directory_is_400(tmp_path): + fake = FakeS3() + fake.put("models/a.bin", b"a") + original = fake.head_object + + def head(**kwargs): + if kwargs["Key"] == "models": + raise _client_error("BadRequest", 400, "HeadObject") + return original(**kwargs) + + fake.head_object = head + summary = _transfer(fake, tmp_path).download("models", tmp_path / "dl") + assert (tmp_path / "dl" / "a.bin").read_bytes() == b"a" + assert summary["files_transferred"] == 1 + + +def test_head_400_outside_download_is_an_error(tmp_path): + fake = FakeS3() + fake.fail = lambda op, kw: _client_error("BadRequest", 400) if op == "head_object" else None + src = tmp_path / "f.txt" + src.write_bytes(b"x") + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "f.txt") + + +@pytest.mark.parametrize( + "exc, expected", + [ + (FileNotFoundError("x"), False), + (_client_error("AccessDenied", 403), False), + (_client_error("InvalidAccessKeyId", 401), False), + (_client_error("NoSuchKey", 404), False), + (nvt.NetworkVolumeTransferError("keys required"), False), + (nvt.NetworkVolumeTransferError("cut", resumable=True), True), + (_client_error("524", 524), True), + (ReadTimeoutError(endpoint_url="https://s3"), True), + ], +) +def test_is_resumable(exc, expected): + assert nvt.is_resumable(exc) is expected + + +def test_is_resumable_runpod_api_errors(): + from runpod import error + + assert not nvt.is_resumable(error.AuthenticationError("bad key")) + + +def test_folder_failure_resumable_only_if_a_file_was(tmp_path): + fake = FakeS3() + root = tmp_path / "data" + root.mkdir() + (root / "a.txt").write_bytes(b"1") + fake.fail = lambda op, kw: _client_error("AccessDenied", 403) if op == "put_object" else None + with pytest.raises(nvt.NetworkVolumeTransferError) as denied: + _transfer(fake, tmp_path).upload(root, None) + assert denied.value.resumable is False + fake.fail = lambda op, kw: ConnectionResetError("reset") if op == "put_object" else None + with pytest.raises(nvt.NetworkVolumeTransferError) as reset: + _transfer(fake, tmp_path).upload(root, None) + assert reset.value.resumable is True + + def test_folder_download_rejects_path_traversal(tmp_path): fake = FakeS3() fake.put("models/ok.bin", b"a") @@ -554,8 +617,18 @@ def fake_upload(volume_id, source, target, **kwargs): def test_cli_download_failure_says_resume(): from runpod.rp_cli.main import cli - with patch.object(nvt, "download", side_effect=RuntimeError("boom")): + with patch.object(nvt, "download", side_effect=ConnectionResetError("boom")): result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b"]) assert result.exit_code != 0 assert "boom" in result.output assert "re-run the same command to resume" in result.output + + +def test_cli_missing_path_has_no_resume_hint(): + from runpod.rp_cli.main import cli + + with patch.object(nvt, "download", side_effect=FileNotFoundError("no such file")): + result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b"]) + assert result.exit_code != 0 + assert "no such file" in result.output + assert "resume" not in result.output From 25113dd90e96ead2807c94897f242b7f1a0dbfdd Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 16:37:15 -0400 Subject: [PATCH 03/23] fix: stop a download whose remote changes midway, stamp remote mtime on downloads --- runpod/network_volume_transfer.py | 22 ++++++++++++++--- tests/test_network_volume_transfer.py | 34 ++++++++++++++++++++++++++- 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index a6635fc7d..b0212d025 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -509,7 +509,9 @@ def guarded(index: int) -> None: # -- download - def download_file(self, key: str, size: int, etag: str, dest: Path) -> None: + def download_file( + self, key: str, size: int, etag: str, dest: Path, modified: Any = None + ) -> None: label = str(dest) if dest.is_file() and dest.stat().st_size == size: self._emit("skip", label, size) @@ -554,6 +556,13 @@ def get(): response = self.client.get_object( Bucket=self.bucket, Key=key, Range=f"bytes={start}-{end}" ) + got = _strip_etag(response.get("ETag")) + if etag and got and got != etag: + raise NetworkVolumeTransferError( + f"{key} changed on the volume during download; " + "run the command again to restart it", + resumable=True, + ) data = response["Body"].read() if len(data) != end - start + 1: raise _Retryable(f"short read for {key} bytes {start}-{end}") @@ -586,6 +595,9 @@ def get(): f"{dest}: size mismatch after download", resumable=True ) os.replace(tmp, dest) + if modified is not None: + # stamp the remote time so uploading the file back counts as unchanged + os.utime(dest, (time.time(), modified.timestamp())) _remove(sidecar) self._emit("done", label) @@ -600,7 +612,11 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in if dest.is_dir(): dest = dest / name self.download_file( - key, int(head["ContentLength"]), _strip_etag(head.get("ETag")), dest + key, + int(head["ContentLength"]), + _strip_etag(head.get("ETag")), + dest, + head.get("LastModified"), ) return self.summary @@ -621,7 +637,7 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in self._run_files( [ lambda d=d, o=o: self.download_file( - o["Key"], int(o["Size"]), _strip_etag(o.get("ETag")), d + o["Key"], int(o["Size"]), _strip_etag(o.get("ETag")), d, o.get("LastModified") ) for d, o in jobs ], diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index cc5dd73f5..587313938 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -64,7 +64,7 @@ def get_object(self, Bucket, Key, Range): self._call("get_object", Key=Key, Range=Range) data, _ = self.objects[Key] start, end = (int(x) for x in Range.split("=")[1].split("-")) - return {"Body": io.BytesIO(data[start : end + 1])} + return {"Body": io.BytesIO(data[start : end + 1]), "ETag": self._etag(data)} def list_objects_v2(self, Bucket, Prefix, ContinuationToken=None): self._call("list_objects_v2", Prefix=Prefix, ContinuationToken=ContinuationToken) @@ -409,6 +409,38 @@ def test_download_restarts_when_remote_changed(tmp_path): assert dest.read_bytes() == DATA[::-1] +def test_download_stops_when_remote_replaced_midway(tmp_path): + fake = FakeS3() + fake.put("big.bin", DATA) + dest = tmp_path / "big.bin" + original = fake.get_object + + def replace_after_first(**kwargs): + response = original(**kwargs) + fake.put("big.bin", DATA[::-1]) + return response + + fake.get_object = replace_after_first + with pytest.raises(nvt.NetworkVolumeTransferError, match="changed on the volume") as info: + _transfer(fake, tmp_path).download("big.bin", dest) + assert info.value.resumable + assert not dest.exists() + + +def test_downloaded_folder_uploads_back_as_unchanged(tmp_path): + fake = FakeS3() + old = datetime(2024, 1, 1, tzinfo=timezone.utc) + fake.put("data/a.bin", DATA, modified=old) + fake.put("data/sub/b.txt", b"b", modified=old) + local = tmp_path / "data" + _transfer(fake, tmp_path).download("data", local) + assert (local / "a.bin").stat().st_mtime == old.timestamp() + fake.calls.clear() + summary = _transfer(fake, tmp_path).upload(local, "data") + assert summary["files_skipped"] == 2 + assert not fake.ops("put_object") and not fake.ops("upload_part") + + def test_download_short_read_is_retried(tmp_path): fake = FakeS3() fake.put("f.bin", DATA) From 67b9f2a21589f0fe89fce8160898b44365e8e79f Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 16:56:45 -0400 Subject: [PATCH 04/23] fix: harden network volume transfer resume and completion checks - a failed complete is only accepted when it was ambiguous and the object's etag equals the multipart etag of our parts - parts are reused only when their md5 etag matches the local bytes - abort when the local file changes mid-upload or state can't be saved - download skips need same size and a remote copy no newer than local - sidecar pinned to its remote key; refuse x + x.rpdownload collisions - reject '..' remote paths - requests share one concurrency limit across files and parts --- runpod/network_volume_transfer.py | 151 +++++++++++++++++++----- tests/test_network_volume_transfer.py | 160 ++++++++++++++++++++++++-- 2 files changed, 274 insertions(+), 37 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index b0212d025..86f91ffa8 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -205,6 +205,19 @@ def _is_md5_etag(etag: str) -> bool: return len(etag) == 32 and all(c in "0123456789abcdef" for c in etag.lower()) +def _multipart_etag(parts: List[Dict[str, Any]]) -> str: + """the etag s3 (and runpod, verified) gives an object assembled from + parts: md5 of the parts' binary md5s, then "-". "" when a part + etag is not an md5.""" + digest = hashlib.md5() + for part in parts: + etag = _strip_etag(part["ETag"]) + if not _is_md5_etag(etag): + return "" + digest.update(bytes.fromhex(etag)) + return f"{digest.hexdigest()}-{len(parts)}" + + def _run_pool(tasks: List[Callable[[], Any]], concurrency: int) -> List[Any]: """run tasks on a thread pool. the first failure cancels the rest.""" results = [] @@ -220,6 +233,12 @@ def _run_pool(tasks: List[Callable[[], Any]], concurrency: int) -> List[Any]: return results +def _check_remote_path(remote_path: Optional[str]) -> None: + """refuse remote paths with '..' segments before anything is sent.""" + if remote_path and ".." in remote_path.replace("\\", "/").split("/"): + raise ValueError(f"remote path {remote_path!r} must not contain '..'") + + def _strip_etag(etag: Any) -> str: return str(etag or "").strip('"') @@ -246,10 +265,17 @@ def __init__( self._progress = progress self.state_dir = state_dir or _state_dir() self._lock = threading.Lock() + # bounds requests in flight across all files: folder transfers run + # files in parallel and each large file runs its parts in parallel + self._slots = threading.BoundedSemaphore(self.concurrency) self.summary = {"files_transferred": 0, "files_skipped": 0, "bytes_transferred": 0} # -- bookkeeping + def _request(self, func: Callable[[], Any]) -> Any: + with self._slots: + return func() + def _emit(self, event: str, path: str, nbytes: int = 0) -> None: with self._lock: if event == "advance": @@ -317,9 +343,12 @@ def upload_file(self, local: Path, key: str) -> None: label = str(local) self._emit("start", label, size) if size <= SINGLE_PUT_MAX: - data = local.read_bytes() _with_retries( - lambda: self.client.put_object(Bucket=self.bucket, Key=key, Body=data) + lambda: self._request( + lambda: self.client.put_object( + Bucket=self.bucket, Key=key, Body=local.read_bytes() + ) + ) ) self._emit("advance", label, size) else: @@ -345,8 +374,9 @@ def _resume_parts( expected = min(part_size, size - offset) if expected <= 0 or int(part.get("Size", -1)) != expected: continue + # a part whose content cannot be verified is uploaded again etag = _strip_etag(part.get("ETag")) - if _is_md5_etag(etag) and _md5_hex(local, offset, expected) != etag.lower(): + if not _is_md5_etag(etag) or _md5_hex(local, offset, expected) != etag.lower(): continue done[number] = part.get("ETag") return done @@ -374,12 +404,7 @@ def _multipart_upload( except LookupError: pass else: - try: - self.client.abort_multipart_upload( - Bucket=self.bucket, Key=key, UploadId=state["upload_id"] - ) - except Exception: # noqa: BLE001 - best effort - pass + self._abort(key, state["upload_id"]) if upload_id is None: _remove(state_path) @@ -388,17 +413,22 @@ def _multipart_upload( lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=key) ) upload_id = created["UploadId"] - _write_json_atomic( - state_path, - { - "upload_id": upload_id, - "part_size": part_size, - "file_size": size, - "mtime_ns": mtime_ns, - "bucket": self.bucket, - "key": key, - }, - ) + try: + _write_json_atomic( + state_path, + { + "upload_id": upload_id, + "part_size": part_size, + "file_size": size, + "mtime_ns": mtime_ns, + "bucket": self.bucket, + "key": key, + }, + ) + except OSError: + # without saved state no re-run can find this upload again + self._abort(key, upload_id) + raise already = sum(min(part_size, size - (n - 1) * part_size) for n in done) if already: @@ -418,7 +448,7 @@ def send(): Body=body, ) - response = _with_retries(send) + response = _with_retries(lambda: self._request(send)) self._emit("advance", label, length) return number, response["ETag"] @@ -428,6 +458,17 @@ def send(): ): done[number] = etag + # the parts were read over the whole upload; if the file moved under + # them, they may mix two versions of it + now = local.stat() + if now.st_size != size or now.st_mtime_ns != mtime_ns: + self._abort(key, upload_id) + _remove(state_path) + raise NetworkVolumeTransferError( + f"{local} changed during upload; run the command again to upload the new version", + resumable=True, + ) + parts = [{"PartNumber": n, "ETag": done[n]} for n in sorted(done)] try: _with_retries( @@ -438,13 +479,34 @@ def send(): MultipartUpload={"Parts": parts}, ) ) - except Exception: - # a slow complete can time out after the server finished it - head = self.head(key) - if head is None or int(head.get("ContentLength", -1)) != size: + except Exception as exc: # noqa: BLE001 - re-raised unless it completed + if not self._completed_anyway(exc, key, size, _multipart_etag(parts)): raise _remove(state_path) + def _completed_anyway(self, exc: Exception, key: str, size: int, want_etag: str) -> bool: + """whether a failed complete took effect anyway. + + assembling a large file can outlive the proxy (524) while it still + finishes server-side, and a retry then finds the upload gone. only + those ambiguous failures qualify, and only when the object at the key + is this upload's: its etag equals the multipart etag of our parts, or, + if that cannot be computed, its size matches. + """ + if not (_is_retryable(exc) or _is_not_found(exc)): + return False + head = self.head(key) + if head is None or int(head.get("ContentLength", -1)) != size: + return False + return not want_etag or _strip_etag(head.get("ETag")) == want_etag + + def _abort(self, key: str, upload_id: str) -> None: + """discard an upload that can no longer be resumed (best effort).""" + try: + self.client.abort_multipart_upload(Bucket=self.bucket, Key=key, UploadId=upload_id) + except Exception: # noqa: BLE001 - leaves orphaned parts at worst + pass + def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bool: if not remote: return False @@ -456,6 +518,7 @@ def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bo return modified.timestamp() >= stat.st_mtime def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: + _check_remote_path(remote_path) if not local.exists(): raise FileNotFoundError(str(local)) if local.is_dir(): @@ -513,7 +576,13 @@ def download_file( self, key: str, size: int, etag: str, dest: Path, modified: Any = None ) -> None: label = str(dest) - if dest.is_file() and dest.stat().st_size == size: + # downloads stamp the remote time, so an unchanged file is the same + # size and not older than the remote copy + if ( + dest.is_file() + and dest.stat().st_size == size + and (modified is None or modified.timestamp() <= dest.stat().st_mtime) + ): self._emit("skip", label, size) return self._emit("start", label, size) @@ -526,6 +595,7 @@ def download_file( # is "-", never the md5 of the file. if ( state + and state.get("key") == key and state.get("size") == size and state.get("etag") == etag and state.get("chunk_size") == DOWNLOAD_CHUNK_SIZE @@ -538,7 +608,13 @@ def download_file( handle.truncate(size) _write_json_atomic( sidecar, - {"size": size, "etag": etag, "chunk_size": DOWNLOAD_CHUNK_SIZE, "done": []}, + { + "key": key, + "size": size, + "etag": etag, + "chunk_size": DOWNLOAD_CHUNK_SIZE, + "done": [], + }, ) count = math.ceil(size / DOWNLOAD_CHUNK_SIZE) @@ -568,7 +644,7 @@ def get(): raise _Retryable(f"short read for {key} bytes {start}-{end}") return data - data = _with_retries(get) + data = _with_retries(lambda: self._request(get)) with open(tmp, "r+b") as handle: handle.seek(start) handle.write(data) @@ -579,6 +655,7 @@ def get(): _write_json_atomic( sidecar, { + "key": key, "size": size, "etag": etag, "chunk_size": DOWNLOAD_CHUNK_SIZE, @@ -602,6 +679,7 @@ def get(): self._emit("done", label) def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, int]: + _check_remote_path(remote_path) key = remote_path.lstrip("/") head = ( None if key == "" or key.endswith("/") else self.head(key, directory_ok=True) @@ -634,6 +712,16 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in PurePosixPath(prefix).name or self.bucket ) jobs = [(_safe_target(root, obj["Key"][len(list_prefix):]), obj) for obj in objects] + # a remote "x.rpdownload" would share its partial file with "x" + targets = {dest for dest, _ in jobs} + for dest, _ in jobs: + for suffix in (".rpdownload", ".rpdownload.json"): + twin = dest.with_name(dest.name + suffix) + if twin in targets: + raise NetworkVolumeTransferError( + f"cannot download both {dest} and {twin}: " + "the second name is used for the first's partial download" + ) self._run_files( [ lambda d=d, o=o: self.download_file( @@ -708,7 +796,9 @@ def _make_transfer( retries={"max_attempts": 10, "mode": "standard"}, connect_timeout=60, read_timeout=3600, - max_pool_connections=max(10, concurrency * concurrency), + # transfer requests share `concurrency` slots; metadata calls from + # the file workers can run alongside them + max_pool_connections=max(10, 2 * concurrency), ) client = boto3.session.Session().client( "s3", @@ -768,7 +858,8 @@ def download( a `remote_path` ending in "/" (or one that only exists as a prefix) is a folder. `local_path` defaults to the remote basename. local files of - the same size are skipped, and partial downloads resume on re-run. + the same size and not older than the remote copy are skipped (downloads + stamp the remote time on the file), and partial downloads resume on re-run. returns {files_transferred, files_skipped, bytes_transferred}. """ diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 587313938..dacd76964 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -25,6 +25,7 @@ class FakeS3: def __init__(self, md5_etags=True): self.objects = {} # key -> (data, last_modified) + self.multipart_etags = {} # key -> etag of an object assembled from parts self.uploads = {} # upload_id -> {"key": str, "parts": {n: bytes}} self.calls = [] self.fail = None # fail(op, kwargs) -> exception or None @@ -45,6 +46,10 @@ def _call(self, op, **kwargs): def put(self, key, data, modified=None): self.objects[key] = (data, modified or datetime.now(timezone.utc)) + self.multipart_etags.pop(key, None) + + def object_etag(self, key): + return self.multipart_etags.get(key) or self._etag(self.objects[key][0]) # -- objects @@ -58,13 +63,13 @@ def head_object(self, Bucket, Key): if Key not in self.objects: raise _client_error("404", 404, "HeadObject") data, modified = self.objects[Key] - return {"ContentLength": len(data), "ETag": self._etag(data), "LastModified": modified} + return {"ContentLength": len(data), "ETag": self.object_etag(Key), "LastModified": modified} def get_object(self, Bucket, Key, Range): self._call("get_object", Key=Key, Range=Range) data, _ = self.objects[Key] start, end = (int(x) for x in Range.split("=")[1].split("-")) - return {"Body": io.BytesIO(data[start : end + 1]), "ETag": self._etag(data)} + return {"Body": io.BytesIO(data[start : end + 1]), "ETag": self.object_etag(Key)} def list_objects_v2(self, Bucket, Prefix, ContinuationToken=None): self._call("list_objects_v2", Prefix=Prefix, ContinuationToken=ContinuationToken) @@ -77,7 +82,7 @@ def list_objects_v2(self, Bucket, Prefix, ContinuationToken=None): { "Key": k, "Size": len(self.objects[k][0]), - "ETag": self._etag(self.objects[k][0]), + "ETag": self.object_etag(k), "LastModified": self.objects[k][1], } for k in page @@ -121,8 +126,11 @@ def list_parts(self, Bucket, Key, UploadId, PartNumberMarker=0): def complete_multipart_upload(self, Bucket, Key, UploadId, MultipartUpload): self._call("complete_multipart_upload", UploadId=UploadId) upload = self._upload(UploadId) - data = b"".join(upload["parts"][p["PartNumber"]] for p in MultipartUpload["Parts"]) - self.put(Key, data) + chunks = [upload["parts"][p["PartNumber"]] for p in MultipartUpload["Parts"]] + self.put(Key, b"".join(chunks)) + if self.md5_etags: + digests = b"".join(hashlib.md5(c).digest() for c in chunks) + self.multipart_etags[Key] = f'"{hashlib.md5(digests).hexdigest()}-{len(chunks)}"' del self.uploads[UploadId] return {} @@ -202,7 +210,7 @@ def test_multipart_resume_uploads_only_missing_and_bad_parts(tmp_path): assert not list((tmp_path / "state").glob("*.json")) -def test_multipart_resume_accepts_non_md5_etag_on_size(tmp_path): +def test_multipart_resume_reuploads_parts_without_md5_etag(tmp_path): fake = FakeS3(md5_etags=False) src = tmp_path / "big.bin" src.write_bytes(DATA) @@ -216,7 +224,8 @@ def test_multipart_resume_accepts_non_md5_etag_on_size(tmp_path): fake.fail = None fake.calls.clear() _transfer(fake, tmp_path).upload(src, "big.bin") - assert sorted(kw["PartNumber"] for kw in fake.ops("upload_part")) == [5] + # parts whose content cannot be verified are sent again + assert sorted(kw["PartNumber"] for kw in fake.ops("upload_part")) == [1, 2, 3, 4, 5] assert fake.objects["big.bin"][0] == DATA @@ -664,3 +673,140 @@ def test_cli_missing_path_has_no_resume_hint(): assert result.exit_code != 0 assert "no such file" in result.output assert "resume" not in result.output + + +# ------------------------------------------------- review follow-ups + + +@pytest.mark.parametrize( + "error", + [_client_error("InternalError", 524), _client_error("InvalidPart", 400)], + ids=["ambiguous-524-old-object", "definite-failure"], +) +def test_failed_complete_does_not_accept_an_old_object(tmp_path, error): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + # an older object of the same size already sits at the key + fake.put("big.bin", b"o" * len(DATA), modified=datetime(2020, 1, 1, tzinfo=timezone.utc)) + fake.fail = lambda op, kw: error if op == "complete_multipart_upload" else None + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(list((tmp_path / "state").glob("*.json"))) == 1 + + +def test_local_file_changed_during_upload_aborts(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + + def rewrite(op, kw): + if op == "upload_part" and kw["PartNumber"] == 2: + src.write_bytes(b"n" * len(DATA)) + later = src.stat().st_mtime + 60 + os.utime(src, (later, later)) + return None + + fake.fail = rewrite + with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload"): + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(fake.ops("abort_multipart_upload")) == 1 + assert not fake.ops("complete_multipart_upload") + assert not list((tmp_path / "state").glob("*.json")) + + +def test_unsavable_state_aborts_the_upload(tmp_path, monkeypatch): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + + def no_disk(path, data): + raise OSError("disk full") + + monkeypatch.setattr(nvt, "_write_json_atomic", no_disk) + with pytest.raises(OSError): + _transfer(fake, tmp_path).upload(src, "big.bin") + assert len(fake.ops("abort_multipart_upload")) == 1 + + +def test_download_refetches_newer_same_size_remote(tmp_path): + fake = FakeS3() + fake.put("f.txt", b"version-1") + dest = tmp_path / "f.txt" + _transfer(fake, tmp_path).download("f.txt", dest) + fake.put("f.txt", b"version-2", modified=datetime.now(timezone.utc) + timedelta(minutes=1)) + summary = _transfer(fake, tmp_path).download("f.txt", dest) + assert summary["files_transferred"] == 1 + assert dest.read_bytes() == b"version-2" + + +def test_download_ignores_sidecar_of_another_key(tmp_path): + fake = FakeS3() + fake.put("b.bin", DATA) + dest = tmp_path / "x.bin" + (tmp_path / "x.bin.rpdownload").write_bytes(b"z" * len(DATA)) + (tmp_path / "x.bin.rpdownload.json").write_text( + json.dumps( + { + "key": "a.bin", + "size": len(DATA), + "etag": nvt._strip_etag(fake.object_etag("b.bin")), + "chunk_size": nvt.DOWNLOAD_CHUNK_SIZE, + "done": [0, 1, 2, 3, 4], + } + ) + ) + _transfer(fake, tmp_path).download("b.bin", dest) + assert dest.read_bytes() == DATA + + +def test_download_refuses_partial_name_collision(tmp_path): + fake = FakeS3() + fake.put("d/x", b"x") + fake.put("d/x.rpdownload", b"y") + with pytest.raises(nvt.NetworkVolumeTransferError, match="partial download"): + _transfer(fake, tmp_path).download("d/", tmp_path / "out") + + +@pytest.mark.parametrize("op", ["upload", "download"]) +def test_dotdot_remote_path_rejected(tmp_path, op): + fake = FakeS3() + src = tmp_path / "f.txt" + src.write_bytes(b"x") + transfer = _transfer(fake, tmp_path) + with pytest.raises(ValueError, match=r"\.\."): + if op == "upload": + transfer.upload(src, "a/../b") + else: + transfer.download("a/../b", tmp_path / "out") + assert not fake.calls + assert not nvt.is_resumable(ValueError("x")) + + +def test_requests_in_flight_never_exceed_concurrency(tmp_path): + import threading + import time as _time + + fake = FakeS3() + active = {"now": 0, "peak": 0} + lock = threading.Lock() + original = fake.upload_part + + def tracked(**kwargs): + with lock: + active["now"] += 1 + active["peak"] = max(active["peak"], active["now"]) + _time.sleep(0.01) + try: + return original(**kwargs) + finally: + with lock: + active["now"] -= 1 + + fake.upload_part = tracked + folder = tmp_path / "data" + folder.mkdir() + for i in range(4): + (folder / f"f{i}.bin").write_bytes(DATA) + _transfer(fake, tmp_path, concurrency=2).upload(folder, "data") + assert active["peak"] <= 2 From 978268e4b8c3c2638ffcdf96d77c0a39cb45fa07 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 16:57:23 -0400 Subject: [PATCH 05/23] chore: address code scanning notes in network volume transfer --- runpod/network_volume_transfer.py | 13 ++++++------- tests/test_network_volume_transfer.py | 10 ++++------ 2 files changed, 10 insertions(+), 13 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 86f91ffa8..284f028a5 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -175,10 +175,7 @@ def _read_json(path: Path) -> Optional[Dict[str, Any]]: def _remove(path: Path) -> None: - try: - path.unlink() - except FileNotFoundError: - pass + path.unlink(missing_ok=True) def _read_range(path: Path, offset: int, length: int) -> bytes: @@ -402,7 +399,8 @@ def _multipart_upload( done = self._resume_parts(local, key, state, size, part_size) upload_id = state["upload_id"] except LookupError: - pass + # the server no longer has this upload: start a new one + upload_id = None else: self._abort(key, state["upload_id"]) if upload_id is None: @@ -504,8 +502,9 @@ def _abort(self, key: str, upload_id: str) -> None: """discard an upload that can no longer be resumed (best effort).""" try: self.client.abort_multipart_upload(Bucket=self.bucket, Key=key, UploadId=upload_id) - except Exception: # noqa: BLE001 - leaves orphaned parts at worst - pass + except Exception: # noqa: BLE001 + # a failed abort only leaves orphaned parts behind + return def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bool: if not remote: diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index dacd76964..6c4f557ee 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -11,7 +11,9 @@ from botocore.exceptions import ClientError, ReadTimeoutError from click.testing import CliRunner -from runpod import network_volume_transfer as nvt +import runpod +import runpod.network_volume_transfer as nvt +import runpod.error def _client_error(code, status, op="Op"): @@ -538,9 +540,7 @@ def test_is_resumable(exc, expected): def test_is_resumable_runpod_api_errors(): - from runpod import error - - assert not nvt.is_resumable(error.AuthenticationError("bad key")) + assert not nvt.is_resumable(runpod.error.AuthenticationError("bad key")) def test_folder_failure_resumable_only_if_a_file_was(tmp_path): @@ -623,8 +623,6 @@ def test_public_functions_delegate(tmp_path): def test_lazy_exports(): - import runpod - assert runpod.upload_to_network_volume is nvt.upload assert runpod.download_from_network_volume is nvt.download From d95e161bb5ae0ae0d323c1a534052ed6adcedbf2 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 17:04:35 -0400 Subject: [PATCH 06/23] fix: close remaining network volume transfer review gaps - an ambiguous failed complete needs our multipart etag to match - a single put fails if the file changed across retries - partial-name collisions also cover directories and case folding - every s3 call takes a concurrency slot --- runpod/network_volume_transfer.py | 83 +++++++++++++++++++-------- tests/test_network_volume_transfer.py | 68 +++++++++++++++++++++- 2 files changed, 126 insertions(+), 25 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 284f028a5..fb3b1e886 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -298,7 +298,9 @@ def head(self, key: str, directory_ok: bool = False) -> Optional[Dict[str, Any]] """ try: return _with_retries( - lambda: self.client.head_object(Bucket=self.bucket, Key=key) + lambda: self._request( + lambda: self.client.head_object(Bucket=self.bucket, Key=key) + ) ) except Exception as exc: # noqa: BLE001 if _is_not_found(exc) or (directory_ok and _error_code(exc)[1] == 400): @@ -310,7 +312,9 @@ def list_prefix(self, prefix: str) -> List[Dict[str, Any]]: kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Prefix": prefix} while True: try: - page = _with_retries(lambda: self.client.list_objects_v2(**kwargs)) + page = _with_retries( + lambda: self._request(lambda: self.client.list_objects_v2(**kwargs)) + ) except Exception as exc: # noqa: BLE001 if _is_not_found(exc): return objects @@ -325,7 +329,9 @@ def _list_parts(self, key: str, upload_id: str) -> List[Dict[str, Any]]: parts: List[Dict[str, Any]] = [] kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Key": key, "UploadId": upload_id} while True: - page = _with_retries(lambda: self.client.list_parts(**kwargs)) + page = _with_retries( + lambda: self._request(lambda: self.client.list_parts(**kwargs)) + ) parts.extend(page.get("Parts") or []) marker = page.get("NextPartNumberMarker") if not page.get("IsTruncated") or not marker: @@ -347,6 +353,14 @@ def upload_file(self, local: Path, key: str) -> None: ) ) ) + # a retry rereads the file; if it changed meanwhile, the object + # may not be the version that was measured + now = local.stat() + if now.st_size != size or now.st_mtime_ns != stat.st_mtime_ns: + raise NetworkVolumeTransferError( + f"{local} changed during upload; run the command again to upload the new version", + resumable=True, + ) self._emit("advance", label, size) else: self._multipart_upload(local, key, size, stat.st_mtime_ns, label) @@ -408,7 +422,9 @@ def _multipart_upload( if upload_id is None: created = _with_retries( - lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=key) + lambda: self._request( + lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=key) + ) ) upload_id = created["UploadId"] try: @@ -470,11 +486,13 @@ def send(): parts = [{"PartNumber": n, "ETag": done[n]} for n in sorted(done)] try: _with_retries( - lambda: self.client.complete_multipart_upload( - Bucket=self.bucket, - Key=key, - UploadId=upload_id, - MultipartUpload={"Parts": parts}, + lambda: self._request( + lambda: self.client.complete_multipart_upload( + Bucket=self.bucket, + Key=key, + UploadId=upload_id, + MultipartUpload={"Parts": parts}, + ) ) ) except Exception as exc: # noqa: BLE001 - re-raised unless it completed @@ -488,20 +506,25 @@ def _completed_anyway(self, exc: Exception, key: str, size: int, want_etag: str) assembling a large file can outlive the proxy (524) while it still finishes server-side, and a retry then finds the upload gone. only those ambiguous failures qualify, and only when the object at the key - is this upload's: its etag equals the multipart etag of our parts, or, - if that cannot be computed, its size matches. + is this upload's: its size matches and its etag equals the multipart + etag of our parts. """ if not (_is_retryable(exc) or _is_not_found(exc)): return False head = self.head(key) if head is None or int(head.get("ContentLength", -1)) != size: return False - return not want_etag or _strip_etag(head.get("ETag")) == want_etag + # without our multipart etag the object's identity cannot be confirmed + return bool(want_etag) and _strip_etag(head.get("ETag")) == want_etag def _abort(self, key: str, upload_id: str) -> None: """discard an upload that can no longer be resumed (best effort).""" try: - self.client.abort_multipart_upload(Bucket=self.bucket, Key=key, UploadId=upload_id) + self._request( + lambda: self.client.abort_multipart_upload( + Bucket=self.bucket, Key=key, UploadId=upload_id + ) + ) except Exception: # noqa: BLE001 # a failed abort only leaves orphaned parts behind return @@ -711,16 +734,7 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in PurePosixPath(prefix).name or self.bucket ) jobs = [(_safe_target(root, obj["Key"][len(list_prefix):]), obj) for obj in objects] - # a remote "x.rpdownload" would share its partial file with "x" - targets = {dest for dest, _ in jobs} - for dest, _ in jobs: - for suffix in (".rpdownload", ".rpdownload.json"): - twin = dest.with_name(dest.name + suffix) - if twin in targets: - raise NetworkVolumeTransferError( - f"cannot download both {dest} and {twin}: " - "the second name is used for the first's partial download" - ) + _check_partial_names([dest for dest, _ in jobs], root) self._run_files( [ lambda d=d, o=o: self.download_file( @@ -733,6 +747,29 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in return self.summary +def _check_partial_names(targets: List[Path], root: Path) -> None: + """refuse targets where a file's partial or sidecar name is also a target, + or a directory a target lives in ("x" beside "x.rpdownload" or + "x.rpdownload/y"). names are compared case-folded, since the destination + may be case-insensitive; that can refuse a listing a case-sensitive + filesystem would take, which is the safe direction.""" + taken: Dict[str, Path] = {} + for target in targets: + rel = target.relative_to(root) + for path in (rel, *rel.parents): + if str(path) != ".": + taken[str(path).casefold()] = target + for target in targets: + rel = target.relative_to(root) + for suffix in (".rpdownload", ".rpdownload.json"): + twin = str(rel.with_name(rel.name + suffix)).casefold() + if twin in taken: + raise NetworkVolumeTransferError( + f"cannot download both {target} and {taken[twin]}: " + f"{rel.name + suffix} is used for the first's partial download" + ) + + def _safe_target(root: Path, rel: str) -> Path: """local path for a remote relative key, refusing anything outside root.""" parts = PurePosixPath(rel).parts diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 6c4f557ee..16da8793b 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -758,12 +758,16 @@ def test_download_ignores_sidecar_of_another_key(tmp_path): assert dest.read_bytes() == DATA -def test_download_refuses_partial_name_collision(tmp_path): +@pytest.mark.parametrize( + "other", ["d/x.rpdownload", "d/x.rpdownload.json", "d/x.rpdownload/y", "d/X.RPDOWNLOAD"] +) +def test_download_refuses_partial_name_collision(tmp_path, other): fake = FakeS3() fake.put("d/x", b"x") - fake.put("d/x.rpdownload", b"y") + fake.put(other, b"y") with pytest.raises(nvt.NetworkVolumeTransferError, match="partial download"): _transfer(fake, tmp_path).download("d/", tmp_path / "out") + assert not fake.ops("get_object") @pytest.mark.parametrize("op", ["upload", "download"]) @@ -808,3 +812,63 @@ def tracked(**kwargs): (folder / f"f{i}.bin").write_bytes(DATA) _transfer(fake, tmp_path, concurrency=2).upload(folder, "data") assert active["peak"] <= 2 + + +def test_ambiguous_complete_without_verifiable_etag_fails(tmp_path): + fake = FakeS3(md5_etags=False) + src = tmp_path / "big.bin" + src.write_bytes(DATA) + fake.put("big.bin", b"o" * len(DATA), modified=datetime(2020, 1, 1, tzinfo=timezone.utc)) + fake.fail = lambda op, kw: ( + _client_error("InternalError", 524) if op == "complete_multipart_upload" else None + ) + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "big.bin") + + +def test_single_put_detects_change_between_retries(tmp_path): + fake = FakeS3() + src = tmp_path / "small.txt" + src.write_bytes(b"original") + attempts = [] + + def change_then_fail(op, kw): + if op == "put_object" and not attempts: + attempts.append(1) + src.write_bytes(b"CHANGED!") + later = src.stat().st_mtime + 60 + os.utime(src, (later, later)) + return _client_error("ServiceUnavailable", 503) + return None + + fake.fail = change_then_fail + with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload"): + _transfer(fake, tmp_path).upload(src, "small.txt") + + +def test_all_requests_share_the_concurrency_limit(tmp_path): + import threading + import time as _time + + fake = FakeS3() + active = {"now": 0, "peak": 0} + lock = threading.Lock() + original_call = fake._call + + def tracked_call(op, **kwargs): + with lock: + active["now"] += 1 + active["peak"] = max(active["peak"], active["now"]) + _time.sleep(0.005) + with lock: + active["now"] -= 1 + return original_call(op, **kwargs) + + fake._call = tracked_call + folder = tmp_path / "data" + folder.mkdir() + for i in range(8): + (folder / f"f{i}.bin").write_bytes(DATA + bytes([i])) + _transfer(fake, tmp_path, concurrency=2).upload(folder, "data") + _transfer(fake, tmp_path, concurrency=2).download("data/", tmp_path / "out") + assert active["peak"] <= 2 From c92ac38a9774869539271cf66d2ea2f2860a8706 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 17:47:48 -0400 Subject: [PATCH 07/23] fix: address network volume transfer review round 1 - ctrl-c stops files in flight (shared stop event) and the cli prints the resume hint - botocore retries off; _with_retries owns them; 300s socket read timeout - refuse download targets that share a local name (casefold + nfc) and '.' key segments - resume hint only for interruptions and transient failures; a dns failure is neither retried nor resumable and points at --datacenter - uploading a file onto a remote folder explains the trailing slash; a download target ending in / is a folder (also via the public function) - a single put that raced a local edit deletes the object so the re-run sends the new version - --force; skip rule compares whole seconds; missing folder is not found; resumed bytes are not counted as transferred; leftovers not uploaded - escape volume id in the rest path; docs reference and changelog entry --- CHANGELOG.md | 1 + docs/cli/references/command_line_interface.md | 29 ++ runpod/network_volume_transfer.py | 247 +++++++++++++----- runpod/rp_cli/main.py | 24 +- tests/test_network_volume_transfer.py | 221 +++++++++++++--- 5 files changed, 410 insertions(+), 112 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index be9484adf..da416680d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Features +* `rp network-volume upload/download` and `runpod.upload_to_network_volume` / `runpod.download_from_network_volume` transfer files and folders over a network volume's S3-compatible API. Re-running an interrupted transfer resumes it: unchanged files are skipped, multipart uploads reuse MD5-verified parts, and downloads keep finished chunks. * Apps expose `NetworkVolume` and `GlobalVolume` resources with explicit path-keyed mounts, lazy create-if-missing provisioning, and execution-bound filesystem access. Global volumes use GraphQL for lookup and creation. ### Bug Fixes diff --git a/docs/cli/references/command_line_interface.md b/docs/cli/references/command_line_interface.md index b3a32bd98..496d55d88 100644 --- a/docs/cli/references/command_line_interface.md +++ b/docs/cli/references/command_line_interface.md @@ -23,6 +23,10 @@ rp ssh POD_ID # open a terminal on a pod rp pod list rp pod create rp pod connect POD_ID + +# Network volume files +rp network-volume upload VOLUME_ID ./dataset datasets/dataset +rp network-volume download VOLUME_ID outputs/ ./outputs ``` ## Overview @@ -59,3 +63,28 @@ entrypoint or a resource function; interrupt discovery if an import blocks. On Ctrl-C, a dev session waits for its running local entrypoint to finish before deleting endpoints. Blocking entrypoint code can delay shutdown. Cleanup failures exit with an error and report the endpoint IDs that still need deletion. + +### Network volume files + +```bash +rp network-volume upload VOLUME_ID ./model.safetensors models/ +rp network-volume upload VOLUME_ID ./dataset datasets/dataset +rp network-volume download VOLUME_ID models/model.safetensors +rp network-volume download VOLUME_ID outputs/ ./outputs +``` + +Transfers use the volume's S3-compatible API, so they need S3 API keys +(separate from the API key): set `RUNPOD_S3_ACCESS_KEY` and +`RUNPOD_S3_SECRET_KEY`, or pass `--s3-access-key` / `--s3-secret-key`. The +volume must be in a [datacenter with the S3 API](https://docs.runpod.io/storage/s3-api). +Its datacenter is looked up with your API key unless you pass `--datacenter`. + +To resume an interrupted transfer, run the same command again. Files that +are unchanged (same size, destination not older) are skipped, a large upload +reuses the parts it already sent, and a large download keeps its finished +chunks in `.rpdownload`. `--force` transfers every file. Parts of an +upload that is never resumed stay in the volume's hidden `.s3compat_uploads` +folder and count toward its size. + +The same transfers are available from Python as +`runpod.upload_to_network_volume` and `runpod.download_from_network_volume`. diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index fb3b1e886..1dce409ae 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -17,7 +17,9 @@ import math import os import random +import socket import threading +import unicodedata import time from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path, PurePosixPath @@ -30,14 +32,18 @@ MAX_PARTS = 10000 DOWNLOAD_CHUNK_SIZE = 64 * MIB RETRY_ATTEMPTS = 8 +# runpod's s3 api keeps in-progress multipart parts here (server-created) UPLOADS_DIR = ".s3compat_uploads" +PARTIAL_SUFFIX = ".rpdownload" +SIDECAR_SUFFIX = ".rpdownload.json" ACCESS_KEY_ENV = "RUNPOD_S3_ACCESS_KEY" SECRET_KEY_ENV = "RUNPOD_S3_SECRET_KEY" S3_KEYS_URL = "https://console.runpod.io/user/settings" # progress(event, path, nbytes). events: "start" (nbytes = file size), -# "advance" (nbytes = delta), "skip" (nbytes = file size), "done". +# "advance" (nbytes sent or received), "reuse" (nbytes an earlier run already +# moved: kept parts or chunks), "skip" (nbytes = file size), "done". Progress = Callable[[str, str, int], None] _sleep = time.sleep @@ -55,6 +61,10 @@ class _Retryable(Exception): """a transient failure raised by this module (e.g. a short read).""" +class _Stopped(Exception): + """the transfer was interrupted; raised by workers so they wind down.""" + + def _state_dir() -> Path: return Path(os.path.expanduser("~/.runpod/transfers")) @@ -97,6 +107,9 @@ def _is_retryable(exc: Exception) -> bool: if isinstance(exc, _Retryable): return True + if _is_dns_failure(exc): + # a host that does not resolve is a wrong datacenter, not a blip + return False if isinstance( exc, ( @@ -117,30 +130,36 @@ def _is_retryable(exc: Exception) -> bool: return False +def _is_dns_failure(exc: BaseException) -> bool: + """whether `exc` (or an error it wraps) is a host name that does not resolve.""" + seen = 0 + stack = [exc] + while stack and seen < 20: + err = stack.pop() + seen += 1 + if isinstance(err, socket.gaierror): + return True + kwargs = getattr(err, "kwargs", None) + wrapped = [err.__cause__, err.__context__, getattr(err, "reason", None)] + if isinstance(kwargs, dict): + wrapped.append(kwargs.get("error")) + stack.extend(w for w in wrapped if isinstance(w, BaseException)) + return False + + def is_resumable(exc: BaseException) -> bool: """whether re-running the same transfer can get past `exc`. - missing paths, bad credentials, and refused requests fail the same way - on every run. interrupted and transient failures resume. + an interruption, or a transient failure (network, timeout, 429, 5xx) + that outlasted the retries, resumes. anything else (missing paths, bad + credentials, a wrong datacenter, local disk errors, invalid input) fails + the same way on every run. """ + if isinstance(exc, KeyboardInterrupt): + return True if isinstance(exc, NetworkVolumeTransferError): return exc.resumable - if isinstance(exc, (FileNotFoundError, PermissionError, ValueError)): - return False - from runpod import error # pylint: disable=import-outside-toplevel - - if isinstance(exc, (error.AuthenticationError, error.QueryError)): - return False - code, status = _error_code(exc) - if status in (400, 401, 403, 404) or code in ( - "AccessDenied", - "InvalidAccessKeyId", - "SignatureDoesNotMatch", - "NoSuchBucket", - "NoSuchKey", - ): - return False - return True + return _is_retryable(exc) def _with_retries(func: Callable[[], Any], attempts: int = RETRY_ATTEMPTS) -> Any: @@ -215,21 +234,42 @@ def _multipart_etag(parts: List[Dict[str, Any]]) -> str: return f"{digest.hexdigest()}-{len(parts)}" -def _run_pool(tasks: List[Callable[[], Any]], concurrency: int) -> List[Any]: - """run tasks on a thread pool. the first failure cancels the rest.""" +def _run_pool( + tasks: List[Callable[[], Any]], + concurrency: int, + stop: Optional[threading.Event] = None, +) -> List[Any]: + """run tasks on a thread pool. the first failure cancels the rest. + + an interrupt (ctrl-c) also sets `stop`, before the pool waits for its + running tasks, so workers checking it wind down instead of finishing + every file in flight. + """ results = [] with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool: futures = [pool.submit(task) for task in tasks] try: for future in as_completed(futures): results.append(future.result()) - except BaseException: + except BaseException as exc: + if stop is not None and not isinstance(exc, Exception): + stop.set() for future in futures: future.cancel() raise return results +def _is_download_leftover(path: Path) -> bool: + """a partial download, its sidecar, or the sidecar's temp file.""" + return path.name.endswith((PARTIAL_SUFFIX, SIDECAR_SUFFIX, SIDECAR_SUFFIX + ".tmp")) + + +def _same_second_or_later(a: float, b: float) -> bool: + """a >= b in whole seconds; s3 dates carry no finer precision.""" + return int(a) >= int(b) + + def _check_remote_path(remote_path: Optional[str]) -> None: """refuse remote paths with '..' segments before anything is sent.""" if remote_path and ".." in remote_path.replace("\\", "/").split("/"): @@ -254,6 +294,7 @@ def __init__( concurrency: int = 4, progress: Optional[Progress] = None, state_dir: Optional[Path] = None, + force: bool = False, ): self.client = client self.bucket = bucket @@ -265,14 +306,23 @@ def __init__( # bounds requests in flight across all files: folder transfers run # files in parallel and each large file runs its parts in parallel self._slots = threading.BoundedSemaphore(self.concurrency) + # transfer every file, even ones the skip rule calls unchanged + self.force = force + # set on ctrl-c so every worker winds down + self._stop = threading.Event() self.summary = {"files_transferred": 0, "files_skipped": 0, "bytes_transferred": 0} # -- bookkeeping def _request(self, func: Callable[[], Any]) -> Any: + if self._stop.is_set(): + raise _Stopped("interrupted") with self._slots: return func() + def _pool(self, tasks: List[Callable[[], Any]]) -> List[Any]: + return _run_pool(tasks, self.concurrency, self._stop) + def _emit(self, event: str, path: str, nbytes: int = 0) -> None: with self._lock: if event == "advance": @@ -354,12 +404,21 @@ def upload_file(self, local: Path, key: str) -> None: ) ) # a retry rereads the file; if it changed meanwhile, the object - # may not be the version that was measured + # may hold either version and is newer than the local edit, so a + # re-run would skip it. remove it so the re-run sends the file. now = local.stat() if now.st_size != size or now.st_mtime_ns != stat.st_mtime_ns: + try: + _with_retries( + lambda: self._request( + lambda: self.client.delete_object(Bucket=self.bucket, Key=key) + ) + ) + except Exception: # noqa: BLE001 + # --force on the re-run still sends it + pass raise NetworkVolumeTransferError( - f"{local} changed during upload; run the command again to upload the new version", - resumable=True, + f"{local} changed during upload", resumable=True ) self._emit("advance", label, size) else: @@ -446,7 +505,7 @@ def _multipart_upload( already = sum(min(part_size, size - (n - 1) * part_size) for n in done) if already: - self._emit("advance", label, already) + self._emit("reuse", label, already) def upload_part(number: int) -> Tuple[int, str]: offset = (number - 1) * part_size @@ -467,9 +526,7 @@ def send(): return number, response["ETag"] todo = [n for n in range(1, count + 1) if n not in done] - for number, etag in _run_pool( - [lambda n=n: upload_part(n) for n in todo], self.concurrency - ): + for number, etag in self._pool([lambda n=n: upload_part(n) for n in todo]): done[number] = etag # the parts were read over the whole upload; if the file moved under @@ -537,7 +594,7 @@ def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bo modified = remote.get("LastModified") if size is None or int(size) != stat.st_size or modified is None: return False - return modified.timestamp() >= stat.st_mtime + return _same_second_or_later(modified.timestamp(), stat.st_mtime) def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: _check_remote_path(remote_path) @@ -550,10 +607,12 @@ def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: for obj in self.list_prefix(f"{prefix}/" if prefix else "") } jobs = [] - for path in sorted(p for p in local.rglob("*") if p.is_file()): + for path in sorted( + p for p in local.rglob("*") if p.is_file() and not _is_download_leftover(p) + ): rel = path.relative_to(local).as_posix() key = f"{prefix}/{rel}" if prefix else rel - if self._unchanged_remote(remote.get(key), path): + if not self.force and self._unchanged_remote(remote.get(key), path): self._emit("skip", str(path), path.stat().st_size) continue jobs.append((path, key)) @@ -566,7 +625,16 @@ def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: key = remote_path.lstrip("/") + local.name else: key = remote_path.lstrip("/") - if self._unchanged_remote(self.head(key), local): + try: + remote = self.head(key) + except Exception as exc: # noqa: BLE001 + # runpod answers a head on a directory with 400 + if _error_code(exc)[1] == 400: + raise NetworkVolumeTransferError( + f"{key!r} is a folder on the volume; use {key + '/'!r} to upload into it" + ) from exc + raise + if not self.force and self._unchanged_remote(remote, local): self._emit("skip", str(local), local.stat().st_size) return self.summary self.upload_file(local, key) @@ -585,7 +653,7 @@ def guarded(index: int) -> None: failures.append(f"{jobs[index][0]}: {exc}") resumable.append(is_resumable(exc)) - _run_pool([lambda i=i: guarded(i) for i in range(len(tasks))], self.concurrency) + self._pool([lambda i=i: guarded(i) for i in range(len(tasks))]) if failures: raise NetworkVolumeTransferError( f"{len(failures)} file(s) failed:\n " + "\n ".join(sorted(failures)), @@ -601,16 +669,20 @@ def download_file( # downloads stamp the remote time, so an unchanged file is the same # size and not older than the remote copy if ( - dest.is_file() + not self.force + and dest.is_file() and dest.stat().st_size == size - and (modified is None or modified.timestamp() <= dest.stat().st_mtime) + and ( + modified is None + or _same_second_or_later(dest.stat().st_mtime, modified.timestamp()) + ) ): self._emit("skip", label, size) return self._emit("start", label, size) dest.parent.mkdir(parents=True, exist_ok=True) - tmp = dest.with_name(dest.name + ".rpdownload") - sidecar = dest.with_name(dest.name + ".rpdownload.json") + tmp = dest.with_name(dest.name + PARTIAL_SUFFIX) + sidecar = dest.with_name(dest.name + SIDECAR_SUFFIX) state = _read_json(sidecar) done: set = set() # the etag is only compared for equality. a multipart object's etag @@ -644,7 +716,7 @@ def download_file( min(DOWNLOAD_CHUNK_SIZE, size - i * DOWNLOAD_CHUNK_SIZE) for i in done if i < count ) if already: - self._emit("advance", label, already) + self._emit("reuse", label, already) def fetch(index: int) -> None: start = index * DOWNLOAD_CHUNK_SIZE @@ -687,7 +759,7 @@ def get(): self._emit("advance", label, end - start + 1) todo = [i for i in range(count) if i not in done] - _run_pool([lambda i=i: fetch(i) for i in todo], self.concurrency) + self._pool([lambda i=i: fetch(i) for i in todo]) if tmp.stat().st_size != size: raise NetworkVolumeTransferError( @@ -700,7 +772,9 @@ def get(): _remove(sidecar) self._emit("done", label) - def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, int]: + def download(self, remote_path: str, local_path: Any = None) -> Dict[str, int]: + """`local_path` is a str or Path; a str ending in "/" is a folder to + download into (a Path cannot carry the slash).""" _check_remote_path(remote_path) key = remote_path.lstrip("/") head = ( @@ -709,7 +783,8 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in if head is not None: name = PurePosixPath(key).name dest = Path(local_path) if local_path is not None else Path(name) - if dest.is_dir(): + # a trailing slash names a folder to download into, as on upload + if dest.is_dir() or str(local_path or "").endswith(("/", os.sep)): dest = dest / name self.download_file( key, @@ -728,13 +803,14 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in if not obj["Key"].endswith("/") and UPLOADS_DIR not in PurePosixPath(obj["Key"]).parts ] - if not objects: + # a typo must not look like an empty folder; the whole volume may be empty + if not objects and prefix: raise FileNotFoundError(f"no such file or directory on the volume: {remote_path}") root = Path(local_path) if local_path is not None else Path( PurePosixPath(prefix).name or self.bucket ) jobs = [(_safe_target(root, obj["Key"][len(list_prefix):]), obj) for obj in objects] - _check_partial_names([dest for dest, _ in jobs], root) + _check_local_names([dest for dest, _ in jobs], root) self._run_files( [ lambda d=d, o=o: self.download_file( @@ -747,25 +823,42 @@ def download(self, remote_path: str, local_path: Optional[Path]) -> Dict[str, in return self.summary -def _check_partial_names(targets: List[Path], root: Path) -> None: - """refuse targets where a file's partial or sidecar name is also a target, - or a directory a target lives in ("x" beside "x.rpdownload" or - "x.rpdownload/y"). names are compared case-folded, since the destination - may be case-insensitive; that can refuse a listing a case-sensitive - filesystem would take, which is the safe direction.""" - taken: Dict[str, Path] = {} +def _fold(name: str) -> str: + """the key under which two spellings a filesystem may treat as one name + collide: nfc first (apfs compares names normalization-insensitively), + then unicode case folding.""" + return unicodedata.normalize("NFC", name).casefold() + + +def _check_local_names(targets: List[Path], root: Path) -> None: + """refuse targets that would share a path on disk: two targets that are + one name on a case-insensitive or normalizing filesystem ("Model.bin" and + "model.bin"), or a file whose partial or sidecar name is also a target or + a directory a target lives in ("x" beside "x.rpdownload" or + "x.rpdownload/y"). that can refuse a listing a case-sensitive filesystem + would take, which is the safe direction.""" + files: Dict[str, Path] = {} + dirs: Dict[str, Path] = {} for target in targets: rel = target.relative_to(root) - for path in (rel, *rel.parents): - if str(path) != ".": - taken[str(path).casefold()] = target + key = _fold(str(rel)) + if key in files: + raise NetworkVolumeTransferError( + f"cannot download both {files[key]} and {target}: " + "they are the same file on a case-insensitive filesystem" + ) + files[key] = target + for parent in rel.parents: + if str(parent) != ".": + dirs[_fold(str(parent))] = target for target in targets: rel = target.relative_to(root) - for suffix in (".rpdownload", ".rpdownload.json"): - twin = str(rel.with_name(rel.name + suffix)).casefold() - if twin in taken: + for suffix in (PARTIAL_SUFFIX, SIDECAR_SUFFIX): + twin = _fold(str(rel.with_name(rel.name + suffix))) + other = files.get(twin) or dirs.get(twin) + if other is not None: raise NetworkVolumeTransferError( - f"cannot download both {target} and {taken[twin]}: " + f"cannot download both {target} and {other}: " f"{rel.name + suffix} is used for the first's partial download" ) @@ -777,6 +870,8 @@ def _safe_target(root: Path, rel: str) -> Path: not rel or rel.startswith("/") or "\\" in rel + # PurePosixPath drops "." and empty segments; "a/./b" would alias "a/b" + or any(seg in ("..", ".", "") for seg in rel.split("/")) or any(part in ("..", "") for part in parts) or (parts and ":" in parts[0]) ): @@ -794,7 +889,9 @@ def _safe_target(root: Path, rel: str) -> Path: def _lookup_datacenter(volume_id: str) -> str: from runpod.api.rest import run_rest_request # pylint: disable=import-outside-toplevel - data = run_rest_request("GET", f"/v2/network-volumes/{volume_id}") or {} + from runpod.api.ctl_commands import _path_segment # pylint: disable=import-outside-toplevel + + data = run_rest_request("GET", f"/v2/network-volumes/{_path_segment(volume_id)}") or {} datacenter = data.get("dataCenter") or data.get("dataCenterId") if not datacenter: raise NetworkVolumeTransferError( @@ -811,6 +908,7 @@ def _make_transfer( endpoint_url: Optional[str], concurrency: int, progress: Optional[Progress], + force: bool = False, ) -> _Transfer: import boto3 # pylint: disable=import-outside-toplevel from botocore.config import Config # pylint: disable=import-outside-toplevel @@ -829,11 +927,15 @@ def _make_transfer( s3={"addressing_style": "path"}, request_checksum_calculation="when_required", response_checksum_validation="when_required", - retries={"max_attempts": 10, "mode": "standard"}, + # _with_retries owns retries; botocore's would multiply them and sleep + # while holding a concurrency slot + retries={"total_max_attempts": 1, "mode": "standard"}, connect_timeout=60, - read_timeout=3600, - # transfer requests share `concurrency` slots; metadata calls from - # the file workers can run alongside them + # a socket read silent this long is a stalled connection (runpod's + # proxy answers a slow request with 524 after about 100s) + read_timeout=300, + # every call holds one of `concurrency` slots; headroom for + # connections still being returned to the pool max_pool_connections=max(10, 2 * concurrency), ) client = boto3.session.Session().client( @@ -844,7 +946,7 @@ def _make_transfer( aws_secret_access_key=secret_key, config=config, ) - return _Transfer(client, volume_id, endpoint, concurrency, progress) + return _Transfer(client, volume_id, endpoint, concurrency, progress, force=force) def upload( @@ -858,13 +960,15 @@ def upload( endpoint_url: Optional[str] = None, concurrency: int = 4, progress: Optional[Progress] = None, + force: bool = False, ) -> Dict[str, int]: """upload a file or folder to a network volume. `remote_path` defaults to the local basename. a folder's contents go under `remote_path/`. for a file, a `remote_path` ending in "/" keeps - the local file name. unchanged files (same size, remote not older) - are skipped, and an interrupted multipart upload resumes on re-run. + the local file name. unchanged files (same size, remote not older, in + whole seconds) are skipped unless `force`, and an interrupted multipart + upload resumes on re-run. s3 keys come from the arguments or RUNPOD_S3_ACCESS_KEY and RUNPOD_S3_SECRET_KEY. without `datacenter`, the volume's datacenter is @@ -873,7 +977,7 @@ def upload( returns {files_transferred, files_skipped, bytes_transferred}. """ transfer = _make_transfer( - volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress + volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress, force ) return transfer.upload(Path(local_path), remote_path) @@ -889,20 +993,23 @@ def download( endpoint_url: Optional[str] = None, concurrency: int = 4, progress: Optional[Progress] = None, + force: bool = False, ) -> Dict[str, int]: """download a file or folder from a network volume. a `remote_path` ending in "/" (or one that only exists as a prefix) is a folder. `local_path` defaults to the remote basename. local files of the same size and not older than the remote copy are skipped (downloads - stamp the remote time on the file), and partial downloads resume on re-run. + stamp the remote time on the file) unless `force`, and partial downloads + resume on re-run. a `local_path` ending in "/" is a folder to download into. returns {files_transferred, files_skipped, bytes_transferred}. """ transfer = _make_transfer( - volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress + volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress, force ) - return transfer.download(remote_path, None if local_path is None else Path(local_path)) + # passed through as given: Path() would drop a trailing "/" + return transfer.download(remote_path, local_path) upload_to_network_volume = upload diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index f9b662c2b..69a020f7e 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1095,7 +1095,9 @@ def network_volume(): Uses the S3-compatible API. Set RUNPOD_S3_ACCESS_KEY and RUNPOD_S3_SECRET_KEY (create keys at console.runpod.io/user/settings). - Re-run an interrupted command to resume it. + Re-run an interrupted command to resume it. Parts of an upload that is + never re-run stay in the volume's hidden .s3compat_uploads folder and + count toward its size. """ @@ -1111,7 +1113,11 @@ def _transfer_options(func): click.option("--endpoint-url", default=None, hidden=True), click.option( "--concurrency", default=4, show_default=True, type=click.IntRange(1, 64), - help="Parallel parts, chunks, or files.", + help="Parallel requests.", + ), + click.option( + "--force", is_flag=True, + help="Transfer every file, even ones the size and time check calls unchanged.", ), ] for option in reversed(options): @@ -1144,7 +1150,7 @@ def _run_transfer(direction, volume_id, source, target, opts) -> None: def progress(event, path, nbytes): if event == "start": tasks[path] = bar.add_task(Path(path).name, total=nbytes) - elif event == "advance" and path in tasks: + elif event in ("advance", "reuse") and path in tasks: bar.advance(tasks[path], nbytes) elif event == "done" and path in tasks: bar.remove_task(tasks.pop(path)) @@ -1162,9 +1168,15 @@ def progress(event, path, nbytes): endpoint_url=opts["endpoint_url"], concurrency=opts["concurrency"], progress=progress, + force=opts["force"], ) + except KeyboardInterrupt: + click.echo("interrupted; re-run the same command to resume", err=True) + sys.exit(130) except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" + if nvt._is_dns_failure(exc): # pylint: disable=protected-access + hint = "\nthe s3 endpoint does not resolve: check --datacenter" raise click.ClickException(f"{exc}{hint}") from exc click.echo( @@ -1183,7 +1195,8 @@ def network_volume_upload(volume_id, local_path, remote_path, **opts): """Upload a file or folder to a network volume. REMOTE_PATH defaults to the local name. A folder's contents go under - REMOTE_PATH/. Unchanged files are skipped. + REMOTE_PATH/. Files whose remote copy has the same size and is not older + are skipped; --force sends them anyway. """ _run_transfer("uploaded", volume_id, local_path, remote_path, opts) @@ -1197,7 +1210,8 @@ def network_volume_download(volume_id, remote_path, local_path, **opts): """Download a file or folder from a network volume. A REMOTE_PATH ending in / is a folder. LOCAL_PATH defaults to the - remote name. Local files of the same size are skipped. + remote name. Local files of the same size and not older than the + remote copy are skipped; --force fetches them anyway. """ _run_transfer("downloaded", volume_id, remote_path, local_path, opts) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 16da8793b..823f8d3cc 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -136,6 +136,12 @@ def complete_multipart_upload(self, Bucket, Key, UploadId, MultipartUpload): del self.uploads[UploadId] return {} + def delete_object(self, Bucket, Key): + self._call("delete_object", Key=Key) + self.objects.pop(Key, None) + self.multipart_etags.pop(Key, None) + return {} + def abort_multipart_upload(self, Bucket, Key, UploadId): self._call("abort_multipart_upload", UploadId=UploadId) self.uploads.pop(UploadId, None) @@ -208,7 +214,8 @@ def test_multipart_resume_uploads_only_missing_and_bad_parts(tmp_path): assert sorted(kw["PartNumber"] for kw in fake.ops("upload_part")) == [2, 5] assert not fake.ops("create_multipart_upload") assert fake.objects["big.bin"][0] == DATA - assert summary["bytes_transferred"] == len(DATA) + # only the resent parts (2 and 5) count as transferred + assert summary["bytes_transferred"] == 16 assert not list((tmp_path / "state").glob("*.json")) @@ -400,7 +407,8 @@ def test_download_resumes_from_sidecar(tmp_path): assert dest.read_bytes() == DATA assert not (tmp_path / "out" / "big.bin.rpdownload.json").exists() assert not (tmp_path / "out" / "big.bin.rpdownload").exists() - assert summary["bytes_transferred"] == len(DATA) + # only the missing chunk counts as transferred + assert summary["bytes_transferred"] == 8 def test_download_restarts_when_remote_changed(tmp_path): @@ -487,13 +495,13 @@ def test_folder_download_filters_and_defaults(tmp_path, monkeypatch): fake.put("models/a.bin", b"a") fake.put("models/sub/", b"") fake.put("models/sub/b.bin", b"bb") - fake.put("models/.s3compat_uploads/x", b"junk") + fake.put(f"models/{nvt.UPLOADS_DIR}/x", b"junk") fake.put("other/c.bin", b"c") monkeypatch.chdir(tmp_path) summary = _transfer(fake, tmp_path, concurrency=2).download("models", None) assert (tmp_path / "models" / "a.bin").read_bytes() == b"a" assert (tmp_path / "models" / "sub" / "b.bin").read_bytes() == b"bb" - assert not (tmp_path / "models" / ".s3compat_uploads").exists() + assert not (tmp_path / "models" / nvt.UPLOADS_DIR).exists() assert summary["files_transferred"] == 2 @@ -513,13 +521,15 @@ def head(**kwargs): assert summary["files_transferred"] == 1 -def test_head_400_outside_download_is_an_error(tmp_path): +def test_upload_onto_remote_folder_is_refused(tmp_path): fake = FakeS3() + # runpod answers a head on a directory with 400 fake.fail = lambda op, kw: _client_error("BadRequest", 400) if op == "head_object" else None src = tmp_path / "f.txt" src.write_bytes(b"x") - with pytest.raises(ClientError): - _transfer(fake, tmp_path).upload(src, "f.txt") + with pytest.raises(nvt.NetworkVolumeTransferError, match="'models/'") as info: + _transfer(fake, tmp_path).upload(src, "models") + assert not nvt.is_resumable(info.value) @pytest.mark.parametrize( @@ -533,6 +543,11 @@ def test_head_400_outside_download_is_an_error(tmp_path): (nvt.NetworkVolumeTransferError("cut", resumable=True), True), (_client_error("524", 524), True), (ReadTimeoutError(endpoint_url="https://s3"), True), + (KeyboardInterrupt(), True), + # unknown or local failures repeat identically on a re-run + (RuntimeError("param validation"), False), + (OSError(28, "No space left on device"), False), + (ValueError("x"), False), ], ) def test_is_resumable(exc, expected): @@ -592,6 +607,9 @@ def test_make_transfer_requires_s3_keys(monkeypatch): def test_make_transfer_builds_datacenter_client(monkeypatch): + # a stale local aws profile must not leak into the session + for name in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE"): + monkeypatch.delenv(name, raising=False) monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") with patch( @@ -602,11 +620,19 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): assert transfer.client.meta.endpoint_url == "https://s3api-eu-ro-1.runpod.io/" assert transfer.client.meta.region_name == "EU-RO-1" assert transfer.client.meta.config.request_checksum_calculation == "when_required" + # _with_retries owns retries + assert transfer.client.meta.config.retries["total_max_attempts"] == 1 credentials = transfer.client._request_signer._credentials assert credentials.access_key == "user_x" assert transfer.bucket == "vol123" +def test_lookup_datacenter_escapes_volume_id(): + with patch("runpod.api.rest.run_rest_request", return_value={"dataCenter": "EU-RO-1"}) as rest: + nvt._lookup_datacenter("a/b?c") + rest.assert_called_once_with("GET", "/v2/network-volumes/a%2Fb%3Fc") + + def test_lookup_datacenter_missing(): with patch("runpod.api.rest.run_rest_request", return_value={}): with pytest.raises(nvt.NetworkVolumeTransferError): @@ -622,6 +648,14 @@ def test_public_functions_delegate(tmp_path): assert nvt.download("vol", "y", "x")[0] == "down" +def test_public_download_keeps_trailing_slash(tmp_path, monkeypatch): + fake = FakeS3() + fake.put("w.bin", b"w") + monkeypatch.setattr(nvt, "_make_transfer", lambda *a, **k: _transfer(fake, tmp_path)) + nvt.download("vol", "w.bin", str(tmp_path / "out") + "/") + assert (tmp_path / "out" / "w.bin").read_bytes() == b"w" + + def test_lazy_exports(): assert runpod.upload_to_network_volume is nvt.upload assert runpod.download_from_network_volume is nvt.download @@ -673,7 +707,7 @@ def test_cli_missing_path_has_no_resume_hint(): assert "resume" not in result.output -# ------------------------------------------------- review follow-ups +# ------------------------------------------- completion, collisions, limits @pytest.mark.parametrize( @@ -785,35 +819,6 @@ def test_dotdot_remote_path_rejected(tmp_path, op): assert not nvt.is_resumable(ValueError("x")) -def test_requests_in_flight_never_exceed_concurrency(tmp_path): - import threading - import time as _time - - fake = FakeS3() - active = {"now": 0, "peak": 0} - lock = threading.Lock() - original = fake.upload_part - - def tracked(**kwargs): - with lock: - active["now"] += 1 - active["peak"] = max(active["peak"], active["now"]) - _time.sleep(0.01) - try: - return original(**kwargs) - finally: - with lock: - active["now"] -= 1 - - fake.upload_part = tracked - folder = tmp_path / "data" - folder.mkdir() - for i in range(4): - (folder / f"f{i}.bin").write_bytes(DATA) - _transfer(fake, tmp_path, concurrency=2).upload(folder, "data") - assert active["peak"] <= 2 - - def test_ambiguous_complete_without_verifiable_etag_fails(tmp_path): fake = FakeS3(md5_etags=False) src = tmp_path / "big.bin" @@ -872,3 +877,145 @@ def tracked_call(op, **kwargs): _transfer(fake, tmp_path, concurrency=2).upload(folder, "data") _transfer(fake, tmp_path, concurrency=2).download("data/", tmp_path / "out") assert active["peak"] <= 2 + + +def test_dns_failure_is_not_retried_or_resumable(): + import socket as _socket + + from botocore.exceptions import EndpointConnectionError + + exc = EndpointConnectionError( + endpoint_url="https://s3api-xx-9.runpod.io/", error=_socket.gaierror(8, "nodename nor servname") + ) + assert not nvt._is_retryable(exc) + assert not nvt.is_resumable(exc) + reachable = EndpointConnectionError( + endpoint_url="https://s3api-eu-ro-1.runpod.io/", error=ConnectionRefusedError() + ) + assert nvt._is_retryable(reachable) + + +def test_single_put_change_with_older_mtime_is_resent_on_rerun(tmp_path): + fake = FakeS3() + src = tmp_path / "small.txt" + src.write_bytes(b"original") + old = src.stat().st_mtime - 3600 + os.utime(src, (old, old)) + attempts = [] + + def edit_then_fail(op, kw): + if op == "put_object" and not attempts: + attempts.append(1) + # a same-size edit whose mtime is older than the object the retry + # writes, so the skip rule alone would keep the object + src.write_bytes(b"CHANGED!") + past = src.stat().st_mtime - 1800 + os.utime(src, (past, past)) + return _client_error("ServiceUnavailable", 503) + return None + + fake.fail = edit_then_fail + with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload"): + _transfer(fake, tmp_path).upload(src, "small.txt") + assert "small.txt" not in fake.objects + fake.fail = None + _transfer(fake, tmp_path).upload(src, "small.txt") + assert fake.objects["small.txt"][0] == b"CHANGED!" + + +def test_missing_folder_is_not_found_but_empty_volume_is_fine(tmp_path): + fake = FakeS3() + fake.put("other/x", b"x") + with pytest.raises(FileNotFoundError): + _transfer(fake, tmp_path).download("missing/", tmp_path / "out") + empty = FakeS3() + summary = _transfer(empty, tmp_path).download("/", tmp_path / "all") + assert summary["files_transferred"] == 0 + + +def test_force_transfers_unchanged_files(tmp_path): + fake = FakeS3() + folder = tmp_path / "data" + folder.mkdir() + (folder / "a.txt").write_bytes(b"alpha") + _transfer(fake, tmp_path).upload(folder, "data") + forced = nvt._Transfer(fake, "vol", "https://s3", 1, None, state_dir=tmp_path / "state", force=True) + assert forced.upload(folder, "data")["files_transferred"] == 1 + out = tmp_path / "out" + _transfer(fake, tmp_path).download("data/", out) + forced = nvt._Transfer(fake, "vol", "https://s3", 1, None, state_dir=tmp_path / "state", force=True) + assert forced.download("data/", out)["files_transferred"] == 1 + + +@pytest.mark.parametrize( + "pair", + [("d/Model.bin", "d/model.bin"), ("d/s", "d/\u017f"), ("d/\u00e9", "d/e\u0301")], + ids=["case", "unicode-fold", "normalization"], +) +def test_download_refuses_targets_sharing_a_local_name(tmp_path, pair): + fake = FakeS3() + fake.put(pair[0], b"a") + fake.put(pair[1], b"b") + with pytest.raises(nvt.NetworkVolumeTransferError, match="same file"): + _transfer(fake, tmp_path).download("d/", tmp_path / "out") + assert not fake.ops("get_object") + + +def test_dot_segment_key_is_refused(tmp_path): + with pytest.raises(nvt.NetworkVolumeTransferError, match="unsafe"): + nvt._safe_target(tmp_path, "a/./b") + + +def test_download_into_missing_folder_with_trailing_slash(tmp_path): + fake = FakeS3() + fake.put("w.bin", b"w") + _transfer(fake, tmp_path).download("w.bin", str(tmp_path / "out") + "/") + assert (tmp_path / "out" / "w.bin").read_bytes() == b"w" + + +def test_folder_upload_skips_download_leftovers(tmp_path): + fake = FakeS3() + folder = tmp_path / "data" + folder.mkdir() + (folder / "a.bin").write_bytes(b"a") + (folder / "b.bin.rpdownload").write_bytes(b"partial") + (folder / "b.bin.rpdownload.json").write_text("{}") + summary = _transfer(fake, tmp_path).upload(folder, "data") + assert summary["files_transferred"] == 1 + assert sorted(fake.objects) == ["data/a.bin"] + + +def test_skip_rule_compares_whole_seconds(tmp_path): + fake = FakeS3() + stamp = datetime(2026, 1, 1, 12, 0, 0, tzinfo=timezone.utc) + fake.put("f.txt", b"abc", modified=stamp) + local = tmp_path / "f.txt" + local.write_bytes(b"abc") + # sub-second local precision against a whole-second server date + os.utime(local, (stamp.timestamp(), stamp.timestamp() + 0.7)) + assert _transfer(fake, tmp_path).upload(local, "f.txt")["files_skipped"] == 1 + assert _transfer(fake, tmp_path).download("f.txt", local)["files_skipped"] == 1 + + +def test_ctrl_c_stops_files_in_flight(tmp_path): + fake = FakeS3() + folder = tmp_path / "data" + folder.mkdir() + for i in range(6): + (folder / f"f{i}.bin").write_bytes(DATA + bytes([i])) + calls = [] + original = fake.upload_part + + def interrupt_once(**kwargs): + calls.append(1) + if len(calls) == 1: + raise KeyboardInterrupt + return original(**kwargs) + + fake.upload_part = interrupt_once + transfer = _transfer(fake, tmp_path, concurrency=2) + with pytest.raises(KeyboardInterrupt): + transfer.upload(folder, "data") + assert transfer._stop.is_set() + # 6 files x 5 parts would be 30 parts; workers stop at their next request + assert len(calls) < 10 From b923852d9d53b3aacd32bff7282801cece872e2d Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 18:08:03 -0400 Subject: [PATCH 08/23] fix: address network volume transfer review round 2 - an unresolvable endpoint fails before the transfer; a dns failure mid-transfer is a network drop and is retried (drops _is_dns_failure) - retry backoff ends early on ctrl-c (stop-aware wait) - single put changed mid-upload: keep the object if the file vanished, and say to re-run with --force if the delete fails - head 400 on an upload target is called a folder only if a listing confirms it; NoSuchUpload and 416 resume - refuse a file that would also be a folder; collision wording covers normalization; leftovers are reported as skipped - state writes sync the directory; cli tests for ctrl-c and the endpoint check --- runpod/network_volume_transfer.py | 162 +++++++++++++++++--------- runpod/rp_cli/main.py | 2 - tests/test_network_volume_transfer.py | 132 +++++++++++++++++++-- 3 files changed, 232 insertions(+), 64 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 1dce409ae..2fed1b5e8 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -107,9 +107,6 @@ def _is_retryable(exc: Exception) -> bool: if isinstance(exc, _Retryable): return True - if _is_dns_failure(exc): - # a host that does not resolve is a wrong datacenter, not a blip - return False if isinstance( exc, ( @@ -130,23 +127,6 @@ def _is_retryable(exc: Exception) -> bool: return False -def _is_dns_failure(exc: BaseException) -> bool: - """whether `exc` (or an error it wraps) is a host name that does not resolve.""" - seen = 0 - stack = [exc] - while stack and seen < 20: - err = stack.pop() - seen += 1 - if isinstance(err, socket.gaierror): - return True - kwargs = getattr(err, "kwargs", None) - wrapped = [err.__cause__, err.__context__, getattr(err, "reason", None)] - if isinstance(kwargs, dict): - wrapped.append(kwargs.get("error")) - stack.extend(w for w in wrapped if isinstance(w, BaseException)) - return False - - def is_resumable(exc: BaseException) -> bool: """whether re-running the same transfer can get past `exc`. @@ -159,18 +139,37 @@ def is_resumable(exc: BaseException) -> bool: return True if isinstance(exc, NetworkVolumeTransferError): return exc.resumable + # final-looking errors a re-run gets past: the multipart upload vanished + # (e.g. after a complete whose 524 hid that it finished), or a range that + # no longer fits an object that changed mid-download + code, status = _error_code(exc) + if code == "NoSuchUpload" or status == 416: + return True return _is_retryable(exc) -def _with_retries(func: Callable[[], Any], attempts: int = RETRY_ATTEMPTS) -> Any: - """call func, retrying transient failures with exponential backoff.""" +def _wait(delay: float, stop: Optional[threading.Event] = None) -> None: + """sleep for a retry backoff, cut short when `stop` is set.""" + if stop is None: + _sleep(delay) + else: + stop.wait(delay) + + +def _with_retries( + func: Callable[[], Any], + attempts: int = RETRY_ATTEMPTS, + stop: Optional[threading.Event] = None, +) -> Any: + """call func, retrying transient failures with exponential backoff. a set + `stop` (ctrl-c) ends a backoff early; the next attempt then raises.""" for attempt in range(attempts): try: return func() except Exception as exc: # noqa: BLE001 - filtered by _is_retryable if attempt == attempts - 1 or not _is_retryable(exc): raise - _sleep(min(2**attempt, 30) + random.uniform(0, 1)) + _wait(min(2**attempt, 30) + random.uniform(0, 1), stop) raise AssertionError("unreachable") @@ -182,6 +181,18 @@ def _write_json_atomic(path: Path, data: Dict[str, Any]) -> None: handle.flush() os.fsync(handle.fileno()) os.replace(tmp, path) + # sync the directory too, so the rename survives a power loss (best + # effort: not every platform can open a directory) + try: + fd = os.open(path.parent, os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + except OSError: + pass + finally: + os.close(fd) def _read_json(path: Path) -> Optional[Dict[str, Any]]: @@ -320,6 +331,9 @@ def _request(self, func: Callable[[], Any]) -> Any: with self._slots: return func() + def _retry(self, func: Callable[[], Any]) -> Any: + return _with_retries(func, stop=self._stop) + def _pool(self, tasks: List[Callable[[], Any]]) -> List[Any]: return _run_pool(tasks, self.concurrency, self._stop) @@ -347,7 +361,7 @@ def head(self, key: str, directory_ok: bool = False) -> Optional[Dict[str, Any]] HeadObject on a directory with 400 rather than 404. """ try: - return _with_retries( + return self._retry( lambda: self._request( lambda: self.client.head_object(Bucket=self.bucket, Key=key) ) @@ -362,7 +376,7 @@ def list_prefix(self, prefix: str) -> List[Dict[str, Any]]: kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Prefix": prefix} while True: try: - page = _with_retries( + page = self._retry( lambda: self._request(lambda: self.client.list_objects_v2(**kwargs)) ) except Exception as exc: # noqa: BLE001 @@ -375,11 +389,26 @@ def list_prefix(self, prefix: str) -> List[Dict[str, Any]]: return objects kwargs["ContinuationToken"] = token + def _is_folder(self, key: str) -> bool: + """whether something is listed under `key/` (runpod lists even an + empty folder's own marker).""" + try: + page = self._retry( + lambda: self._request( + lambda: self.client.list_objects_v2( + Bucket=self.bucket, Prefix=f"{key}/", MaxKeys=1 + ) + ) + ) + except Exception: # noqa: BLE001 - unknown: report the original error + return False + return bool(page.get("Contents")) + def _list_parts(self, key: str, upload_id: str) -> List[Dict[str, Any]]: parts: List[Dict[str, Any]] = [] kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Key": key, "UploadId": upload_id} while True: - page = _with_retries( + page = self._retry( lambda: self._request(lambda: self.client.list_parts(**kwargs)) ) parts.extend(page.get("Parts") or []) @@ -396,7 +425,7 @@ def upload_file(self, local: Path, key: str) -> None: label = str(local) self._emit("start", label, size) if size <= SINGLE_PUT_MAX: - _with_retries( + self._retry( lambda: self._request( lambda: self.client.put_object( Bucket=self.bucket, Key=key, Body=local.read_bytes() @@ -406,17 +435,20 @@ def upload_file(self, local: Path, key: str) -> None: # a retry rereads the file; if it changed meanwhile, the object # may hold either version and is newer than the local edit, so a # re-run would skip it. remove it so the re-run sends the file. - now = local.stat() + try: + now = local.stat() + except OSError as exc: + # gone or unreadable: the object is what was read, keep it + raise NetworkVolumeTransferError(f"{local} is gone after its upload: {exc}") from exc if now.st_size != size or now.st_mtime_ns != stat.st_mtime_ns: try: - _with_retries( - lambda: self._request( - lambda: self.client.delete_object(Bucket=self.bucket, Key=key) - ) - ) - except Exception: # noqa: BLE001 - # --force on the re-run still sends it - pass + # not via self._retry: this cleanup runs even after ctrl-c + _with_retries(lambda: self.client.delete_object(Bucket=self.bucket, Key=key)) + except Exception as exc: # noqa: BLE001 + raise NetworkVolumeTransferError( + f"{local} changed during upload and the remote copy could not be " + f"removed ({exc}); re-run with --force" + ) from exc raise NetworkVolumeTransferError( f"{local} changed during upload", resumable=True ) @@ -480,7 +512,7 @@ def _multipart_upload( _remove(state_path) if upload_id is None: - created = _with_retries( + created = self._retry( lambda: self._request( lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=key) ) @@ -521,7 +553,7 @@ def send(): Body=body, ) - response = _with_retries(lambda: self._request(send)) + response = self._retry(lambda: self._request(send)) self._emit("advance", label, length) return number, response["ETag"] @@ -535,14 +567,11 @@ def send(): if now.st_size != size or now.st_mtime_ns != mtime_ns: self._abort(key, upload_id) _remove(state_path) - raise NetworkVolumeTransferError( - f"{local} changed during upload; run the command again to upload the new version", - resumable=True, - ) + raise NetworkVolumeTransferError(f"{local} changed during upload", resumable=True) parts = [{"PartNumber": n, "ETag": done[n]} for n in sorted(done)] try: - _with_retries( + self._retry( lambda: self._request( lambda: self.client.complete_multipart_upload( Bucket=self.bucket, @@ -607,9 +636,11 @@ def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: for obj in self.list_prefix(f"{prefix}/" if prefix else "") } jobs = [] - for path in sorted( - p for p in local.rglob("*") if p.is_file() and not _is_download_leftover(p) - ): + for path in sorted(p for p in local.rglob("*") if p.is_file()): + if _is_download_leftover(path): + # an interrupted download's partial file is not data + self._emit("skip", str(path), path.stat().st_size) + continue rel = path.relative_to(local).as_posix() key = f"{prefix}/{rel}" if prefix else rel if not self.force and self._unchanged_remote(remote.get(key), path): @@ -628,8 +659,9 @@ def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: try: remote = self.head(key) except Exception as exc: # noqa: BLE001 - # runpod answers a head on a directory with 400 - if _error_code(exc)[1] == 400: + # runpod answers a head on a directory with 400, but a head has no + # body, so a 400 for any other reason looks the same: confirm + if _error_code(exc)[1] == 400 and self._is_folder(key): raise NetworkVolumeTransferError( f"{key!r} is a folder on the volume; use {key + '/'!r} to upload into it" ) from exc @@ -738,7 +770,7 @@ def get(): raise _Retryable(f"short read for {key} bytes {start}-{end}") return data - data = _with_retries(lambda: self._request(get)) + data = self._retry(lambda: self._request(get)) with open(tmp, "r+b") as handle: handle.seek(start) handle.write(data) @@ -833,7 +865,8 @@ def _fold(name: str) -> str: def _check_local_names(targets: List[Path], root: Path) -> None: """refuse targets that would share a path on disk: two targets that are one name on a case-insensitive or normalizing filesystem ("Model.bin" and - "model.bin"), or a file whose partial or sidecar name is also a target or + "model.bin"), a file whose name is also a folder another file lives in + ("a" beside "A/b"), or a file whose partial or sidecar name is also a target or a directory a target lives in ("x" beside "x.rpdownload" or "x.rpdownload/y"). that can refuse a listing a case-sensitive filesystem would take, which is the safe direction.""" @@ -844,13 +877,19 @@ def _check_local_names(targets: List[Path], root: Path) -> None: key = _fold(str(rel)) if key in files: raise NetworkVolumeTransferError( - f"cannot download both {files[key]} and {target}: " - "they are the same file on a case-insensitive filesystem" + f"cannot download both {files[key]} and {target}: they would be one " + "file on a case-insensitive or normalizing filesystem" ) files[key] = target for parent in rel.parents: if str(parent) != ".": dirs[_fold(str(parent))] = target + for key, target in files.items(): + if key in dirs: + raise NetworkVolumeTransferError( + f"cannot download both {target} and {dirs[key]}: " + "the first would have to be a file and a folder" + ) for target in targets: rel = target.relative_to(root) for suffix in (PARTIAL_SUFFIX, SIDECAR_SUFFIX): @@ -900,6 +939,24 @@ def _lookup_datacenter(volume_id: str) -> str: return str(datacenter) +def _check_endpoint_resolves(endpoint: str) -> None: + """fail fast on a datacenter whose s3 endpoint does not resolve, before + any transfer starts. once a transfer is running, a failed lookup is a + network drop and is retried like any other.""" + from urllib.parse import urlparse # pylint: disable=import-outside-toplevel + + host = urlparse(endpoint).hostname + if not host: + raise NetworkVolumeTransferError(f"invalid s3 endpoint {endpoint!r}") + try: + socket.getaddrinfo(host, 443) + except socket.gaierror as exc: + raise NetworkVolumeTransferError( + f"the s3 endpoint {host} does not resolve: check --datacenter (only " + "datacenters with the s3 api are supported) or your network connection" + ) from exc + + def _make_transfer( volume_id: str, datacenter: Optional[str], @@ -922,6 +979,7 @@ def _make_transfer( ) datacenter = datacenter or _lookup_datacenter(volume_id) endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" + _check_endpoint_resolves(endpoint) concurrency = max(1, int(concurrency)) config = Config( s3={"addressing_style": "path"}, diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index 69a020f7e..b51ba185a 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1175,8 +1175,6 @@ def progress(event, path, nbytes): sys.exit(130) except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" - if nvt._is_dns_failure(exc): # pylint: disable=protected-access - hint = "\nthe s3 endpoint does not resolve: check --datacenter" raise click.ClickException(f"{exc}{hint}") from exc click.echo( diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 823f8d3cc..b4c18a7a8 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -13,6 +13,8 @@ import runpod import runpod.network_volume_transfer as nvt + +REAL_WAIT = nvt._wait import runpod.error @@ -73,7 +75,7 @@ def get_object(self, Bucket, Key, Range): start, end = (int(x) for x in Range.split("=")[1].split("-")) return {"Body": io.BytesIO(data[start : end + 1]), "ETag": self.object_etag(Key)} - def list_objects_v2(self, Bucket, Prefix, ContinuationToken=None): + def list_objects_v2(self, Bucket, Prefix, ContinuationToken=None, MaxKeys=None): self._call("list_objects_v2", Prefix=Prefix, ContinuationToken=ContinuationToken) keys = sorted(k for k in self.objects if k.startswith(Prefix)) start = int(ContinuationToken or 0) @@ -153,6 +155,9 @@ def ops(self, op): @pytest.fixture(autouse=True) def small_sizes(monkeypatch): monkeypatch.setattr(nvt, "_sleep", lambda _: None) + monkeypatch.setattr(nvt, "_wait", lambda delay, stop=None: None) + # transfers built through _make_transfer must not hit dns in tests + monkeypatch.setattr(nvt, "_check_endpoint_resolves", lambda endpoint: None) monkeypatch.setattr(nvt, "SINGLE_PUT_MAX", 10) monkeypatch.setattr(nvt, "DEFAULT_PART_SIZE", 8) monkeypatch.setattr(nvt, "DOWNLOAD_CHUNK_SIZE", 8) @@ -523,6 +528,7 @@ def head(**kwargs): def test_upload_onto_remote_folder_is_refused(tmp_path): fake = FakeS3() + fake.put("models/a.bin", b"a") # runpod answers a head on a directory with 400 fake.fail = lambda op, kw: _client_error("BadRequest", 400) if op == "head_object" else None src = tmp_path / "f.txt" @@ -697,6 +703,26 @@ def test_cli_download_failure_says_resume(): assert "re-run the same command to resume" in result.output +def test_cli_ctrl_c_exits_130_with_resume_hint(): + from runpod.rp_cli.main import cli + + with patch.object(nvt, "upload", side_effect=KeyboardInterrupt): + result = CliRunner().invoke(cli, ["network-volume", "upload", "vol", "."]) + assert result.exit_code == 130 + assert "interrupted; re-run the same command to resume" in result.output + + +def test_cli_unresolvable_endpoint_has_no_resume_hint(): + from runpod.rp_cli.main import cli + + err = nvt.NetworkVolumeTransferError("the s3 endpoint s3api-x.invalid does not resolve: check --datacenter") + with patch.object(nvt, "download", side_effect=err): + result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b", "--datacenter", "X"]) + assert result.exit_code == 1 + assert "check --datacenter" in result.output + assert "re-run" not in result.output + + def test_cli_missing_path_has_no_resume_hint(): from runpod.rp_cli.main import cli @@ -879,20 +905,106 @@ def tracked_call(op, **kwargs): assert active["peak"] <= 2 -def test_dns_failure_is_not_retried_or_resumable(): +def test_unresolvable_endpoint_fails_before_the_transfer(monkeypatch): + monkeypatch.undo() # the autouse fixture stubs the check out + with pytest.raises(nvt.NetworkVolumeTransferError, match="--datacenter"): + nvt._check_endpoint_resolves("https://s3api-nope.invalid/") + nvt._check_endpoint_resolves("http://127.0.0.1:9/") + + +def test_dns_failure_mid_transfer_is_a_network_drop(): import socket as _socket from botocore.exceptions import EndpointConnectionError exc = EndpointConnectionError( - endpoint_url="https://s3api-xx-9.runpod.io/", error=_socket.gaierror(8, "nodename nor servname") - ) - assert not nvt._is_retryable(exc) - assert not nvt.is_resumable(exc) - reachable = EndpointConnectionError( - endpoint_url="https://s3api-eu-ro-1.runpod.io/", error=ConnectionRefusedError() + endpoint_url="https://s3api-eu-ro-1.runpod.io/", error=_socket.gaierror(8, "nodename nor servname") ) - assert nvt._is_retryable(reachable) + assert nvt._is_retryable(exc) + assert nvt.is_resumable(exc) + + +def test_head_400_that_is_not_a_folder_is_reported_as_is(tmp_path): + fake = FakeS3() + fake.fail = lambda op, kw: _client_error("BadRequest", 400) if op == "head_object" else None + src = tmp_path / "f.txt" + src.write_bytes(b"x") + with pytest.raises(ClientError): + _transfer(fake, tmp_path).upload(src, "f.txt") + + +def test_vanished_local_file_keeps_the_uploaded_object(tmp_path): + fake = FakeS3() + src = tmp_path / "small.txt" + src.write_bytes(b"original") + original = fake.put_object + + def put_then_vanish(**kwargs): + result = original(**kwargs) + src.unlink() + return result + + fake.put_object = put_then_vanish + with pytest.raises(nvt.NetworkVolumeTransferError, match="gone"): + _transfer(fake, tmp_path).upload(src, "small.txt") + assert fake.objects["small.txt"][0] == b"original" + assert not fake.ops("delete_object") + + +def test_failed_delete_after_change_asks_for_force(tmp_path): + fake = FakeS3() + src = tmp_path / "small.txt" + src.write_bytes(b"original") + original = fake.put_object + + def put_then_edit(**kwargs): + result = original(**kwargs) + src.write_bytes(b"CHANGED!") + later = src.stat().st_mtime + 60 + os.utime(src, (later, later)) + return result + + fake.put_object = put_then_edit + fake.fail = lambda op, kw: _client_error("AccessDenied", 403) if op == "delete_object" else None + with pytest.raises(nvt.NetworkVolumeTransferError, match="--force") as info: + _transfer(fake, tmp_path).upload(src, "small.txt") + assert not nvt.is_resumable(info.value) + + +@pytest.mark.parametrize( + "exc", [_client_error("NoSuchUpload", 404), _client_error("InvalidRange", 416)] +) +def test_vanished_upload_and_stale_range_resume(exc): + assert nvt.is_resumable(exc) + + +def test_download_refuses_file_that_is_also_a_folder(tmp_path): + fake = FakeS3() + fake.put("d/a", b"file") + fake.put("d/A/b", b"nested") + with pytest.raises(nvt.NetworkVolumeTransferError, match="file and a folder"): + _transfer(fake, tmp_path).download("d/", tmp_path / "out") + + +def test_retry_backoff_ends_when_stopped(monkeypatch): + import threading + import time as _time + + monkeypatch.setattr(nvt, "_wait", REAL_WAIT) + stop = threading.Event() + calls = [] + + def fail(): + calls.append(1) + if len(calls) == 1: + stop.set() # ctrl-c arrives during the first attempt + raise ConnectionError("down") + + started = _time.monotonic() + with pytest.raises(ConnectionError): + nvt._with_retries(fail, attempts=3, stop=stop) + # the backoffs (1s, 2s plus jitter) are cut short + assert _time.monotonic() - started < 0.5 def test_single_put_change_with_older_mtime_is_resent_on_rerun(tmp_path): @@ -956,7 +1068,7 @@ def test_download_refuses_targets_sharing_a_local_name(tmp_path, pair): fake = FakeS3() fake.put(pair[0], b"a") fake.put(pair[1], b"b") - with pytest.raises(nvt.NetworkVolumeTransferError, match="same file"): + with pytest.raises(nvt.NetworkVolumeTransferError, match="one file"): _transfer(fake, tmp_path).download("d/", tmp_path / "out") assert not fake.ops("get_object") From 8c260d25f9ddad81593f4d89116db5a8886a2ec2 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 18:14:55 -0400 Subject: [PATCH 09/23] fix: endpoint pre-check skips proxied endpoints; stub dns in its tests - behind an https proxy (honoured by botocore) the proxy resolves the endpoint, so the local lookup is skipped - endpoint check tests stub getaddrinfo and proxies instead of real dns - explain the ignored directory fsync failure --- runpod/network_volume_transfer.py | 13 +++++++-- tests/test_network_volume_transfer.py | 40 ++++++++++++++++++++++++--- 2 files changed, 46 insertions(+), 7 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 2fed1b5e8..fc32f2bbd 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -190,7 +190,8 @@ def _write_json_atomic(path: Path, data: Dict[str, Any]) -> None: try: os.fsync(fd) except OSError: - pass + # some platforms refuse to fsync a directory; the file itself is synced + return finally: os.close(fd) @@ -942,12 +943,18 @@ def _lookup_datacenter(volume_id: str) -> str: def _check_endpoint_resolves(endpoint: str) -> None: """fail fast on a datacenter whose s3 endpoint does not resolve, before any transfer starts. once a transfer is running, a failed lookup is a - network drop and is retried like any other.""" + network drop and is retried like any other. behind an https proxy (which + botocore honours from the environment) the proxy resolves the name, so + the local lookup is skipped.""" from urllib.parse import urlparse # pylint: disable=import-outside-toplevel + from urllib.request import getproxies, proxy_bypass # pylint: disable=import-outside-toplevel - host = urlparse(endpoint).hostname + parsed = urlparse(endpoint) + host = parsed.hostname if not host: raise NetworkVolumeTransferError(f"invalid s3 endpoint {endpoint!r}") + if getproxies().get(parsed.scheme or "https") and not proxy_bypass(host): + return try: socket.getaddrinfo(host, 443) except socket.gaierror as exc: diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index b4c18a7a8..0d69c1868 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -15,6 +15,7 @@ import runpod.network_volume_transfer as nvt REAL_WAIT = nvt._wait +REAL_ENDPOINT_CHECK = nvt._check_endpoint_resolves import runpod.error @@ -905,11 +906,42 @@ def tracked_call(op, **kwargs): assert active["peak"] <= 2 -def test_unresolvable_endpoint_fails_before_the_transfer(monkeypatch): - monkeypatch.undo() # the autouse fixture stubs the check out +@pytest.fixture +def endpoint_check(monkeypatch): + """the real endpoint check with dns and proxies stubbed (the autouse + fixture replaces the check itself).""" + import socket as _socket + + monkeypatch.setattr(nvt, "_check_endpoint_resolves", REAL_ENDPOINT_CHECK) + for name in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "NO_PROXY", "no_proxy"): + monkeypatch.delenv(name, raising=False) + state = {"error": None} + + def fake_getaddrinfo(host, port, *args, **kwargs): + if state["error"] is not None: + raise state["error"] + return [(_socket.AF_INET, _socket.SOCK_STREAM, 6, "", ("127.0.0.1", port))] + + monkeypatch.setattr(_socket, "getaddrinfo", fake_getaddrinfo) + return state + + +def test_unresolvable_endpoint_fails_before_the_transfer(endpoint_check): + import socket as _socket + + endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") with pytest.raises(nvt.NetworkVolumeTransferError, match="--datacenter"): - nvt._check_endpoint_resolves("https://s3api-nope.invalid/") - nvt._check_endpoint_resolves("http://127.0.0.1:9/") + nvt._check_endpoint_resolves("https://s3api-nope-1.runpod.io/") + endpoint_check["error"] = None + nvt._check_endpoint_resolves("https://s3api-eu-ro-1.runpod.io/") + + +def test_endpoint_check_skipped_behind_a_proxy(endpoint_check, monkeypatch): + import socket as _socket + + endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") + monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:3128") + nvt._check_endpoint_resolves("https://s3api-eu-ro-1.runpod.io/") def test_dns_failure_mid_transfer_is_a_network_drop(): From 0c8b4d39a5598aba4a634b372eab42fee0e66a62 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 18:22:01 -0400 Subject: [PATCH 10/23] fix: make the ctrl-c test deterministic and close its stop race - a request re-checks the stop flag after it gets a slot, so nothing starts after an interrupt lands - ctrl-c test blocks later parts on the stop event and asserts a bound that holds under any scheduling - backoff test restores the real sleep so a revert is caught - skip the endpoint dns check for an explicit endpoint_url - move the saved helpers below the imports --- runpod/network_volume_transfer.py | 8 +++++++- tests/test_network_volume_transfer.py | 27 +++++++++++++++++++-------- 2 files changed, 26 insertions(+), 9 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index fc32f2bbd..9ae9115d0 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -330,6 +330,9 @@ def _request(self, func: Callable[[], Any]) -> Any: if self._stop.is_set(): raise _Stopped("interrupted") with self._slots: + # checked again: an interrupt may land while waiting for the slot + if self._stop.is_set(): + raise _Stopped("interrupted") return func() def _retry(self, func: Callable[[], Any]) -> Any: @@ -986,7 +989,10 @@ def _make_transfer( ) datacenter = datacenter or _lookup_datacenter(volume_id) endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" - _check_endpoint_resolves(endpoint) + if not endpoint_url: + # an explicit endpoint (a local stand-in, a custom route) is the + # caller's; only the derived datacenter endpoint is checked + _check_endpoint_resolves(endpoint) concurrency = max(1, int(concurrency)) config = Config( s3={"addressing_style": "path"}, diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 0d69c1868..db4a6394b 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -13,10 +13,11 @@ import runpod import runpod.network_volume_transfer as nvt +import runpod.error +# the real helpers, before the autouse fixture stubs them REAL_WAIT = nvt._wait REAL_ENDPOINT_CHECK = nvt._check_endpoint_resolves -import runpod.error def _client_error(code, status, op="Op"): @@ -1023,6 +1024,8 @@ def test_retry_backoff_ends_when_stopped(monkeypatch): import time as _time monkeypatch.setattr(nvt, "_wait", REAL_WAIT) + # the real sleep too, so a revert to an uninterruptible sleep is caught + monkeypatch.setattr(nvt, "_sleep", _time.sleep) stop = threading.Event() calls = [] @@ -1148,18 +1151,26 @@ def test_ctrl_c_stops_files_in_flight(tmp_path): for i in range(6): (folder / f"f{i}.bin").write_bytes(DATA + bytes([i])) calls = [] + lock = __import__("threading").Lock() original = fake.upload_part + transfer = _transfer(fake, tmp_path, concurrency=2) - def interrupt_once(**kwargs): - calls.append(1) - if len(calls) == 1: + def interrupt_first(**kwargs): + with lock: + calls.append(1) + first = len(calls) == 1 + if first: raise KeyboardInterrupt + # every other part waits for the interrupt to land, so what runs + # after it is deterministic: only parts already in flight finish + transfer._stop.wait(5) return original(**kwargs) - fake.upload_part = interrupt_once - transfer = _transfer(fake, tmp_path, concurrency=2) + fake.upload_part = interrupt_first with pytest.raises(KeyboardInterrupt): transfer.upload(folder, "data") assert transfer._stop.is_set() - # 6 files x 5 parts would be 30 parts; workers stop at their next request - assert len(calls) < 10 + # 6 files x 5 parts would be 30 parts. the interrupted part, plus at most + # one part per request slot that started before the interrupt landed; + # nothing starts after it + assert len(calls) <= 3 From f1fa981267e314fbe69badf7b7dd3d143a08cdda Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 18:39:39 -0400 Subject: [PATCH 11/23] refactor: move transfers onto NetworkVolume and rp volume upload and download are NetworkVolume methods instead of top-level functions, and the cli commands live under rp volume instead of a new rp network-volume group. the volume is found by id or name with the api key, so --datacenter is gone. --- CHANGELOG.md | 2 +- docs/cli/references/command_line_interface.md | 25 ++-- runpod/__init__.py | 4 - runpod/apps/volume.py | 79 +++++++++++ runpod/network_volume_transfer.py | 82 +++++------ runpod/rp_cli/main.py | 131 +----------------- runpod/rp_cli/volume.py | 126 +++++++++++++++++ tests/test_network_volume_transfer.py | 100 ++++++++----- 8 files changed, 325 insertions(+), 224 deletions(-) create mode 100644 runpod/rp_cli/volume.py diff --git a/CHANGELOG.md b/CHANGELOG.md index da416680d..fb4db3893 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Features -* `rp network-volume upload/download` and `runpod.upload_to_network_volume` / `runpod.download_from_network_volume` transfer files and folders over a network volume's S3-compatible API. Re-running an interrupted transfer resumes it: unchanged files are skipped, multipart uploads reuse MD5-verified parts, and downloads keep finished chunks. +* `rp volume upload/download` and `NetworkVolume.upload` / `NetworkVolume.download` transfer files and folders over a network volume's S3-compatible API. Re-running an interrupted transfer resumes it: unchanged files are skipped, multipart uploads reuse MD5-verified parts, and downloads keep finished chunks. * Apps expose `NetworkVolume` and `GlobalVolume` resources with explicit path-keyed mounts, lazy create-if-missing provisioning, and execution-bound filesystem access. Global volumes use GraphQL for lookup and creation. ### Bug Fixes diff --git a/docs/cli/references/command_line_interface.md b/docs/cli/references/command_line_interface.md index 496d55d88..9ac4e857c 100644 --- a/docs/cli/references/command_line_interface.md +++ b/docs/cli/references/command_line_interface.md @@ -25,8 +25,8 @@ rp pod create rp pod connect POD_ID # Network volume files -rp network-volume upload VOLUME_ID ./dataset datasets/dataset -rp network-volume download VOLUME_ID outputs/ ./outputs +rp volume upload VOLUME ./dataset datasets/dataset +rp volume download VOLUME outputs/ ./outputs ``` ## Overview @@ -67,17 +67,17 @@ exit with an error and report the endpoint IDs that still need deletion. ### Network volume files ```bash -rp network-volume upload VOLUME_ID ./model.safetensors models/ -rp network-volume upload VOLUME_ID ./dataset datasets/dataset -rp network-volume download VOLUME_ID models/model.safetensors -rp network-volume download VOLUME_ID outputs/ ./outputs +rp volume upload VOLUME ./model.safetensors models/ +rp volume upload VOLUME ./dataset datasets/dataset +rp volume download VOLUME models/model.safetensors +rp volume download VOLUME outputs/ ./outputs ``` Transfers use the volume's S3-compatible API, so they need S3 API keys (separate from the API key): set `RUNPOD_S3_ACCESS_KEY` and `RUNPOD_S3_SECRET_KEY`, or pass `--s3-access-key` / `--s3-secret-key`. The volume must be in a [datacenter with the S3 API](https://docs.runpod.io/storage/s3-api). -Its datacenter is looked up with your API key unless you pass `--datacenter`. +`VOLUME` is the volume's ID or name, looked up with your API key. To resume an interrupted transfer, run the same command again. Files that are unchanged (same size, destination not older) are skipped, a large upload @@ -86,5 +86,12 @@ chunks in `.rpdownload`. `--force` transfers every file. Parts of an upload that is never resumed stay in the volume's hidden `.s3compat_uploads` folder and count toward its size. -The same transfers are available from Python as -`runpod.upload_to_network_volume` and `runpod.download_from_network_volume`. +The same transfers are available from Python: + +```python +from runpod import NetworkVolume + +volume = NetworkVolume("my-volume") +volume.upload("./dataset", "datasets/dataset") +volume.download("outputs/", "./outputs") +``` diff --git a/runpod/__init__.py b/runpod/__init__.py index e9ddd406d..427023537 100644 --- a/runpod/__init__.py +++ b/runpod/__init__.py @@ -68,10 +68,6 @@ "patch": "runpod.apps.markers", "post": "runpod.apps.markers", "put": "runpod.apps.markers", - # network volume file transfers - "download_from_network_volume": "runpod.network_volume_transfer", - "network_volume_transfer": "runpod.network_volume_transfer", - "upload_to_network_volume": "runpod.network_volume_transfer", # logger "RunPodLogger": "runpod.serverless.modules.rp_logger", # submodule diff --git a/runpod/apps/volume.py b/runpod/apps/volume.py index 487093587..888463200 100644 --- a/runpod/apps/volume.py +++ b/runpod/apps/volume.py @@ -84,6 +84,85 @@ def name(self) -> str: def __repr__(self) -> str: return f"" + def upload( + self, + local_path, + remote_path: Optional[str] = None, + *, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + concurrency: int = 4, + progress=None, + force: bool = False, + endpoint_url: Optional[str] = None, + ) -> Dict[str, int]: + """upload a file or folder to this volume through its s3 api. + + the volume must exist; it is found by id or name with the runpod + api key, which also gives its datacenter. `remote_path` defaults to + the local basename. a folder's contents go under `remote_path/`. for + a file, a `remote_path` ending in "/" keeps the local file name. + unchanged files (same size, remote not older, in whole seconds) are + skipped unless `force`, and an interrupted multipart upload resumes + on re-run. + + s3 keys come from the arguments or RUNPOD_S3_ACCESS_KEY and + RUNPOD_S3_SECRET_KEY. `progress(event, path, nbytes)` receives + "start", "advance", "reuse", "skip", and "done" events. + + returns {files_transferred, files_skipped, bytes_transferred}. + """ + from runpod import network_volume_transfer as nvt + + return nvt.upload( + self.reference, + local_path, + remote_path, + access_key=access_key, + secret_key=secret_key, + endpoint_url=endpoint_url, + concurrency=concurrency, + progress=progress, + force=force, + ) + + def download( + self, + remote_path: str, + local_path=None, + *, + access_key: Optional[str] = None, + secret_key: Optional[str] = None, + concurrency: int = 4, + progress=None, + force: bool = False, + endpoint_url: Optional[str] = None, + ) -> Dict[str, int]: + """download a file or folder from this volume through its s3 api. + + a `remote_path` ending in "/" (or one that only exists as a prefix) + is a folder. `local_path` defaults to the remote basename; one + ending in "/" is a folder to download into. local files of the same + size and not older than the remote copy are skipped (downloads stamp + the remote time on the file) unless `force`, and partial downloads + resume on re-run. keys and `progress` work as in `upload`. + + returns {files_transferred, files_skipped, bytes_transferred}. + """ + from runpod import network_volume_transfer as nvt + + return nvt.download( + self.reference, + remote_path, + local_path, + access_key=access_key, + secret_key=secret_key, + endpoint_url=endpoint_url, + concurrency=concurrency, + progress=progress, + force=force, + ) + class GlobalVolume(Volume): """a global volume resolved by name or id without datacenter placement.""" diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 9ae9115d0..189f7a1e2 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -1,7 +1,8 @@ """resumable file and folder transfers to and from network volumes. transfers go through the runpod s3-compatible api. the bucket is the -network volume id and the endpoint is per datacenter. +network volume id and the endpoint is per datacenter. call these +through NetworkVolume.upload and NetworkVolume.download. uploads above 64 MiB use multipart. the upload id is kept in a small state file under ~/.runpod/transfers, so re-running an interrupted @@ -929,18 +930,32 @@ def _safe_target(root: Path, rel: str) -> Path: # ---------------------------------------------------------------- client -def _lookup_datacenter(volume_id: str) -> str: +def _resolve_volume(reference: str) -> Tuple[str, str]: + """the id and datacenter of the network volume with this id or name.""" from runpod.api.rest import run_rest_request # pylint: disable=import-outside-toplevel - - from runpod.api.ctl_commands import _path_segment # pylint: disable=import-outside-toplevel - - data = run_rest_request("GET", f"/v2/network-volumes/{_path_segment(volume_id)}") or {} - datacenter = data.get("dataCenter") or data.get("dataCenterId") + from runpod.apps.utils.lookup import find_by_id_or_name # pylint: disable=import-outside-toplevel + + volumes: List[Dict[str, Any]] = [] + params: Dict[str, str] = {} + while True: + data = run_rest_request("GET", "/v2/network-volumes", params=params) or {} + volumes.extend(data.get("networkVolumes") or []) + page = data.get("pagination") or {} + cursor = page.get("nextCursor") + if not page.get("hasNextPage") or not cursor or cursor == params.get("cursor"): + break + params = {"cursor": cursor} + record = find_by_id_or_name( + volumes, reference, noun="network volumes", error=NetworkVolumeTransferError + ) + if record is None: + raise NetworkVolumeTransferError(f"network volume {reference!r} not found") + datacenter = record.get("dataCenter") or record.get("dataCenterId") if not datacenter: raise NetworkVolumeTransferError( - f"could not find the datacenter of network volume {volume_id}" + f"could not find the datacenter of network volume {reference!r}" ) - return str(datacenter) + return record["id"], str(datacenter) def _check_endpoint_resolves(endpoint: str) -> None: @@ -962,14 +977,13 @@ def _check_endpoint_resolves(endpoint: str) -> None: socket.getaddrinfo(host, 443) except socket.gaierror as exc: raise NetworkVolumeTransferError( - f"the s3 endpoint {host} does not resolve: check --datacenter (only " - "datacenters with the s3 api are supported) or your network connection" + f"the s3 endpoint {host} does not resolve: the volume's datacenter " + "may not support the s3 api, or check your network connection" ) from exc def _make_transfer( - volume_id: str, - datacenter: Optional[str], + volume: str, access_key: Optional[str], secret_key: Optional[str], endpoint_url: Optional[str], @@ -987,7 +1001,7 @@ def _make_transfer( f"s3 api keys required: set {ACCESS_KEY_ENV} and {SECRET_KEY_ENV} " f"(create keys at {S3_KEYS_URL})" ) - datacenter = datacenter or _lookup_datacenter(volume_id) + volume_id, datacenter = _resolve_volume(volume) endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" if not endpoint_url: # an explicit endpoint (a local stand-in, a custom route) is the @@ -1021,11 +1035,10 @@ def _make_transfer( def upload( - volume_id: str, + volume: str, local_path, remote_path: Optional[str] = None, *, - datacenter: Optional[str] = None, access_key: Optional[str] = None, secret_key: Optional[str] = None, endpoint_url: Optional[str] = None, @@ -1033,32 +1046,18 @@ def upload( progress: Optional[Progress] = None, force: bool = False, ) -> Dict[str, int]: - """upload a file or folder to a network volume. - - `remote_path` defaults to the local basename. a folder's contents go - under `remote_path/`. for a file, a `remote_path` ending in "/" keeps - the local file name. unchanged files (same size, remote not older, in - whole seconds) are skipped unless `force`, and an interrupted multipart - upload resumes on re-run. - - s3 keys come from the arguments or RUNPOD_S3_ACCESS_KEY and - RUNPOD_S3_SECRET_KEY. without `datacenter`, the volume's datacenter is - looked up with the runpod api key. - - returns {files_transferred, files_skipped, bytes_transferred}. - """ + """backs NetworkVolume.upload, which documents the behaviour.""" transfer = _make_transfer( - volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress, force + volume, access_key, secret_key, endpoint_url, concurrency, progress, force ) return transfer.upload(Path(local_path), remote_path) def download( - volume_id: str, + volume: str, remote_path: str, local_path=None, *, - datacenter: Optional[str] = None, access_key: Optional[str] = None, secret_key: Optional[str] = None, endpoint_url: Optional[str] = None, @@ -1066,22 +1065,9 @@ def download( progress: Optional[Progress] = None, force: bool = False, ) -> Dict[str, int]: - """download a file or folder from a network volume. - - a `remote_path` ending in "/" (or one that only exists as a prefix) is - a folder. `local_path` defaults to the remote basename. local files of - the same size and not older than the remote copy are skipped (downloads - stamp the remote time on the file) unless `force`, and partial downloads - resume on re-run. a `local_path` ending in "/" is a folder to download into. - - returns {files_transferred, files_skipped, bytes_transferred}. - """ + """backs NetworkVolume.download, which documents the behaviour.""" transfer = _make_transfer( - volume_id, datacenter, access_key, secret_key, endpoint_url, concurrency, progress, force + volume, access_key, secret_key, endpoint_url, concurrency, progress, force ) # passed through as given: Path() would drop a trailing "/" return transfer.download(remote_path, local_path) - - -upload_to_network_volume = upload -download_from_network_volume = download diff --git a/runpod/rp_cli/main.py b/runpod/rp_cli/main.py index b51ba185a..4c6936f4c 100644 --- a/runpod/rp_cli/main.py +++ b/runpod/rp_cli/main.py @@ -1086,142 +1086,15 @@ def undeploy(app_name, env_name, yes): ) -# -------------------------------------------------------- network volumes - - -@cli.group(name="network-volume") -def network_volume(): - """Upload and download files on network volumes (resumable). - - Uses the S3-compatible API. Set RUNPOD_S3_ACCESS_KEY and - RUNPOD_S3_SECRET_KEY (create keys at console.runpod.io/user/settings). - Re-run an interrupted command to resume it. Parts of an upload that is - never re-run stay in the volume's hidden .s3compat_uploads folder and - count toward its size. - """ - - -def _transfer_options(func): - options = [ - click.option( - "--datacenter", - default=None, - help="Volume datacenter, e.g. EU-RO-1 (looked up with your API key if omitted).", - ), - click.option("--s3-access-key", default=None, help="S3 API access key."), - click.option("--s3-secret-key", default=None, help="S3 API secret key."), - click.option("--endpoint-url", default=None, hidden=True), - click.option( - "--concurrency", default=4, show_default=True, type=click.IntRange(1, 64), - help="Parallel requests.", - ), - click.option( - "--force", is_flag=True, - help="Transfer every file, even ones the size and time check calls unchanged.", - ), - ] - for option in reversed(options): - func = option(func) - return func - - -def _run_transfer(direction, volume_id, source, target, opts) -> None: - from rich.console import Console - from rich.progress import ( - BarColumn, - DownloadColumn, - Progress, - TextColumn, - TransferSpeedColumn, - ) - - from runpod import network_volume_transfer as nvt - - tasks = {} - with Progress( - TextColumn("{task.description}"), - BarColumn(), - DownloadColumn(), - TransferSpeedColumn(), - console=Console(stderr=True), - transient=True, - ) as bar: - - def progress(event, path, nbytes): - if event == "start": - tasks[path] = bar.add_task(Path(path).name, total=nbytes) - elif event in ("advance", "reuse") and path in tasks: - bar.advance(tasks[path], nbytes) - elif event == "done" and path in tasks: - bar.remove_task(tasks.pop(path)) - bar.console.print(f" {direction} {path}") - - func = nvt.upload if direction == "uploaded" else nvt.download - try: - summary = func( - volume_id, - source, - target, - datacenter=opts["datacenter"], - access_key=opts["s3_access_key"], - secret_key=opts["s3_secret_key"], - endpoint_url=opts["endpoint_url"], - concurrency=opts["concurrency"], - progress=progress, - force=opts["force"], - ) - except KeyboardInterrupt: - click.echo("interrupted; re-run the same command to resume", err=True) - sys.exit(130) - except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly - hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" - raise click.ClickException(f"{exc}{hint}") from exc - - click.echo( - f"{direction} {summary['files_transferred']} file(s), " - f"{summary['bytes_transferred']} bytes; " - f"skipped {summary['files_skipped']} unchanged" - ) - - -@network_volume.command(name="upload") -@click.argument("volume_id") -@click.argument("local_path", type=click.Path(exists=True)) -@click.argument("remote_path", required=False) -@_transfer_options -def network_volume_upload(volume_id, local_path, remote_path, **opts): - """Upload a file or folder to a network volume. - - REMOTE_PATH defaults to the local name. A folder's contents go under - REMOTE_PATH/. Files whose remote copy has the same size and is not older - are skipped; --force sends them anyway. - """ - _run_transfer("uploaded", volume_id, local_path, remote_path, opts) - - -@network_volume.command(name="download") -@click.argument("volume_id") -@click.argument("remote_path") -@click.argument("local_path", required=False) -@_transfer_options -def network_volume_download(volume_id, remote_path, local_path, **opts): - """Download a file or folder from a network volume. - - A REMOTE_PATH ending in / is a folder. LOCAL_PATH defaults to the - remote name. Local files of the same size and not older than the - remote copy are skipped; --force fetches them anyway. - """ - _run_transfer("downloaded", volume_id, remote_path, local_path, opts) - - # -------------------------------------------------------- legacy groups def _mount_groups() -> None: from runpod.cli.groups.pod.commands import pod_cli from runpod.cli.groups.ssh.commands import ssh_cli + from runpod.rp_cli.volume import volume - for command in (ssh_cli, pod_cli): + for command in (ssh_cli, pod_cli, volume): cli.add_command(command) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py new file mode 100644 index 000000000..e5894862d --- /dev/null +++ b/runpod/rp_cli/volume.py @@ -0,0 +1,126 @@ +import sys +from pathlib import Path + +import click + + +@click.group() +def volume(): + """Upload and download files on network volumes (resumable). + + Uses the S3-compatible API. Set RUNPOD_S3_ACCESS_KEY and + RUNPOD_S3_SECRET_KEY (create keys at console.runpod.io/user/settings). + Re-run an interrupted command to resume it. Parts of an upload that is + never re-run stay in the volume's hidden .s3compat_uploads folder and + count toward its size. + """ + + +def _transfer_options(func): + options = [ + click.option("--s3-access-key", default=None, help="S3 API access key."), + click.option("--s3-secret-key", default=None, help="S3 API secret key."), + click.option("--endpoint-url", default=None, hidden=True), + click.option( + "--concurrency", default=4, show_default=True, type=click.IntRange(1, 64), + help="Parallel requests.", + ), + click.option( + "--force", is_flag=True, + help="Transfer every file, even ones the size and time check calls unchanged.", + ), + ] + for option in reversed(options): + func = option(func) + return func + + +def _run_transfer(direction, reference, source, target, opts) -> None: + from rich.console import Console + from rich.progress import ( + BarColumn, + DownloadColumn, + Progress, + TextColumn, + TransferSpeedColumn, + ) + + from runpod import network_volume_transfer as nvt + from runpod.apps import NetworkVolume + + tasks = {} + with Progress( + TextColumn("{task.description}"), + BarColumn(), + DownloadColumn(), + TransferSpeedColumn(), + console=Console(stderr=True), + transient=True, + ) as bar: + + def progress(event, path, nbytes): + if event == "start": + tasks[path] = bar.add_task(Path(path).name, total=nbytes) + elif event in ("advance", "reuse") and path in tasks: + bar.advance(tasks[path], nbytes) + elif event == "done" and path in tasks: + bar.remove_task(tasks.pop(path)) + bar.console.print(f" {direction} {path}") + + target_volume = NetworkVolume(reference, create=False) + func = target_volume.upload if direction == "uploaded" else target_volume.download + try: + summary = func( + source, + target, + access_key=opts["s3_access_key"], + secret_key=opts["s3_secret_key"], + endpoint_url=opts["endpoint_url"], + concurrency=opts["concurrency"], + progress=progress, + force=opts["force"], + ) + except KeyboardInterrupt: + click.echo("interrupted; re-run the same command to resume", err=True) + sys.exit(130) + except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly + hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" + raise click.ClickException(f"{exc}{hint}") from exc + + click.echo( + f"{direction} {summary['files_transferred']} file(s), " + f"{summary['bytes_transferred']} bytes; " + f"skipped {summary['files_skipped']} unchanged" + ) + + +@volume.command(name="upload") +@click.argument("volume_ref", metavar="VOLUME") +@click.argument("local_path", type=click.Path(exists=True)) +@click.argument("remote_path", required=False) +@_transfer_options +def volume_upload(volume_ref, local_path, remote_path, **opts): + """Upload a file or folder to a network volume. + + VOLUME is the volume's id or name. REMOTE_PATH defaults to the local + name. A folder's contents go under REMOTE_PATH/. Files whose remote copy + has the same size and is not older are skipped; --force sends them + anyway. + """ + _run_transfer("uploaded", volume_ref, local_path, remote_path, opts) + + +@volume.command(name="download") +@click.argument("volume_ref", metavar="VOLUME") +@click.argument("remote_path") +@click.argument("local_path", required=False) +@_transfer_options +def volume_download(volume_ref, remote_path, local_path, **opts): + """Download a file or folder from a network volume. + + VOLUME is the volume's id or name. A REMOTE_PATH ending in / is a + folder. LOCAL_PATH defaults to the remote name. Local files of the same + size and not older than the remote copy are skipped; --force fetches + them anyway. + """ + _run_transfer("downloaded", volume_ref, remote_path, local_path, opts) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index db4a6394b..fdb525ee7 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -172,6 +172,16 @@ def _transfer(fake, tmp_path, concurrency=1, events=None): ) +_VOLUME = {"id": "vol123", "name": "models", "dataCenterId": "EU-RO-1"} + + +def _volumes_page(*volumes, cursor=None): + return { + "networkVolumes": list(volumes), + "pagination": {"hasNextPage": cursor is not None, "nextCursor": cursor}, + } + + DATA = bytes(range(40)) # 5 parts / chunks of 8 bytes @@ -611,7 +621,7 @@ def test_make_transfer_requires_s3_keys(monkeypatch): monkeypatch.setenv("AWS_ACCESS_KEY_ID", "aws") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "aws") with pytest.raises(nvt.NetworkVolumeTransferError, match="RUNPOD_S3_ACCESS_KEY"): - nvt.upload("vol", ".", datacenter="EU-RO-1") + nvt.upload("vol", ".") def test_make_transfer_builds_datacenter_client(monkeypatch): @@ -620,11 +630,9 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): monkeypatch.delenv(name, raising=False) monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") - with patch( - "runpod.api.rest.run_rest_request", return_value={"dataCenter": "EU-RO-1"} - ) as rest: - transfer = nvt._make_transfer("vol123", None, None, None, None, 4, None) - rest.assert_called_once_with("GET", "/v2/network-volumes/vol123") + with patch("runpod.api.rest.run_rest_request", return_value=_volumes_page(_VOLUME)) as rest: + transfer = nvt._make_transfer("models", None, None, None, 4, None) + rest.assert_called_once_with("GET", "/v2/network-volumes", params={}) assert transfer.client.meta.endpoint_url == "https://s3api-eu-ro-1.runpod.io/" assert transfer.client.meta.region_name == "EU-RO-1" assert transfer.client.meta.config.request_checksum_calculation == "when_required" @@ -635,38 +643,64 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): assert transfer.bucket == "vol123" -def test_lookup_datacenter_escapes_volume_id(): - with patch("runpod.api.rest.run_rest_request", return_value={"dataCenter": "EU-RO-1"}) as rest: - nvt._lookup_datacenter("a/b?c") - rest.assert_called_once_with("GET", "/v2/network-volumes/a%2Fb%3Fc") +def test_resolve_volume_by_name_across_pages(): + pages = [ + _volumes_page({"id": "other", "name": "x", "dataCenterId": "US-KS-2"}, cursor="c1"), + _volumes_page(_VOLUME), + ] + with patch("runpod.api.rest.run_rest_request", side_effect=pages) as rest: + assert nvt._resolve_volume("models") == ("vol123", "EU-RO-1") + assert rest.call_args_list[1].kwargs["params"] == {"cursor": "c1"} -def test_lookup_datacenter_missing(): - with patch("runpod.api.rest.run_rest_request", return_value={}): - with pytest.raises(nvt.NetworkVolumeTransferError): - nvt._lookup_datacenter("vol") +def test_resolve_volume_by_id(): + with patch("runpod.api.rest.run_rest_request", return_value=_volumes_page(_VOLUME)): + assert nvt._resolve_volume("vol123") == ("vol123", "EU-RO-1") -def test_public_functions_delegate(tmp_path): +@pytest.mark.parametrize( + "volumes, match", + [ + ([], "not found"), + ([{**_VOLUME, "id": "a"}, {**_VOLUME, "id": "b"}], "reference by id"), + ([{"id": "vol123", "name": "models"}], "datacenter"), + ], + ids=["missing", "ambiguous-name", "no-datacenter"], +) +def test_resolve_volume_errors(volumes, match): + with patch("runpod.api.rest.run_rest_request", return_value=_volumes_page(*volumes)): + with pytest.raises(nvt.NetworkVolumeTransferError, match=match): + nvt._resolve_volume("models") + + +def test_network_volume_methods_delegate(): + calls = [] fake_transfer = type( "T", (), {"upload": lambda self, l, r: ("up", l, r), "download": lambda self, r, l: ("down", r, l)} )() - with patch.object(nvt, "_make_transfer", return_value=fake_transfer): - assert nvt.upload("vol", "x", "y")[0] == "up" - assert nvt.download("vol", "y", "x")[0] == "down" + + def make(*args, **kwargs): + calls.append(args) + return fake_transfer + + volume = runpod.NetworkVolume("models", create=False) + with patch.object(nvt, "_make_transfer", side_effect=make): + assert volume.upload("x", "y", concurrency=2)[0] == "up" + assert volume.download("y", "x")[0] == "down" + assert calls[0][0] == "models" and calls[0][4] == 2 -def test_public_download_keeps_trailing_slash(tmp_path, monkeypatch): +def test_network_volume_download_keeps_trailing_slash(tmp_path, monkeypatch): fake = FakeS3() fake.put("w.bin", b"w") monkeypatch.setattr(nvt, "_make_transfer", lambda *a, **k: _transfer(fake, tmp_path)) - nvt.download("vol", "w.bin", str(tmp_path / "out") + "/") + runpod.NetworkVolume("vol").download("w.bin", str(tmp_path / "out") + "/") assert (tmp_path / "out" / "w.bin").read_bytes() == b"w" -def test_lazy_exports(): - assert runpod.upload_to_network_volume is nvt.upload - assert runpod.download_from_network_volume is nvt.download +def test_no_module_level_transfer_exports(): + for name in ("upload_to_network_volume", "download_from_network_volume", "network_volume_transfer"): + assert name not in runpod.__all__ # ------------------------------------------------------------------- cli @@ -678,18 +712,18 @@ def test_cli_upload_prints_summary(tmp_path): src = tmp_path / "f.txt" src.write_text("x") - def fake_upload(volume_id, source, target, **kwargs): + def fake_upload(volume, source, target, **kwargs): progress = kwargs["progress"] progress("start", source, 1) progress("advance", source, 1) progress("done", source, 0) - assert kwargs["datacenter"] == "EU-RO-1" and kwargs["concurrency"] == 2 + assert volume == "models" and kwargs["concurrency"] == 2 return {"files_transferred": 1, "files_skipped": 0, "bytes_transferred": 1} with patch.object(nvt, "upload", side_effect=fake_upload): result = CliRunner().invoke( cli, - ["network-volume", "upload", "vol", str(src), "--datacenter", "EU-RO-1", "--concurrency", "2"], + ["volume", "upload", "models", str(src), "--concurrency", "2"], ) assert result.exit_code == 0, result.output assert "uploaded 1 file(s), 1 bytes; skipped 0 unchanged" in result.output @@ -699,7 +733,7 @@ def test_cli_download_failure_says_resume(): from runpod.rp_cli.main import cli with patch.object(nvt, "download", side_effect=ConnectionResetError("boom")): - result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b"]) + result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code != 0 assert "boom" in result.output assert "re-run the same command to resume" in result.output @@ -709,7 +743,7 @@ def test_cli_ctrl_c_exits_130_with_resume_hint(): from runpod.rp_cli.main import cli with patch.object(nvt, "upload", side_effect=KeyboardInterrupt): - result = CliRunner().invoke(cli, ["network-volume", "upload", "vol", "."]) + result = CliRunner().invoke(cli, ["volume", "upload", "vol", "."]) assert result.exit_code == 130 assert "interrupted; re-run the same command to resume" in result.output @@ -717,11 +751,11 @@ def test_cli_ctrl_c_exits_130_with_resume_hint(): def test_cli_unresolvable_endpoint_has_no_resume_hint(): from runpod.rp_cli.main import cli - err = nvt.NetworkVolumeTransferError("the s3 endpoint s3api-x.invalid does not resolve: check --datacenter") + err = nvt.NetworkVolumeTransferError("the s3 endpoint s3api-x.invalid does not resolve") with patch.object(nvt, "download", side_effect=err): - result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b", "--datacenter", "X"]) + result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code == 1 - assert "check --datacenter" in result.output + assert "does not resolve" in result.output assert "re-run" not in result.output @@ -729,7 +763,7 @@ def test_cli_missing_path_has_no_resume_hint(): from runpod.rp_cli.main import cli with patch.object(nvt, "download", side_effect=FileNotFoundError("no such file")): - result = CliRunner().invoke(cli, ["network-volume", "download", "vol", "a/b"]) + result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code != 0 assert "no such file" in result.output assert "resume" not in result.output @@ -931,7 +965,7 @@ def test_unresolvable_endpoint_fails_before_the_transfer(endpoint_check): import socket as _socket endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") - with pytest.raises(nvt.NetworkVolumeTransferError, match="--datacenter"): + with pytest.raises(nvt.NetworkVolumeTransferError, match="does not resolve"): nvt._check_endpoint_resolves("https://s3api-nope-1.runpod.io/") endpoint_check["error"] = None nvt._check_endpoint_resolves("https://s3api-eu-ro-1.runpod.io/") From b9c8f926abff816455aa5b70c99409d0211df230 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 18:47:59 -0400 Subject: [PATCH 12/23] fix: list volumes through the shared paginated rest helper also pin the explicit-endpoint branch of the dns check with a test and keep system proxy settings out of the endpoint_check fixture. --- runpod/network_volume_transfer.py | 14 ++----- tests/test_network_volume_transfer.py | 56 +++++++++++++++++++-------- 2 files changed, 42 insertions(+), 28 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 189f7a1e2..c6dc836c8 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -932,19 +932,11 @@ def _safe_target(root: Path, rel: str) -> Path: def _resolve_volume(reference: str) -> Tuple[str, str]: """the id and datacenter of the network volume with this id or name.""" - from runpod.api.rest import run_rest_request # pylint: disable=import-outside-toplevel + from runpod.api.rest import run_rest_paginated_request_async # pylint: disable=import-outside-toplevel + from runpod.apps.context import block # pylint: disable=import-outside-toplevel from runpod.apps.utils.lookup import find_by_id_or_name # pylint: disable=import-outside-toplevel - volumes: List[Dict[str, Any]] = [] - params: Dict[str, str] = {} - while True: - data = run_rest_request("GET", "/v2/network-volumes", params=params) or {} - volumes.extend(data.get("networkVolumes") or []) - page = data.get("pagination") or {} - cursor = page.get("nextCursor") - if not page.get("hasNextPage") or not cursor or cursor == params.get("cursor"): - break - params = {"cursor": cursor} + volumes = block(run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes")) record = find_by_id_or_name( volumes, reference, noun="network volumes", error=NetworkVolumeTransferError ) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index fdb525ee7..dd253a6be 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -175,11 +175,15 @@ def _transfer(fake, tmp_path, concurrency=1, events=None): _VOLUME = {"id": "vol123", "name": "models", "dataCenterId": "EU-RO-1"} -def _volumes_page(*volumes, cursor=None): - return { - "networkVolumes": list(volumes), - "pagination": {"hasNextPage": cursor is not None, "nextCursor": cursor}, - } +def _list_volumes(*volumes): + """patch the paginated volume listing to return these records.""" + calls = [] + + async def fake(path, collection, **kwargs): + calls.append((path, collection)) + return list(volumes) + + return patch("runpod.api.rest.run_rest_paginated_request_async", fake), calls DATA = bytes(range(40)) # 5 parts / chunks of 8 bytes @@ -630,9 +634,10 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): monkeypatch.delenv(name, raising=False) monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") - with patch("runpod.api.rest.run_rest_request", return_value=_volumes_page(_VOLUME)) as rest: + listing, calls = _list_volumes(_VOLUME) + with listing: transfer = nvt._make_transfer("models", None, None, None, 4, None) - rest.assert_called_once_with("GET", "/v2/network-volumes", params={}) + assert calls == [("/v2/network-volumes", "networkVolumes")] assert transfer.client.meta.endpoint_url == "https://s3api-eu-ro-1.runpod.io/" assert transfer.client.meta.region_name == "EU-RO-1" assert transfer.client.meta.config.request_checksum_calculation == "when_required" @@ -643,18 +648,15 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): assert transfer.bucket == "vol123" -def test_resolve_volume_by_name_across_pages(): - pages = [ - _volumes_page({"id": "other", "name": "x", "dataCenterId": "US-KS-2"}, cursor="c1"), - _volumes_page(_VOLUME), - ] - with patch("runpod.api.rest.run_rest_request", side_effect=pages) as rest: +def test_resolve_volume_by_name(): + listing, _ = _list_volumes({"id": "other", "name": "x", "dataCenterId": "US-KS-2"}, _VOLUME) + with listing: assert nvt._resolve_volume("models") == ("vol123", "EU-RO-1") - assert rest.call_args_list[1].kwargs["params"] == {"cursor": "c1"} def test_resolve_volume_by_id(): - with patch("runpod.api.rest.run_rest_request", return_value=_volumes_page(_VOLUME)): + listing, _ = _list_volumes(_VOLUME) + with listing: assert nvt._resolve_volume("vol123") == ("vol123", "EU-RO-1") @@ -668,7 +670,8 @@ def test_resolve_volume_by_id(): ids=["missing", "ambiguous-name", "no-datacenter"], ) def test_resolve_volume_errors(volumes, match): - with patch("runpod.api.rest.run_rest_request", return_value=_volumes_page(*volumes)): + listing, _ = _list_volumes(*volumes) + with listing: with pytest.raises(nvt.NetworkVolumeTransferError, match=match): nvt._resolve_volume("models") @@ -946,10 +949,14 @@ def endpoint_check(monkeypatch): """the real endpoint check with dns and proxies stubbed (the autouse fixture replaces the check itself).""" import socket as _socket + import urllib.request monkeypatch.setattr(nvt, "_check_endpoint_resolves", REAL_ENDPOINT_CHECK) for name in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "NO_PROXY", "no_proxy"): monkeypatch.delenv(name, raising=False) + # proxies from the environment only: macos and windows also read system settings + monkeypatch.setattr(urllib.request, "getproxies", urllib.request.getproxies_environment) + monkeypatch.setattr(urllib.request, "proxy_bypass", urllib.request.proxy_bypass_environment) state = {"error": None} def fake_getaddrinfo(host, port, *args, **kwargs): @@ -979,6 +986,19 @@ def test_endpoint_check_skipped_behind_a_proxy(endpoint_check, monkeypatch): nvt._check_endpoint_resolves("https://s3api-eu-ro-1.runpod.io/") +def test_explicit_endpoint_skips_the_dns_check(endpoint_check, monkeypatch): + import socket as _socket + + monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") + monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") + monkeypatch.setattr(nvt, "_resolve_volume", lambda volume: ("vol123", "NOPE-1")) + endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") + transfer = nvt._make_transfer("vol123", None, None, "http://localhost:9000", 4, None) + assert transfer.client.meta.endpoint_url == "http://localhost:9000" + with pytest.raises(nvt.NetworkVolumeTransferError, match="does not resolve"): + nvt._make_transfer("vol123", None, None, None, 4, None) + + def test_dns_failure_mid_transfer_is_a_network_drop(): import socket as _socket @@ -1185,7 +1205,9 @@ def test_ctrl_c_stops_files_in_flight(tmp_path): for i in range(6): (folder / f"f{i}.bin").write_bytes(DATA + bytes([i])) calls = [] - lock = __import__("threading").Lock() + import threading + + lock = threading.Lock() original = fake.upload_part transfer = _transfer(fake, tmp_path, concurrency=2) From b085cb10eaf7ffbc382cd5952a102d5803af96cb Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 19:04:26 -0400 Subject: [PATCH 13/23] refactor: split transfer methods and fix two folder-download edge cases _Transfer's long upload and download methods are split into one step each, with _Upload and _Download holding a file's transfer state. the module-level upload/download pass-throughs are gone; NetworkVolume calls _make_transfer directly. fixes: the sidecar's .tmp name now counts as a download name collision, and a folder failure is resumable only when every file failure is. --- runpod/apps/volume.py | 27 +- runpod/network_volume_transfer.py | 901 ++++++++++++++------------ runpod/rp_cli/volume.py | 103 +-- tests/test_network_volume_transfer.py | 42 +- 4 files changed, 595 insertions(+), 478 deletions(-) diff --git a/runpod/apps/volume.py b/runpod/apps/volume.py index 888463200..4635c9df6 100644 --- a/runpod/apps/volume.py +++ b/runpod/apps/volume.py @@ -114,17 +114,10 @@ def upload( """ from runpod import network_volume_transfer as nvt - return nvt.upload( - self.reference, - local_path, - remote_path, - access_key=access_key, - secret_key=secret_key, - endpoint_url=endpoint_url, - concurrency=concurrency, - progress=progress, - force=force, + transfer = nvt._make_transfer( + self.reference, access_key, secret_key, endpoint_url, concurrency, progress, force ) + return transfer.upload(Path(local_path), remote_path) def download( self, @@ -151,17 +144,11 @@ def download( """ from runpod import network_volume_transfer as nvt - return nvt.download( - self.reference, - remote_path, - local_path, - access_key=access_key, - secret_key=secret_key, - endpoint_url=endpoint_url, - concurrency=concurrency, - progress=progress, - force=force, + transfer = nvt._make_transfer( + self.reference, access_key, secret_key, endpoint_url, concurrency, progress, force ) + # passed through as given: Path() would drop a trailing "/" + return transfer.download(remote_path, local_path) class GlobalVolume(Volume): diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index c6dc836c8..0b657a07f 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -37,6 +37,8 @@ UPLOADS_DIR = ".s3compat_uploads" PARTIAL_SUFFIX = ".rpdownload" SIDECAR_SUFFIX = ".rpdownload.json" +# every file a download writes next to its target +TRANSFER_SUFFIXES = (PARTIAL_SUFFIX, SIDECAR_SUFFIX, SIDECAR_SUFFIX + ".tmp") ACCESS_KEY_ENV = "RUNPOD_S3_ACCESS_KEY" SECRET_KEY_ENV = "RUNPOD_S3_SECRET_KEY" @@ -275,7 +277,7 @@ def _run_pool( def _is_download_leftover(path: Path) -> bool: """a partial download, its sidecar, or the sidecar's temp file.""" - return path.name.endswith((PARTIAL_SUFFIX, SIDECAR_SUFFIX, SIDECAR_SUFFIX + ".tmp")) + return path.name.endswith(TRANSFER_SUFFIXES) def _same_second_or_later(a: float, b: float) -> bool: @@ -296,8 +298,16 @@ def _strip_etag(etag: Any) -> str: # ------------------------------------------------------------- transfers +# the summary counter each progress event advances +_SUMMARY_FIELDS = { + "advance": "bytes_transferred", + "done": "files_transferred", + "skip": "files_skipped", +} + + class _Transfer: - """one bucket on one endpoint. the public functions build this.""" + """one bucket on one endpoint. _make_transfer builds this.""" def __init__( self, @@ -343,13 +353,10 @@ def _pool(self, tasks: List[Callable[[], Any]]) -> List[Any]: return _run_pool(tasks, self.concurrency, self._stop) def _emit(self, event: str, path: str, nbytes: int = 0) -> None: - with self._lock: - if event == "advance": - self.summary["bytes_transferred"] += nbytes - elif event == "done": - self.summary["files_transferred"] += 1 - elif event == "skip": - self.summary["files_skipped"] += 1 + field = _SUMMARY_FIELDS.get(event) + if field is not None: + with self._lock: + self.summary[field] += nbytes if event == "advance" else 1 if self._progress is not None: self._progress(event, path, nbytes) @@ -424,172 +431,210 @@ def _list_parts(self, key: str, upload_id: str) -> List[Dict[str, Any]]: # -- upload + def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: + _check_remote_path(remote_path) + if not local.exists(): + raise FileNotFoundError(str(local)) + if local.is_dir(): + self._upload_folder(local, remote_path) + else: + self._upload_single_file(local, remote_path) + return self.summary + + def _upload_folder(self, local: Path, remote_path: Optional[str]) -> None: + prefix = (local.resolve().name if remote_path is None else remote_path).strip("/") + remote = {obj["Key"]: obj for obj in self.list_prefix(f"{prefix}/" if prefix else "")} + jobs = [] + for path in sorted(p for p in local.rglob("*") if p.is_file()): + key = _join_key(prefix, path.relative_to(local).as_posix()) + if not self._skip_upload(path, remote.get(key)): + jobs.append((path, key)) + self._run_files([lambda p=p, k=k: self.upload_file(p, k) for p, k in jobs], jobs) + + def _skip_upload(self, path: Path, remote: Optional[Dict[str, Any]]) -> bool: + """report a folder file that needs no upload: an interrupted + download's partial file (not data), or an unchanged file.""" + leftover = _is_download_leftover(path) + if not leftover and (self.force or not self._unchanged_remote(remote, path)): + return False + self._emit("skip", str(path), path.stat().st_size) + return True + + def _upload_single_file(self, local: Path, remote_path: Optional[str]) -> None: + key = _file_key(local, remote_path) + remote = self._head_upload_target(key) + if not self.force and self._unchanged_remote(remote, local): + self._emit("skip", str(local), local.stat().st_size) + return + self.upload_file(local, key) + + def _head_upload_target(self, key: str) -> Optional[Dict[str, Any]]: + try: + return self.head(key) + except Exception as exc: # noqa: BLE001 + # runpod answers a head on a directory with 400, but a head has no + # body, so a 400 for any other reason looks the same: confirm + if _error_code(exc)[1] == 400 and self._is_folder(key): + raise NetworkVolumeTransferError( + f"{key!r} is a folder on the volume; use {key + '/'!r} to upload into it" + ) from exc + raise + + def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bool: + if not remote: + return False + stat = local.stat() + size = remote.get("Size", remote.get("ContentLength")) + modified = remote.get("LastModified") + if size is None or int(size) != stat.st_size or modified is None: + return False + return _same_second_or_later(modified.timestamp(), stat.st_mtime) + def upload_file(self, local: Path, key: str) -> None: stat = local.stat() - size = stat.st_size label = str(local) - self._emit("start", label, size) - if size <= SINGLE_PUT_MAX: - self._retry( - lambda: self._request( - lambda: self.client.put_object( - Bucket=self.bucket, Key=key, Body=local.read_bytes() - ) - ) - ) - # a retry rereads the file; if it changed meanwhile, the object - # may hold either version and is newer than the local edit, so a - # re-run would skip it. remove it so the re-run sends the file. - try: - now = local.stat() - except OSError as exc: - # gone or unreadable: the object is what was read, keep it - raise NetworkVolumeTransferError(f"{local} is gone after its upload: {exc}") from exc - if now.st_size != size or now.st_mtime_ns != stat.st_mtime_ns: - try: - # not via self._retry: this cleanup runs even after ctrl-c - _with_retries(lambda: self.client.delete_object(Bucket=self.bucket, Key=key)) - except Exception as exc: # noqa: BLE001 - raise NetworkVolumeTransferError( - f"{local} changed during upload and the remote copy could not be " - f"removed ({exc}); re-run with --force" - ) from exc - raise NetworkVolumeTransferError( - f"{local} changed during upload", resumable=True - ) - self._emit("advance", label, size) + self._emit("start", label, stat.st_size) + if stat.st_size <= SINGLE_PUT_MAX: + self._single_put(local, key, stat) else: - self._multipart_upload(local, key, size, stat.st_mtime_ns, label) + self._multipart_upload(_Upload(local, key, stat, part_size_for(stat.st_size))) self._emit("done", label) - def _resume_parts( - self, local: Path, key: str, state: Dict[str, Any], size: int, part_size: int - ) -> Dict[int, str]: + def _single_put(self, local: Path, key: str, stat: os.stat_result) -> None: + self._retry( + lambda: self._request( + lambda: self.client.put_object(Bucket=self.bucket, Key=key, Body=local.read_bytes()) + ) + ) + self._delete_if_changed(local, key, stat) + self._emit("advance", str(local), stat.st_size) + + def _delete_if_changed(self, local: Path, key: str, stat: os.stat_result) -> None: + """a retry rereads the file; if it changed meanwhile, the object may + hold either version and is newer than the local edit, so a re-run + would skip it. remove it so the re-run sends the file.""" + try: + now = local.stat() + except OSError as exc: + # gone or unreadable: the object is what was read, keep it + raise NetworkVolumeTransferError(f"{local} is gone after its upload: {exc}") from exc + if _same_version(stat, now): + return + try: + # not via self._retry: this cleanup runs even after ctrl-c + _with_retries(lambda: self.client.delete_object(Bucket=self.bucket, Key=key)) + except Exception as exc: # noqa: BLE001 + raise NetworkVolumeTransferError( + f"{local} changed during upload and the remote copy could not be " + f"removed ({exc}); re-run with --force" + ) from exc + raise NetworkVolumeTransferError(f"{local} changed during upload", resumable=True) + + def _multipart_upload(self, up: "_Upload") -> None: + up.state_path = self._state_path(up.key, up.local) + if not self._resume_upload(up): + self._create_upload(up) + reused = sum(up.part_length(n) for n in up.parts) + if reused: + self._emit("reuse", up.label, reused) + self._send_missing_parts(up) + # the parts were read over the whole upload; if the file moved under + # them, they may mix two versions of it + if not _same_version(up.stat, up.local.stat()): + self._abort(up.key, up.upload_id) + _remove(up.state_path) + raise NetworkVolumeTransferError(f"{up.local} changed during upload", resumable=True) + self._complete_upload(up) + _remove(up.state_path) + + def _resume_upload(self, up: "_Upload") -> bool: + """pick up the saved upload for this file with its verified parts. + false when there is none to resume.""" + state = _read_json(up.state_path) + if not state or not state.get("upload_id"): + return False + if not up.matches(state): + # the local file changed since the upload began: its parts are stale + self._abort(up.key, state["upload_id"]) + _remove(up.state_path) + return False + try: + up.parts = self._verified_parts(up, state["upload_id"]) + except LookupError: + # the server no longer has this upload: start a new one + _remove(up.state_path) + return False + up.upload_id = state["upload_id"] + return True + + def _verified_parts(self, up: "_Upload", upload_id: str) -> Dict[int, str]: """parts of the saved upload that match the local file.""" - upload_id = state["upload_id"] try: - listed = self._list_parts(key, upload_id) + listed = self._list_parts(up.key, upload_id) except Exception as exc: # noqa: BLE001 code, _ = _error_code(exc) if code == "NoSuchUpload" or _is_not_found(exc): raise LookupError(upload_id) from exc raise - done: Dict[int, str] = {} - for part in listed: - number = int(part["PartNumber"]) - offset = (number - 1) * part_size - expected = min(part_size, size - offset) - if expected <= 0 or int(part.get("Size", -1)) != expected: - continue - # a part whose content cannot be verified is uploaded again - etag = _strip_etag(part.get("ETag")) - if not _is_md5_etag(etag) or _md5_hex(local, offset, expected) != etag.lower(): - continue - done[number] = part.get("ETag") - return done - - def _multipart_upload( - self, local: Path, key: str, size: int, mtime_ns: int, label: str - ) -> None: - part_size = part_size_for(size) - count = math.ceil(size / part_size) - state_path = self._state_path(key, local) - state = _read_json(state_path) - upload_id: Optional[str] = None - done: Dict[int, str] = {} - - if state and state.get("upload_id"): - same = ( - state.get("file_size") == size - and state.get("mtime_ns") == mtime_ns - and state.get("part_size") == part_size + return { + int(part["PartNumber"]): part.get("ETag") + for part in listed + if up.part_matches(part) + } + + def _create_upload(self, up: "_Upload") -> None: + created = self._retry( + lambda: self._request( + lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=up.key) ) - if same: - try: - done = self._resume_parts(local, key, state, size, part_size) - upload_id = state["upload_id"] - except LookupError: - # the server no longer has this upload: start a new one - upload_id = None - else: - self._abort(key, state["upload_id"]) - if upload_id is None: - _remove(state_path) - - if upload_id is None: - created = self._retry( - lambda: self._request( - lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=key) - ) - ) - upload_id = created["UploadId"] - try: - _write_json_atomic( - state_path, - { - "upload_id": upload_id, - "part_size": part_size, - "file_size": size, - "mtime_ns": mtime_ns, - "bucket": self.bucket, - "key": key, - }, - ) - except OSError: - # without saved state no re-run can find this upload again - self._abort(key, upload_id) - raise - - already = sum(min(part_size, size - (n - 1) * part_size) for n in done) - if already: - self._emit("reuse", label, already) - - def upload_part(number: int) -> Tuple[int, str]: - offset = (number - 1) * part_size - length = min(part_size, size - offset) + ) + up.upload_id = created["UploadId"] + up.parts = {} + try: + _write_json_atomic(up.state_path, up.state(self.bucket)) + except OSError: + # without saved state no re-run can find this upload again + self._abort(up.key, up.upload_id) + raise - def send(): - body = _read_range(local, offset, length) - return self.client.upload_part( + def _send_missing_parts(self, up: "_Upload") -> None: + todo = [n for n in range(1, up.count + 1) if n not in up.parts] + for number, etag in self._pool([lambda n=n: self._send_part(up, n) for n in todo]): + up.parts[number] = etag + + def _send_part(self, up: "_Upload", number: int) -> Tuple[int, str]: + offset = (number - 1) * up.part_size + length = up.part_length(number) + response = self._retry( + lambda: self._request( + lambda: self.client.upload_part( Bucket=self.bucket, - Key=key, - UploadId=upload_id, + Key=up.key, + UploadId=up.upload_id, PartNumber=number, - Body=body, + Body=_read_range(up.local, offset, length), ) + ) + ) + self._emit("advance", up.label, length) + return number, response["ETag"] - response = self._retry(lambda: self._request(send)) - self._emit("advance", label, length) - return number, response["ETag"] - - todo = [n for n in range(1, count + 1) if n not in done] - for number, etag in self._pool([lambda n=n: upload_part(n) for n in todo]): - done[number] = etag - - # the parts were read over the whole upload; if the file moved under - # them, they may mix two versions of it - now = local.stat() - if now.st_size != size or now.st_mtime_ns != mtime_ns: - self._abort(key, upload_id) - _remove(state_path) - raise NetworkVolumeTransferError(f"{local} changed during upload", resumable=True) - - parts = [{"PartNumber": n, "ETag": done[n]} for n in sorted(done)] + def _complete_upload(self, up: "_Upload") -> None: + parts = [{"PartNumber": n, "ETag": up.parts[n]} for n in sorted(up.parts)] try: self._retry( lambda: self._request( lambda: self.client.complete_multipart_upload( Bucket=self.bucket, - Key=key, - UploadId=upload_id, + Key=up.key, + UploadId=up.upload_id, MultipartUpload={"Parts": parts}, ) ) ) except Exception as exc: # noqa: BLE001 - re-raised unless it completed - if not self._completed_anyway(exc, key, size, _multipart_etag(parts)): + if not self._completed_anyway(exc, up.key, up.stat.st_size, _multipart_etag(parts)): raise - _remove(state_path) def _completed_anyway(self, exc: Exception, key: str, size: int, want_etag: str) -> bool: """whether a failed complete took effect anyway. @@ -620,65 +665,9 @@ def _abort(self, key: str, upload_id: str) -> None: # a failed abort only leaves orphaned parts behind return - def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bool: - if not remote: - return False - stat = local.stat() - size = remote.get("Size", remote.get("ContentLength")) - modified = remote.get("LastModified") - if size is None or int(size) != stat.st_size or modified is None: - return False - return _same_second_or_later(modified.timestamp(), stat.st_mtime) - - def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: - _check_remote_path(remote_path) - if not local.exists(): - raise FileNotFoundError(str(local)) - if local.is_dir(): - prefix = (local.resolve().name if remote_path is None else remote_path).strip("/") - remote = { - obj["Key"]: obj - for obj in self.list_prefix(f"{prefix}/" if prefix else "") - } - jobs = [] - for path in sorted(p for p in local.rglob("*") if p.is_file()): - if _is_download_leftover(path): - # an interrupted download's partial file is not data - self._emit("skip", str(path), path.stat().st_size) - continue - rel = path.relative_to(local).as_posix() - key = f"{prefix}/{rel}" if prefix else rel - if not self.force and self._unchanged_remote(remote.get(key), path): - self._emit("skip", str(path), path.stat().st_size) - continue - jobs.append((path, key)) - self._run_files([lambda p=p, k=k: self.upload_file(p, k) for p, k in jobs], jobs) - return self.summary - - if remote_path is None: - key = local.name - elif remote_path.endswith("/"): - key = remote_path.lstrip("/") + local.name - else: - key = remote_path.lstrip("/") - try: - remote = self.head(key) - except Exception as exc: # noqa: BLE001 - # runpod answers a head on a directory with 400, but a head has no - # body, so a 400 for any other reason looks the same: confirm - if _error_code(exc)[1] == 400 and self._is_folder(key): - raise NetworkVolumeTransferError( - f"{key!r} is a folder on the volume; use {key + '/'!r} to upload into it" - ) from exc - raise - if not self.force and self._unchanged_remote(remote, local): - self._emit("skip", str(local), local.stat().st_size) - return self.summary - self.upload_file(local, key) - return self.summary - def _run_files(self, tasks: List[Callable[[], Any]], jobs: List[Tuple[Any, Any]]) -> None: - """run per-file tasks. every file is attempted; failures are reported together.""" + """run per-file tasks. every file is attempted; failures are reported + together, resumable only when every one of them is.""" failures = [] resumable = [] @@ -694,152 +683,35 @@ def guarded(index: int) -> None: if failures: raise NetworkVolumeTransferError( f"{len(failures)} file(s) failed:\n " + "\n ".join(sorted(failures)), - resumable=any(resumable), + resumable=all(resumable), ) # -- download - def download_file( - self, key: str, size: int, etag: str, dest: Path, modified: Any = None - ) -> None: - label = str(dest) - # downloads stamp the remote time, so an unchanged file is the same - # size and not older than the remote copy - if ( - not self.force - and dest.is_file() - and dest.stat().st_size == size - and ( - modified is None - or _same_second_or_later(dest.stat().st_mtime, modified.timestamp()) - ) - ): - self._emit("skip", label, size) - return - self._emit("start", label, size) - dest.parent.mkdir(parents=True, exist_ok=True) - tmp = dest.with_name(dest.name + PARTIAL_SUFFIX) - sidecar = dest.with_name(dest.name + SIDECAR_SUFFIX) - state = _read_json(sidecar) - done: set = set() - # the etag is only compared for equality. a multipart object's etag - # is "-", never the md5 of the file. - if ( - state - and state.get("key") == key - and state.get("size") == size - and state.get("etag") == etag - and state.get("chunk_size") == DOWNLOAD_CHUNK_SIZE - and tmp.is_file() - and tmp.stat().st_size == size - ): - done = {int(i) for i in state.get("done", [])} - else: - with open(tmp, "wb") as handle: - handle.truncate(size) - _write_json_atomic( - sidecar, - { - "key": key, - "size": size, - "etag": etag, - "chunk_size": DOWNLOAD_CHUNK_SIZE, - "done": [], - }, - ) - - count = math.ceil(size / DOWNLOAD_CHUNK_SIZE) - already = sum( - min(DOWNLOAD_CHUNK_SIZE, size - i * DOWNLOAD_CHUNK_SIZE) for i in done if i < count - ) - if already: - self._emit("reuse", label, already) - - def fetch(index: int) -> None: - start = index * DOWNLOAD_CHUNK_SIZE - end = min(start + DOWNLOAD_CHUNK_SIZE, size) - 1 - - def get(): - response = self.client.get_object( - Bucket=self.bucket, Key=key, Range=f"bytes={start}-{end}" - ) - got = _strip_etag(response.get("ETag")) - if etag and got and got != etag: - raise NetworkVolumeTransferError( - f"{key} changed on the volume during download; " - "run the command again to restart it", - resumable=True, - ) - data = response["Body"].read() - if len(data) != end - start + 1: - raise _Retryable(f"short read for {key} bytes {start}-{end}") - return data - - data = self._retry(lambda: self._request(get)) - with open(tmp, "r+b") as handle: - handle.seek(start) - handle.write(data) - handle.flush() - os.fsync(handle.fileno()) - with self._lock: - done.add(index) - _write_json_atomic( - sidecar, - { - "key": key, - "size": size, - "etag": etag, - "chunk_size": DOWNLOAD_CHUNK_SIZE, - "done": sorted(done), - }, - ) - self._emit("advance", label, end - start + 1) - - todo = [i for i in range(count) if i not in done] - self._pool([lambda i=i: fetch(i) for i in todo]) - - if tmp.stat().st_size != size: - raise NetworkVolumeTransferError( - f"{dest}: size mismatch after download", resumable=True - ) - os.replace(tmp, dest) - if modified is not None: - # stamp the remote time so uploading the file back counts as unchanged - os.utime(dest, (time.time(), modified.timestamp())) - _remove(sidecar) - self._emit("done", label) - def download(self, remote_path: str, local_path: Any = None) -> Dict[str, int]: """`local_path` is a str or Path; a str ending in "/" is a folder to download into (a Path cannot carry the slash).""" _check_remote_path(remote_path) key = remote_path.lstrip("/") - head = ( - None if key == "" or key.endswith("/") else self.head(key, directory_ok=True) - ) + head = None if key == "" or key.endswith("/") else self.head(key, directory_ok=True) if head is not None: - name = PurePosixPath(key).name - dest = Path(local_path) if local_path is not None else Path(name) - # a trailing slash names a folder to download into, as on upload - if dest.is_dir() or str(local_path or "").endswith(("/", os.sep)): - dest = dest / name - self.download_file( - key, - int(head["ContentLength"]), - _strip_etag(head.get("ETag")), - dest, - head.get("LastModified"), - ) - return self.summary + self._download_object(key, head, local_path) + else: + self._download_prefix(key.rstrip("/"), remote_path, local_path) + return self.summary + + def _download_object(self, key: str, head: Dict[str, Any], local_path: Any) -> None: + self.download_file( + key, + int(head["ContentLength"]), + _strip_etag(head.get("ETag")), + _object_dest(key, local_path), + head.get("LastModified"), + ) - prefix = key.rstrip("/") + def _download_prefix(self, prefix: str, remote_path: str, local_path: Any) -> None: list_prefix = f"{prefix}/" if prefix else "" - objects = [ - obj - for obj in self.list_prefix(list_prefix) - if not obj["Key"].endswith("/") - and UPLOADS_DIR not in PurePosixPath(obj["Key"]).parts - ] + objects = self._listed_files(list_prefix) # a typo must not look like an empty folder; the whole volume may be empty if not objects and prefix: raise FileNotFoundError(f"no such file or directory on the volume: {remote_path}") @@ -857,7 +729,224 @@ def download(self, remote_path: str, local_path: Any = None) -> Dict[str, int]: ], jobs, ) - return self.summary + + def _listed_files(self, list_prefix: str) -> List[Dict[str, Any]]: + """the files under a prefix, without folder markers and in-progress + multipart parts.""" + return [ + obj + for obj in self.list_prefix(list_prefix) + if not obj["Key"].endswith("/") + and UPLOADS_DIR not in PurePosixPath(obj["Key"]).parts + ] + + def download_file( + self, key: str, size: int, etag: str, dest: Path, modified: Any = None + ) -> None: + job = _Download(key, size, etag, dest, modified) + if self._unchanged_local(job): + self._emit("skip", job.label, size) + return + self._emit("start", job.label, size) + dest.parent.mkdir(parents=True, exist_ok=True) + job.completed_chunks = job.load_or_init() + reused = sum(job.chunk_length(i) for i in job.completed_chunks) + if reused: + self._emit("reuse", job.label, reused) + todo = [i for i in range(job.count) if i not in job.completed_chunks] + self._pool([lambda i=i: self._fetch_chunk(job, i) for i in todo]) + self._finish_download(job) + + def _unchanged_local(self, job: "_Download") -> bool: + """downloads stamp the remote time, so an unchanged file is the same + size and not older than the remote copy.""" + if self.force or not job.dest.is_file(): + return False + stat = job.dest.stat() + if stat.st_size != job.size: + return False + return job.modified is None or _same_second_or_later( + stat.st_mtime, job.modified.timestamp() + ) + + def _fetch_chunk(self, job: "_Download", index: int) -> None: + start = index * DOWNLOAD_CHUNK_SIZE + length = job.chunk_length(index) + data = self._retry(lambda: self._request(lambda: self._get_range(job, start, length))) + _write_at(job.tmp, start, data) + with self._lock: + job.completed_chunks.add(index) + job.save_state() + self._emit("advance", job.label, length) + + def _get_range(self, job: "_Download", start: int, length: int) -> bytes: + end = start + length - 1 + response = self.client.get_object( + Bucket=self.bucket, Key=job.key, Range=f"bytes={start}-{end}" + ) + got = _strip_etag(response.get("ETag")) + if job.etag and got and got != job.etag: + raise NetworkVolumeTransferError( + f"{job.key} changed on the volume during download; " + "run the command again to restart it", + resumable=True, + ) + data = response["Body"].read() + if len(data) != length: + raise _Retryable(f"short read for {job.key} bytes {start}-{end}") + return data + + def _finish_download(self, job: "_Download") -> None: + if job.tmp.stat().st_size != job.size: + raise NetworkVolumeTransferError( + f"{job.dest}: size mismatch after download", resumable=True + ) + os.replace(job.tmp, job.dest) + if job.modified is not None: + # stamp the remote time so uploading the file back counts as unchanged + os.utime(job.dest, (time.time(), job.modified.timestamp())) + _remove(job.sidecar) + self._emit("done", job.label) + + +class _Upload: + """one multipart upload of a local file, pinned to its size and mtime.""" + + def __init__(self, local: Path, key: str, stat: os.stat_result, part_size: int): + self.local = local + self.key = key + self.stat = stat + self.part_size = part_size + self.count = math.ceil(stat.st_size / part_size) + self.label = str(local) + self.state_path = Path() + self.upload_id = "" + self.parts: Dict[int, str] = {} # part number -> etag + + def part_length(self, number: int) -> int: + return min(self.part_size, self.stat.st_size - (number - 1) * self.part_size) + + def matches(self, state: Dict[str, Any]) -> bool: + """whether a saved state belongs to this version of the file.""" + return ( + state.get("file_size") == self.stat.st_size + and state.get("mtime_ns") == self.stat.st_mtime_ns + and state.get("part_size") == self.part_size + ) + + def part_matches(self, part: Dict[str, Any]) -> bool: + """whether a listed part holds this file's bytes. a part whose + content cannot be verified is uploaded again.""" + number = int(part["PartNumber"]) + expected = self.part_length(number) + if expected <= 0 or int(part.get("Size", -1)) != expected: + return False + etag = _strip_etag(part.get("ETag")) + offset = (number - 1) * self.part_size + return _is_md5_etag(etag) and _md5_hex(self.local, offset, expected) == etag.lower() + + def state(self, bucket: str) -> Dict[str, Any]: + return { + "upload_id": self.upload_id, + "part_size": self.part_size, + "file_size": self.stat.st_size, + "mtime_ns": self.stat.st_mtime_ns, + "bucket": bucket, + "key": self.key, + } + + +class _Download: + """one file download into `.rpdownload`, with its finished chunks + recorded in a sidecar json.""" + + def __init__(self, key: str, size: int, etag: str, dest: Path, modified: Any): + self.key = key + self.size = size + self.etag = etag + self.dest = dest + self.modified = modified + self.label = str(dest) + self.tmp = dest.with_name(dest.name + PARTIAL_SUFFIX) + self.sidecar = dest.with_name(dest.name + SIDECAR_SUFFIX) + self.count = math.ceil(size / DOWNLOAD_CHUNK_SIZE) + self.completed_chunks: set = set() + + def chunk_length(self, index: int) -> int: + return min(DOWNLOAD_CHUNK_SIZE, self.size - index * DOWNLOAD_CHUNK_SIZE) + + def load_or_init(self) -> set: + """the chunks an earlier run finished, or a fresh partial file.""" + state = _read_json(self.sidecar) + if self._resumes(state): + return {int(i) for i in state.get("done", []) if 0 <= int(i) < self.count} + with open(self.tmp, "wb") as handle: + handle.truncate(self.size) + self.save_state() + return set() + + def _resumes(self, state: Optional[Dict[str, Any]]) -> bool: + # the etag is only compared for equality. a multipart object's etag + # is "-", never the md5 of the file. + return bool( + state + and state.get("key") == self.key + and state.get("size") == self.size + and state.get("etag") == self.etag + and state.get("chunk_size") == DOWNLOAD_CHUNK_SIZE + and self.tmp.is_file() + and self.tmp.stat().st_size == self.size + ) + + def save_state(self) -> None: + _write_json_atomic( + self.sidecar, + { + "key": self.key, + "size": self.size, + "etag": self.etag, + "chunk_size": DOWNLOAD_CHUNK_SIZE, + "done": sorted(self.completed_chunks), + }, + ) + + +def _same_version(a: os.stat_result, b: os.stat_result) -> bool: + """whether two stats of a file show the same size and mtime.""" + return a.st_size == b.st_size and a.st_mtime_ns == b.st_mtime_ns + + +def _join_key(prefix: str, rel: str) -> str: + return f"{prefix}/{rel}" if prefix else rel + + +def _file_key(local: Path, remote_path: Optional[str]) -> str: + """the key for a single-file upload. it defaults to the local name, and + a remote path ending in "/" keeps the local name inside that folder.""" + if remote_path is None: + return local.name + key = remote_path.lstrip("/") + return key + local.name if remote_path.endswith("/") else key + + +def _object_dest(key: str, local_path: Any) -> Path: + """the local file for a single-object download. a trailing slash names + a folder to download into, as on upload.""" + name = PurePosixPath(key).name + dest = Path(local_path) if local_path is not None else Path(name) + if dest.is_dir() or str(local_path or "").endswith(("/", os.sep)): + return dest / name + return dest + + +def _write_at(path: Path, offset: int, data: bytes) -> None: + """write and sync data at offset, so it is on disk before the sidecar + claims it.""" + with open(path, "r+b") as handle: + handle.seek(offset) + handle.write(data) + handle.flush() + os.fsync(handle.fileno()) def _fold(name: str) -> str: @@ -871,10 +960,17 @@ def _check_local_names(targets: List[Path], root: Path) -> None: """refuse targets that would share a path on disk: two targets that are one name on a case-insensitive or normalizing filesystem ("Model.bin" and "model.bin"), a file whose name is also a folder another file lives in - ("a" beside "A/b"), or a file whose partial or sidecar name is also a target or - a directory a target lives in ("x" beside "x.rpdownload" or - "x.rpdownload/y"). that can refuse a listing a case-sensitive filesystem - would take, which is the safe direction.""" + ("a" beside "A/b"), or a file whose partial, sidecar, or sidecar temp + name is also a target or a directory a target lives in ("x" beside + "x.rpdownload" or "x.rpdownload/y"). that can refuse a listing a + case-sensitive filesystem would take, which is the safe direction.""" + files, dirs = _index_names(targets, root) + _check_file_folder_clash(files, dirs) + _check_transfer_names(targets, root, files, dirs) + + +def _index_names(targets: List[Path], root: Path) -> Tuple[Dict[str, Path], Dict[str, Path]]: + """targets and the folders they live in, keyed by folded name.""" files: Dict[str, Path] = {} dirs: Dict[str, Path] = {} for target in targets: @@ -886,18 +982,27 @@ def _check_local_names(targets: List[Path], root: Path) -> None: "file on a case-insensitive or normalizing filesystem" ) files[key] = target - for parent in rel.parents: - if str(parent) != ".": - dirs[_fold(str(parent))] = target - for key, target in files.items(): - if key in dirs: - raise NetworkVolumeTransferError( - f"cannot download both {target} and {dirs[key]}: " - "the first would have to be a file and a folder" - ) + dirs.update({_fold(str(parent)): target for parent in rel.parents if str(parent) != "."}) + return files, dirs + + +def _check_file_folder_clash(files: Dict[str, Path], dirs: Dict[str, Path]) -> None: + clash = next((key for key in files if key in dirs), None) + if clash is not None: + raise NetworkVolumeTransferError( + f"cannot download both {files[clash]} and {dirs[clash]}: " + "the first would have to be a file and a folder" + ) + + +def _check_transfer_names( + targets: List[Path], root: Path, files: Dict[str, Path], dirs: Dict[str, Path] +) -> None: + """refuse a target whose download working files would land on another + target or a folder one lives in.""" for target in targets: rel = target.relative_to(root) - for suffix in (PARTIAL_SUFFIX, SIDECAR_SUFFIX): + for suffix in TRANSFER_SUFFIXES: twin = _fold(str(rel.with_name(rel.name + suffix))) other = files.get(twin) or dirs.get(twin) if other is not None: @@ -974,18 +1079,7 @@ def _check_endpoint_resolves(endpoint: str) -> None: ) from exc -def _make_transfer( - volume: str, - access_key: Optional[str], - secret_key: Optional[str], - endpoint_url: Optional[str], - concurrency: int, - progress: Optional[Progress], - force: bool = False, -) -> _Transfer: - import boto3 # pylint: disable=import-outside-toplevel - from botocore.config import Config # pylint: disable=import-outside-toplevel - +def _credentials(access_key: Optional[str], secret_key: Optional[str]) -> Tuple[str, str]: access_key = access_key or os.environ.get(ACCESS_KEY_ENV) secret_key = secret_key or os.environ.get(SECRET_KEY_ENV) if not access_key or not secret_key: @@ -993,14 +1087,13 @@ def _make_transfer( f"s3 api keys required: set {ACCESS_KEY_ENV} and {SECRET_KEY_ENV} " f"(create keys at {S3_KEYS_URL})" ) - volume_id, datacenter = _resolve_volume(volume) - endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" - if not endpoint_url: - # an explicit endpoint (a local stand-in, a custom route) is the - # caller's; only the derived datacenter endpoint is checked - _check_endpoint_resolves(endpoint) - concurrency = max(1, int(concurrency)) - config = Config( + return access_key, secret_key + + +def _client_config(concurrency: int) -> Any: + from botocore.config import Config # pylint: disable=import-outside-toplevel + + return Config( s3={"addressing_style": "path"}, request_checksum_calculation="when_required", response_checksum_validation="when_required", @@ -1015,51 +1108,35 @@ def _make_transfer( # connections still being returned to the pool max_pool_connections=max(10, 2 * concurrency), ) + + +def _make_transfer( + volume: str, + access_key: Optional[str], + secret_key: Optional[str], + endpoint_url: Optional[str], + concurrency: int, + progress: Optional[Progress], + force: bool = False, +) -> _Transfer: + """a transfer to the network volume with this id or name. backs + NetworkVolume.upload and NetworkVolume.download.""" + import boto3 # pylint: disable=import-outside-toplevel + + access_key, secret_key = _credentials(access_key, secret_key) + volume_id, datacenter = _resolve_volume(volume) + endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" + if not endpoint_url: + # an explicit endpoint (a local stand-in, a custom route) is the + # caller's; only the derived datacenter endpoint is checked + _check_endpoint_resolves(endpoint) + concurrency = max(1, int(concurrency)) client = boto3.session.Session().client( "s3", region_name=datacenter, endpoint_url=endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, - config=config, + config=_client_config(concurrency), ) return _Transfer(client, volume_id, endpoint, concurrency, progress, force=force) - - -def upload( - volume: str, - local_path, - remote_path: Optional[str] = None, - *, - access_key: Optional[str] = None, - secret_key: Optional[str] = None, - endpoint_url: Optional[str] = None, - concurrency: int = 4, - progress: Optional[Progress] = None, - force: bool = False, -) -> Dict[str, int]: - """backs NetworkVolume.upload, which documents the behaviour.""" - transfer = _make_transfer( - volume, access_key, secret_key, endpoint_url, concurrency, progress, force - ) - return transfer.upload(Path(local_path), remote_path) - - -def download( - volume: str, - remote_path: str, - local_path=None, - *, - access_key: Optional[str] = None, - secret_key: Optional[str] = None, - endpoint_url: Optional[str] = None, - concurrency: int = 4, - progress: Optional[Progress] = None, - force: bool = False, -) -> Dict[str, int]: - """backs NetworkVolume.download, which documents the behaviour.""" - transfer = _make_transfer( - volume, access_key, secret_key, endpoint_url, concurrency, progress, force - ) - # passed through as given: Path() would drop a trailing "/" - return transfer.download(remote_path, local_path) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index e5894862d..0b1d1e20f 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -35,7 +35,29 @@ def _transfer_options(func): return func -def _run_transfer(direction, reference, source, target, opts) -> None: +class _ProgressRenderer: + """draws one bar per file in flight and prints each finished file.""" + + def __init__(self, bar, action_label): + self._bar = bar + self._action_label = action_label + self._tasks = {} + + def __call__(self, event, path, nbytes): + if event == "start": + self._tasks[path] = self._bar.add_task(Path(path).name, total=nbytes) + return + task = self._tasks.get(path) + if task is None: + return + if event in ("advance", "reuse"): + self._bar.advance(task, nbytes) + elif event == "done": + self._bar.remove_task(self._tasks.pop(path)) + self._bar.console.print(f" {self._action_label} {path}") + + +def _progress_bar(): from rich.console import Console from rich.progress import ( BarColumn, @@ -45,55 +67,58 @@ def _run_transfer(direction, reference, source, target, opts) -> None: TransferSpeedColumn, ) - from runpod import network_volume_transfer as nvt - from runpod.apps import NetworkVolume - - tasks = {} - with Progress( + return Progress( TextColumn("{task.description}"), BarColumn(), DownloadColumn(), TransferSpeedColumn(), console=Console(stderr=True), transient=True, - ) as bar: - - def progress(event, path, nbytes): - if event == "start": - tasks[path] = bar.add_task(Path(path).name, total=nbytes) - elif event in ("advance", "reuse") and path in tasks: - bar.advance(tasks[path], nbytes) - elif event == "done" and path in tasks: - bar.remove_task(tasks.pop(path)) - bar.console.print(f" {direction} {path}") - - target_volume = NetworkVolume(reference, create=False) - func = target_volume.upload if direction == "uploaded" else target_volume.download - try: - summary = func( - source, - target, - access_key=opts["s3_access_key"], - secret_key=opts["s3_secret_key"], - endpoint_url=opts["endpoint_url"], - concurrency=opts["concurrency"], - progress=progress, - force=opts["force"], - ) - except KeyboardInterrupt: - click.echo("interrupted; re-run the same command to resume", err=True) - sys.exit(130) - except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly - hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" - raise click.ClickException(f"{exc}{hint}") from exc + ) + +def _invoke_transfer(transfer_method, source, target, opts, progress): + """run the transfer, turning failures into cli errors with a resume hint + when re-running would resume.""" + from runpod import network_volume_transfer as nvt + + try: + return transfer_method( + source, + target, + access_key=opts["s3_access_key"], + secret_key=opts["s3_secret_key"], + endpoint_url=opts["endpoint_url"], + concurrency=opts["concurrency"], + progress=progress, + force=opts["force"], + ) + except KeyboardInterrupt: + click.echo("interrupted; re-run the same command to resume", err=True) + sys.exit(130) + except Exception as exc: # noqa: BLE001 - surface transfer errors cleanly + hint = "\nre-run the same command to resume" if nvt.is_resumable(exc) else "" + raise click.ClickException(f"{exc}{hint}") from exc + + +def _run_transfer(action_label, transfer_method, source, target, opts) -> None: + with _progress_bar() as bar: + summary = _invoke_transfer( + transfer_method, source, target, opts, _ProgressRenderer(bar, action_label) + ) click.echo( - f"{direction} {summary['files_transferred']} file(s), " + f"{action_label} {summary['files_transferred']} file(s), " f"{summary['bytes_transferred']} bytes; " f"skipped {summary['files_skipped']} unchanged" ) +def _volume(reference): + from runpod.apps import NetworkVolume + + return NetworkVolume(reference, create=False) + + @volume.command(name="upload") @click.argument("volume_ref", metavar="VOLUME") @click.argument("local_path", type=click.Path(exists=True)) @@ -107,7 +132,7 @@ def volume_upload(volume_ref, local_path, remote_path, **opts): has the same size and is not older are skipped; --force sends them anyway. """ - _run_transfer("uploaded", volume_ref, local_path, remote_path, opts) + _run_transfer("uploaded", _volume(volume_ref).upload, local_path, remote_path, opts) @volume.command(name="download") @@ -123,4 +148,4 @@ def volume_download(volume_ref, remote_path, local_path, **opts): size and not older than the remote copy are skipped; --force fetches them anyway. """ - _run_transfer("downloaded", volume_ref, remote_path, local_path, opts) + _run_transfer("downloaded", _volume(volume_ref).download, remote_path, local_path, opts) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index dd253a6be..8edc7b3f2 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -388,6 +388,26 @@ def test_folder_upload_default_prefix_and_failure_report(tmp_path): assert "data/ok.txt" in fake.objects +def test_folder_failure_is_resumable_only_when_every_failure_is(tmp_path): + fake = FakeS3() + root = tmp_path / "data" + root.mkdir() + (root / "cut.txt").write_bytes(b"1") + (root / "denied.txt").write_bytes(b"2") + errors = { + "data/cut.txt": _client_error("ServiceUnavailable", 503), + "data/denied.txt": _client_error("AccessDenied", 403), + } + fake.fail = lambda op, kw: errors.get(kw.get("Key")) if op == "put_object" else None + with pytest.raises(nvt.NetworkVolumeTransferError, match="2 file") as mixed: + _transfer(fake, tmp_path).upload(root, None) + assert not nvt.is_resumable(mixed.value) + del errors["data/denied.txt"] + with pytest.raises(nvt.NetworkVolumeTransferError, match="1 file") as cut: + _transfer(fake, tmp_path).upload(root, None) + assert nvt.is_resumable(cut.value) + + def test_single_file_upload_skips_unchanged(tmp_path): fake = FakeS3() src = tmp_path / "f.txt" @@ -625,7 +645,7 @@ def test_make_transfer_requires_s3_keys(monkeypatch): monkeypatch.setenv("AWS_ACCESS_KEY_ID", "aws") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "aws") with pytest.raises(nvt.NetworkVolumeTransferError, match="RUNPOD_S3_ACCESS_KEY"): - nvt.upload("vol", ".") + nvt._make_transfer("vol", None, None, None, 4, None) def test_make_transfer_builds_datacenter_client(monkeypatch): @@ -716,6 +736,7 @@ def test_cli_upload_prints_summary(tmp_path): src.write_text("x") def fake_upload(volume, source, target, **kwargs): + volume = volume.reference progress = kwargs["progress"] progress("start", source, 1) progress("advance", source, 1) @@ -723,7 +744,7 @@ def fake_upload(volume, source, target, **kwargs): assert volume == "models" and kwargs["concurrency"] == 2 return {"files_transferred": 1, "files_skipped": 0, "bytes_transferred": 1} - with patch.object(nvt, "upload", side_effect=fake_upload): + with patch.object(runpod.NetworkVolume, "upload", autospec=True, side_effect=fake_upload): result = CliRunner().invoke( cli, ["volume", "upload", "models", str(src), "--concurrency", "2"], @@ -735,7 +756,7 @@ def fake_upload(volume, source, target, **kwargs): def test_cli_download_failure_says_resume(): from runpod.rp_cli.main import cli - with patch.object(nvt, "download", side_effect=ConnectionResetError("boom")): + with patch.object(runpod.NetworkVolume, "download", side_effect=ConnectionResetError("boom")): result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code != 0 assert "boom" in result.output @@ -745,7 +766,7 @@ def test_cli_download_failure_says_resume(): def test_cli_ctrl_c_exits_130_with_resume_hint(): from runpod.rp_cli.main import cli - with patch.object(nvt, "upload", side_effect=KeyboardInterrupt): + with patch.object(runpod.NetworkVolume, "upload", side_effect=KeyboardInterrupt): result = CliRunner().invoke(cli, ["volume", "upload", "vol", "."]) assert result.exit_code == 130 assert "interrupted; re-run the same command to resume" in result.output @@ -755,7 +776,7 @@ def test_cli_unresolvable_endpoint_has_no_resume_hint(): from runpod.rp_cli.main import cli err = nvt.NetworkVolumeTransferError("the s3 endpoint s3api-x.invalid does not resolve") - with patch.object(nvt, "download", side_effect=err): + with patch.object(runpod.NetworkVolume, "download", side_effect=err): result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code == 1 assert "does not resolve" in result.output @@ -765,7 +786,7 @@ def test_cli_unresolvable_endpoint_has_no_resume_hint(): def test_cli_missing_path_has_no_resume_hint(): from runpod.rp_cli.main import cli - with patch.object(nvt, "download", side_effect=FileNotFoundError("no such file")): + with patch.object(runpod.NetworkVolume, "download", side_effect=FileNotFoundError("no such file")): result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code != 0 assert "no such file" in result.output @@ -858,7 +879,14 @@ def test_download_ignores_sidecar_of_another_key(tmp_path): @pytest.mark.parametrize( - "other", ["d/x.rpdownload", "d/x.rpdownload.json", "d/x.rpdownload/y", "d/X.RPDOWNLOAD"] + "other", + [ + "d/x.rpdownload", + "d/x.rpdownload.json", + "d/x.rpdownload.json.tmp", + "d/x.rpdownload/y", + "d/X.RPDOWNLOAD", + ], ) def test_download_refuses_partial_name_collision(tmp_path, other): fake = FakeS3() From a4bc86a78f056eeb5577d06e08f13cc1b01c31c7 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 19:09:01 -0400 Subject: [PATCH 14/23] fix: drop a stale sidecar before emptying a partial download matches runpodctl: a stop between emptying the partial and saving fresh state left the old sidecar vouching for discarded chunks. --- runpod/network_volume_transfer.py | 3 +++ tests/test_network_volume_transfer.py | 21 +++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 0b657a07f..bb09cf085 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -880,6 +880,9 @@ def load_or_init(self) -> set: state = _read_json(self.sidecar) if self._resumes(state): return {int(i) for i in state.get("done", []) if 0 <= int(i) < self.count} + # drop the old sidecar first: a stop between emptying the partial and + # saving fresh state must not leave it vouching for discarded chunks + _remove(self.sidecar) with open(self.tmp, "wb") as handle: handle.truncate(self.size) self.save_state() diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 8edc7b3f2..c1b485d7c 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -1258,3 +1258,24 @@ def interrupt_first(**kwargs): # one part per request slot that started before the interrupt landed; # nothing starts after it assert len(calls) <= 3 + + +def test_discarded_chunks_drop_the_sidecar_first(tmp_path): + dest = tmp_path / "f.bin" + job = nvt._Download("f.bin", len(DATA), "e", dest, None) + # a sidecar that vouches for every chunk beside a partial of the wrong size + job.sidecar.write_text(json.dumps( + {"key": "f.bin", "size": len(DATA), "etag": "e", "chunk_size": nvt.DOWNLOAD_CHUNK_SIZE, "done": list(range(job.count))} + )) + job.tmp.write_bytes(b"short") + sidecar_seen = [] + real_open = open + + def watch_open(path, *args, **kwargs): + if path == job.tmp: + sidecar_seen.append(job.sidecar.exists()) + return real_open(path, *args, **kwargs) + + with patch("builtins.open", watch_open): + assert job.load_or_init() == set() + assert sidecar_seen == [False] From f7299a98e3e1bcb621acde9a176094b0fe363ac4 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 19:36:12 -0400 Subject: [PATCH 15/23] fix: refuse symlinked download work files a symlinked .rpdownload, sidecar, or sidecar temp file aimed the write at another file; python has no root confinement, so it could be any file. matches runpodctl. --- runpod/network_volume_transfer.py | 11 +++++++++++ tests/test_network_volume_transfer.py | 16 ++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index bb09cf085..20e041af8 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -749,6 +749,7 @@ def download_file( return self._emit("start", job.label, size) dest.parent.mkdir(parents=True, exist_ok=True) + job.refuse_linked_work_files() job.completed_chunks = job.load_or_init() reused = sum(job.chunk_length(i) for i in job.completed_chunks) if reused: @@ -872,6 +873,16 @@ def __init__(self, key: str, size: int, etag: str, dest: Path, modified: Any): self.count = math.ceil(size / DOWNLOAD_CHUNK_SIZE) self.completed_chunks: set = set() + def refuse_linked_work_files(self) -> None: + """refuse a partial, sidecar, or sidecar temp file that exists as + anything but a regular file: a symlink would aim the write elsewhere.""" + for suffix in TRANSFER_SUFFIXES: + path = self.dest.with_name(self.dest.name + suffix) + if path.is_symlink() or (path.exists() and not path.is_file()): + raise NetworkVolumeTransferError( + f"{path} exists and is not a regular file; remove it and rerun" + ) + def chunk_length(self, index: int) -> int: return min(DOWNLOAD_CHUNK_SIZE, self.size - index * DOWNLOAD_CHUNK_SIZE) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index c1b485d7c..7e8f1f00e 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -1279,3 +1279,19 @@ def watch_open(path, *args, **kwargs): with patch("builtins.open", watch_open): assert job.load_or_init() == set() assert sidecar_seen == [False] + + +@pytest.mark.parametrize("suffix", nvt.TRANSFER_SUFFIXES) +def test_download_refuses_symlinked_work_files(tmp_path, suffix): + fake = FakeS3() + fake.put("x", DATA) # several chunks, so the sidecar path is taken + other = tmp_path / "other" + other.write_bytes(b"keep me") + try: + (tmp_path / ("x" + suffix)).symlink_to("other") + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + with pytest.raises(nvt.NetworkVolumeTransferError, match="not a regular file") as info: + _transfer(fake, tmp_path).download("x", tmp_path / "x") + assert not nvt.is_resumable(info.value) + assert other.read_bytes() == b"keep me" From a5e18c368a32a0f8f8998d8e55d02841e75c2cca Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 20:43:40 -0400 Subject: [PATCH 16/23] refactor: condense the network volume transfer module botocore's standard retry mode replaces the custom retry and backoff stack, the endpoint dns pre-check is gone (the datacenter comes from the volume record), one thread-pool helper replaces three, and one-use helpers are folded into their callers. CHANGELOG.md is back to main's: release-please owns it. --- CHANGELOG.md | 1 - runpod/network_volume_transfer.py | 1349 ++++++++++--------------- tests/test_network_volume_transfer.py | 213 +--- 3 files changed, 606 insertions(+), 957 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fb4db3893..be9484adf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,6 @@ ### Features -* `rp volume upload/download` and `NetworkVolume.upload` / `NetworkVolume.download` transfer files and folders over a network volume's S3-compatible API. Re-running an interrupted transfer resumes it: unchanged files are skipped, multipart uploads reuse MD5-verified parts, and downloads keep finished chunks. * Apps expose `NetworkVolume` and `GlobalVolume` resources with explicit path-keyed mounts, lazy create-if-missing provisioning, and execution-bound filesystem access. Global volumes use GraphQL for lookup and creation. ### Bug Fixes diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 20e041af8..2f13cc900 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -1,27 +1,25 @@ """resumable file and folder transfers to and from network volumes. -transfers go through the runpod s3-compatible api. the bucket is the -network volume id and the endpoint is per datacenter. call these -through NetworkVolume.upload and NetworkVolume.download. - -uploads above 64 MiB use multipart. the upload id is kept in a small -state file under ~/.runpod/transfers, so re-running an interrupted -upload sends only the parts the server does not already hold. downloads -fetch 64 MiB ranges in parallel into `.rpdownload`, recording -finished ranges in a sidecar json, so a re-run fetches only what is -missing. folder transfers skip files that are already present and -unchanged. +transfers go through the runpod s3-compatible api: the bucket is the +network volume id and the endpoint is per datacenter. call these through +NetworkVolume.upload and NetworkVolume.download. + +uploads above 64 MiB use multipart. the upload id is kept in a small state +file under ~/.runpod/transfers, so re-running an interrupted upload sends +only the parts the server does not already hold. downloads fetch 64 MiB +ranges into `.rpdownload`, recording finished ranges in a sidecar +json, so a re-run fetches only what is missing. folder transfers skip +files that are already present and unchanged. botocore retries failed +requests. """ import hashlib import json import math import os -import random -import socket import threading -import unicodedata import time +import unicodedata from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path, PurePosixPath from typing import Any, Callable, Dict, List, Optional, Tuple @@ -32,7 +30,6 @@ MAX_PART_SIZE = 500 * MIB # runpod's maximum multipart part size MAX_PARTS = 10000 DOWNLOAD_CHUNK_SIZE = 64 * MIB -RETRY_ATTEMPTS = 8 # runpod's s3 api keeps in-progress multipart parts here (server-created) UPLOADS_DIR = ".s3compat_uploads" PARTIAL_SUFFIX = ".rpdownload" @@ -49,8 +46,6 @@ # moved: kept parts or chunks), "skip" (nbytes = file size), "done". Progress = Callable[[str, str, int], None] -_sleep = time.sleep - class NetworkVolumeTransferError(RuntimeError): """a transfer failed. `resumable` says whether re-running can help.""" @@ -60,21 +55,6 @@ def __init__(self, message: str, resumable: bool = False): self.resumable = resumable -class _Retryable(Exception): - """a transient failure raised by this module (e.g. a short read).""" - - -class _Stopped(Exception): - """the transfer was interrupted; raised by workers so they wind down.""" - - -def _state_dir() -> Path: - return Path(os.path.expanduser("~/.runpod/transfers")) - - -# --------------------------------------------------------------- helpers - - def part_size_for(size: int) -> int: """part size for a multipart upload of `size` bytes.""" part_size = DEFAULT_PART_SIZE @@ -87,96 +67,63 @@ def part_size_for(size: int) -> int: return part_size -def _error_code(exc: Exception) -> Tuple[str, int]: +def _error(exc: BaseException) -> Tuple[str, int]: + """the s3 error code and http status of a botocore ClientError.""" response = getattr(exc, "response", None) or {} code = str(response.get("Error", {}).get("Code", "")) - status = response.get("ResponseMetadata", {}).get("HTTPStatusCode") or 0 - return code, int(status) + return code, int(response.get("ResponseMetadata", {}).get("HTTPStatusCode") or 0) -def _is_not_found(exc: Exception) -> bool: - code, status = _error_code(exc) +def _is_not_found(exc: BaseException) -> bool: + code, status = _error(exc) return status == 404 or code in ("404", "NoSuchKey", "NotFound") -def _is_retryable(exc: Exception) -> bool: +def is_resumable(exc: BaseException) -> bool: + """whether re-running the same transfer can get past `exc`: an + interruption, or a network or server failure that outlasted botocore's + retries. bad credentials, missing paths and local disk errors fail the + same way on every run.""" from botocore.exceptions import ( # pylint: disable=import-outside-toplevel - ClientError, ConnectionError as BotoConnectionError, HTTPClientError, IncompleteReadError, ResponseStreamingError, ) - if isinstance(exc, _Retryable): - return True - if isinstance( - exc, - ( - BotoConnectionError, - HTTPClientError, - IncompleteReadError, - ResponseStreamingError, - ConnectionError, - TimeoutError, - ), - ): - return True - if isinstance(exc, ClientError): - code, status = _error_code(exc) - if status == 429 or status >= 500: - return True - return code in ("SlowDown", "RequestTimeout", "InternalError", "524") - return False - - -def is_resumable(exc: BaseException) -> bool: - """whether re-running the same transfer can get past `exc`. - - an interruption, or a transient failure (network, timeout, 429, 5xx) - that outlasted the retries, resumes. anything else (missing paths, bad - credentials, a wrong datacenter, local disk errors, invalid input) fails - the same way on every run. - """ if isinstance(exc, KeyboardInterrupt): return True if isinstance(exc, NetworkVolumeTransferError): return exc.resumable - # final-looking errors a re-run gets past: the multipart upload vanished - # (e.g. after a complete whose 524 hid that it finished), or a range that - # no longer fits an object that changed mid-download - code, status = _error_code(exc) - if code == "NoSuchUpload" or status == 416: + transient = ( + BotoConnectionError, + HTTPClientError, + IncompleteReadError, + ResponseStreamingError, + ) + if isinstance(exc, transient + (ConnectionError, TimeoutError)): return True - return _is_retryable(exc) - - -def _wait(delay: float, stop: Optional[threading.Event] = None) -> None: - """sleep for a retry backoff, cut short when `stop` is set.""" - if stop is None: - _sleep(delay) - else: - stop.wait(delay) + # NoSuchUpload: the upload vanished (a complete whose 524 hid that it + # finished); 416: a range that no longer fits an object that changed + code, status = _error(exc) + return ( + code in ("NoSuchUpload", "SlowDown", "RequestTimeout") + or status in (416, 429) + or status >= 500 + ) -def _with_retries( - func: Callable[[], Any], - attempts: int = RETRY_ATTEMPTS, - stop: Optional[threading.Event] = None, -) -> Any: - """call func, retrying transient failures with exponential backoff. a set - `stop` (ctrl-c) ends a backoff early; the next attempt then raises.""" - for attempt in range(attempts): - try: - return func() - except Exception as exc: # noqa: BLE001 - filtered by _is_retryable - if attempt == attempts - 1 or not _is_retryable(exc): - raise - _wait(min(2**attempt, 30) + random.uniform(0, 1), stop) - raise AssertionError("unreachable") +def _read_json(path: Path) -> Optional[Dict[str, Any]]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None -def _write_json_atomic(path: Path, data: Dict[str, Any]) -> None: +def _write_json(path: Path, data: Dict[str, Any]) -> None: + """replace path atomically, synced, so a crash leaves the old state or + the new one.""" path.parent.mkdir(parents=True, exist_ok=True) tmp = path.with_name(path.name + ".tmp") with open(tmp, "w", encoding="utf-8") as handle: @@ -184,55 +131,23 @@ def _write_json_atomic(path: Path, data: Dict[str, Any]) -> None: handle.flush() os.fsync(handle.fileno()) os.replace(tmp, path) - # sync the directory too, so the rename survives a power loss (best - # effort: not every platform can open a directory) - try: - fd = os.open(path.parent, os.O_RDONLY) - except OSError: - return - try: - os.fsync(fd) - except OSError: - # some platforms refuse to fsync a directory; the file itself is synced - return - finally: - os.close(fd) - - -def _read_json(path: Path) -> Optional[Dict[str, Any]]: - try: - with open(path, encoding="utf-8") as handle: - data = json.load(handle) - except (OSError, ValueError): - return None - return data if isinstance(data, dict) else None - - -def _remove(path: Path) -> None: - path.unlink(missing_ok=True) - - -def _read_range(path: Path, offset: int, length: int) -> bytes: - with open(path, "rb") as handle: - handle.seek(offset) - return handle.read(length) def _md5_hex(path: Path, offset: int, length: int) -> str: digest = hashlib.md5() # noqa: S324 - etag comparison, not security with open(path, "rb") as handle: handle.seek(offset) - remaining = length - while remaining > 0: - block = handle.read(min(8 * MIB, remaining)) - if not block: - break + while length > 0 and (block := handle.read(min(8 * MIB, length))): digest.update(block) - remaining -= len(block) + length -= len(block) return digest.hexdigest() -def _is_md5_etag(etag: str) -> bool: +def _etag(value: Any) -> str: + return str(value or "").strip('"') + + +def _is_md5(etag: str) -> bool: return len(etag) == 32 and all(c in "0123456789abcdef" for c in etag.lower()) @@ -240,64 +155,27 @@ def _multipart_etag(parts: List[Dict[str, Any]]) -> str: """the etag s3 (and runpod, verified) gives an object assembled from parts: md5 of the parts' binary md5s, then "-". "" when a part etag is not an md5.""" - digest = hashlib.md5() - for part in parts: - etag = _strip_etag(part["ETag"]) - if not _is_md5_etag(etag): - return "" - digest.update(bytes.fromhex(etag)) + etags = [_etag(part["ETag"]) for part in parts] + if not all(_is_md5(etag) for etag in etags): + return "" + digest = hashlib.md5(b"".join(bytes.fromhex(etag) for etag in etags)) # noqa: S324 return f"{digest.hexdigest()}-{len(parts)}" -def _run_pool( - tasks: List[Callable[[], Any]], - concurrency: int, - stop: Optional[threading.Event] = None, -) -> List[Any]: - """run tasks on a thread pool. the first failure cancels the rest. - - an interrupt (ctrl-c) also sets `stop`, before the pool waits for its - running tasks, so workers checking it wind down instead of finishing - every file in flight. - """ - results = [] - with ThreadPoolExecutor(max_workers=max(1, concurrency)) as pool: - futures = [pool.submit(task) for task in tasks] - try: - for future in as_completed(futures): - results.append(future.result()) - except BaseException as exc: - if stop is not None and not isinstance(exc, Exception): - stop.set() - for future in futures: - future.cancel() - raise - return results - - -def _is_download_leftover(path: Path) -> bool: - """a partial download, its sidecar, or the sidecar's temp file.""" - return path.name.endswith(TRANSFER_SUFFIXES) +def _same_version(a: os.stat_result, b: os.stat_result) -> bool: + return a.st_size == b.st_size and a.st_mtime_ns == b.st_mtime_ns -def _same_second_or_later(a: float, b: float) -> bool: +def _not_older(a: float, b: float) -> bool: """a >= b in whole seconds; s3 dates carry no finer precision.""" return int(a) >= int(b) def _check_remote_path(remote_path: Optional[str]) -> None: - """refuse remote paths with '..' segments before anything is sent.""" if remote_path and ".." in remote_path.replace("\\", "/").split("/"): raise ValueError(f"remote path {remote_path!r} must not contain '..'") -def _strip_etag(etag: Any) -> str: - return str(etag or "").strip('"') - - -# ------------------------------------------------------------- transfers - - # the summary counter each progress event advances _SUMMARY_FIELDS = { "advance": "bytes_transferred", @@ -323,34 +201,42 @@ def __init__( self.bucket = bucket self.endpoint = endpoint self.concurrency = max(1, int(concurrency)) + self.state_dir = state_dir or Path.home() / ".runpod" / "transfers" + # transfer every file, even ones the skip rule calls unchanged + self.force = force + self.summary = { + "files_transferred": 0, + "files_skipped": 0, + "bytes_transferred": 0, + } self._progress = progress - self.state_dir = state_dir or _state_dir() self._lock = threading.Lock() - # bounds requests in flight across all files: folder transfers run - # files in parallel and each large file runs its parts in parallel + # folder transfers run files in parallel and each large file runs its + # parts in parallel; this bounds the requests in flight across both self._slots = threading.BoundedSemaphore(self.concurrency) - # transfer every file, even ones the skip rule calls unchanged - self.force = force - # set on ctrl-c so every worker winds down + # set on ctrl-c so queued work stops instead of finishing files in flight self._stop = threading.Event() - self.summary = {"files_transferred": 0, "files_skipped": 0, "bytes_transferred": 0} - # -- bookkeeping - - def _request(self, func: Callable[[], Any]) -> Any: - if self._stop.is_set(): - raise _Stopped("interrupted") + def _call(self, method: str, **kwargs: Any) -> Any: with self._slots: - # checked again: an interrupt may land while waiting for the slot if self._stop.is_set(): - raise _Stopped("interrupted") - return func() - - def _retry(self, func: Callable[[], Any]) -> Any: - return _with_retries(func, stop=self._stop) + raise NetworkVolumeTransferError("interrupted", resumable=True) + return getattr(self.client, method)(Bucket=self.bucket, **kwargs) - def _pool(self, tasks: List[Callable[[], Any]]) -> List[Any]: - return _run_pool(tasks, self.concurrency, self._stop) + def _run(self, tasks: List[Callable[[], Any]]) -> List[Any]: + """run tasks in parallel; the first failure cancels the ones not started.""" + pool = ThreadPoolExecutor(max_workers=self.concurrency) + try: + return [ + future.result() + for future in as_completed([pool.submit(t) for t in tasks]) + ] + except BaseException as exc: + if not isinstance(exc, Exception): + self._stop.set() + raise + finally: + pool.shutdown(cancel_futures=True) def _emit(self, event: str, path: str, nbytes: int = 0) -> None: field = _SUMMARY_FIELDS.get(event) @@ -360,74 +246,83 @@ def _emit(self, event: str, path: str, nbytes: int = 0) -> None: if self._progress is not None: self._progress(event, path, nbytes) - def _state_path(self, key: str, local: Path) -> Path: - ident = f"{self.endpoint}|{self.bucket}|{key}|{local.resolve()}" - return self.state_dir / (hashlib.sha256(ident.encode()).hexdigest() + ".json") + def _run_files(self, jobs: List[Tuple[Any, Callable[[], Any]]]) -> None: + """run every file; failures are reported together, resumable only + when every one of them is.""" + failures: List[Tuple[str, bool]] = [] + + def guarded(label: Any, task: Callable[[], Any]) -> None: + try: + task() + except Exception as exc: # noqa: BLE001 - collected and re-raised + with self._lock: + failures.append((f"{label}: {exc}", is_resumable(exc))) + + self._run([lambda l=label, t=task: guarded(l, t) for label, task in jobs]) + if failures: + raise NetworkVolumeTransferError( + f"{len(failures)} file(s) failed:\n " + + "\n ".join(sorted(m for m, _ in failures)), + resumable=all(ok for _, ok in failures), + ) # -- remote metadata def head(self, key: str, directory_ok: bool = False) -> Optional[Dict[str, Any]]: - """object metadata, or None when there is no file at `key`. - - with `directory_ok`, a 400 also means "not a file": runpod answers - HeadObject on a directory with 400 rather than 404. - """ + """object metadata, or None when there is no file at `key`. with + `directory_ok`, a 400 also means "not a file": runpod answers + HeadObject on a directory with 400 rather than 404.""" try: - return self._retry( - lambda: self._request( - lambda: self.client.head_object(Bucket=self.bucket, Key=key) - ) - ) + return self._call("head_object", Key=key) except Exception as exc: # noqa: BLE001 - if _is_not_found(exc) or (directory_ok and _error_code(exc)[1] == 400): + if _is_not_found(exc) or (directory_ok and _error(exc)[1] == 400): return None raise - def list_prefix(self, prefix: str) -> List[Dict[str, Any]]: - objects: List[Dict[str, Any]] = [] - kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Prefix": prefix} + def _pages( + self, method: str, items: str, marker: str, arg: str, **kwargs: Any + ) -> List[Any]: + found: List[Any] = [] while True: - try: - page = self._retry( - lambda: self._request(lambda: self.client.list_objects_v2(**kwargs)) - ) - except Exception as exc: # noqa: BLE001 - if _is_not_found(exc): - return objects - raise - objects.extend(page.get("Contents") or []) - token = page.get("NextContinuationToken") - if not page.get("IsTruncated") or not token: - return objects - kwargs["ContinuationToken"] = token + page = self._call(method, **kwargs) + found.extend(page.get(items) or []) + if not page.get("IsTruncated") or not page.get(marker): + return found + kwargs[arg] = page[marker] + + def list_files(self, prefix: str) -> List[Dict[str, Any]]: + """the files under prefix, without folder markers and in-progress + multipart parts.""" + try: + listed = self._pages( + "list_objects_v2", + "Contents", + "NextContinuationToken", + "ContinuationToken", + Prefix=prefix, + ) + except Exception as exc: # noqa: BLE001 + if _is_not_found(exc): + return [] + raise + return [ + obj + for obj in listed + if not obj["Key"].endswith("/") + and UPLOADS_DIR not in PurePosixPath(obj["Key"]).parts + ] def _is_folder(self, key: str) -> bool: """whether something is listed under `key/` (runpod lists even an empty folder's own marker).""" try: - page = self._retry( - lambda: self._request( - lambda: self.client.list_objects_v2( - Bucket=self.bucket, Prefix=f"{key}/", MaxKeys=1 - ) + return bool( + self._call("list_objects_v2", Prefix=f"{key}/", MaxKeys=1).get( + "Contents" ) ) except Exception: # noqa: BLE001 - unknown: report the original error return False - return bool(page.get("Contents")) - - def _list_parts(self, key: str, upload_id: str) -> List[Dict[str, Any]]: - parts: List[Dict[str, Any]] = [] - kwargs: Dict[str, Any] = {"Bucket": self.bucket, "Key": key, "UploadId": upload_id} - while True: - page = self._retry( - lambda: self._request(lambda: self.client.list_parts(**kwargs)) - ) - parts.extend(page.get("Parts") or []) - marker = page.get("NextPartNumberMarker") - if not page.get("IsTruncated") or not marker: - return parts - kwargs["PartNumberMarker"] = marker # -- upload @@ -438,254 +333,107 @@ def upload(self, local: Path, remote_path: Optional[str]) -> Dict[str, int]: if local.is_dir(): self._upload_folder(local, remote_path) else: - self._upload_single_file(local, remote_path) + self._upload_one(local, remote_path) return self.summary def _upload_folder(self, local: Path, remote_path: Optional[str]) -> None: - prefix = (local.resolve().name if remote_path is None else remote_path).strip("/") - remote = {obj["Key"]: obj for obj in self.list_prefix(f"{prefix}/" if prefix else "")} + prefix = (local.resolve().name if remote_path is None else remote_path).strip( + "/" + ) + remote = { + obj["Key"]: obj for obj in self.list_files(f"{prefix}/" if prefix else "") + } jobs = [] for path in sorted(p for p in local.rglob("*") if p.is_file()): - key = _join_key(prefix, path.relative_to(local).as_posix()) - if not self._skip_upload(path, remote.get(key)): - jobs.append((path, key)) - self._run_files([lambda p=p, k=k: self.upload_file(p, k) for p, k in jobs], jobs) - - def _skip_upload(self, path: Path, remote: Optional[Dict[str, Any]]) -> bool: - """report a folder file that needs no upload: an interrupted - download's partial file (not data), or an unchanged file.""" - leftover = _is_download_leftover(path) - if not leftover and (self.force or not self._unchanged_remote(remote, path)): - return False - self._emit("skip", str(path), path.stat().st_size) - return True - - def _upload_single_file(self, local: Path, remote_path: Optional[str]) -> None: - key = _file_key(local, remote_path) - remote = self._head_upload_target(key) - if not self.force and self._unchanged_remote(remote, local): - self._emit("skip", str(local), local.stat().st_size) - return - self.upload_file(local, key) - - def _head_upload_target(self, key: str) -> Optional[Dict[str, Any]]: + key = path.relative_to(local).as_posix() + key = f"{prefix}/{key}" if prefix else key + # an interrupted download's working files are not data + if path.name.endswith(TRANSFER_SUFFIXES) or self._unchanged_remote( + remote.get(key), path + ): + self._emit("skip", str(path), path.stat().st_size) + else: + jobs.append((path, lambda p=path, k=key: self.upload_file(p, k))) + self._run_files(jobs) + + def _upload_one(self, local: Path, remote_path: Optional[str]) -> None: + # a remote path ending in "/" keeps the local name inside that folder + key = local.name if remote_path is None else remote_path.lstrip("/") + if remote_path is not None and remote_path.endswith("/"): + key += local.name try: - return self.head(key) + remote = self.head(key) except Exception as exc: # noqa: BLE001 # runpod answers a head on a directory with 400, but a head has no # body, so a 400 for any other reason looks the same: confirm - if _error_code(exc)[1] == 400 and self._is_folder(key): + if _error(exc)[1] == 400 and self._is_folder(key): raise NetworkVolumeTransferError( - f"{key!r} is a folder on the volume; use {key + '/'!r} to upload into it" + f"{key!r} is a folder on the volume; " + f"use {key + '/'!r} to upload into it" ) from exc raise + if self._unchanged_remote(remote, local): + self._emit("skip", str(local), local.stat().st_size) + return + self.upload_file(local, key) def _unchanged_remote(self, remote: Optional[Dict[str, Any]], local: Path) -> bool: - if not remote: + if self.force or not remote or remote.get("LastModified") is None: return False stat = local.stat() size = remote.get("Size", remote.get("ContentLength")) - modified = remote.get("LastModified") - if size is None or int(size) != stat.st_size or modified is None: + if size is None or int(size) != stat.st_size: return False - return _same_second_or_later(modified.timestamp(), stat.st_mtime) + return _not_older(remote["LastModified"].timestamp(), stat.st_mtime) def upload_file(self, local: Path, key: str) -> None: stat = local.stat() - label = str(local) - self._emit("start", label, stat.st_size) + self._emit("start", str(local), stat.st_size) if stat.st_size <= SINGLE_PUT_MAX: self._single_put(local, key, stat) else: - self._multipart_upload(_Upload(local, key, stat, part_size_for(stat.st_size))) - self._emit("done", label) + _Upload(self, local, key, stat).run() + self._emit("done", str(local)) def _single_put(self, local: Path, key: str, stat: os.stat_result) -> None: - self._retry( - lambda: self._request( - lambda: self.client.put_object(Bucket=self.bucket, Key=key, Body=local.read_bytes()) - ) - ) - self._delete_if_changed(local, key, stat) - self._emit("advance", str(local), stat.st_size) - - def _delete_if_changed(self, local: Path, key: str, stat: os.stat_result) -> None: - """a retry rereads the file; if it changed meanwhile, the object may - hold either version and is newer than the local edit, so a re-run - would skip it. remove it so the re-run sends the file.""" + self._call("put_object", Key=key, Body=local.read_bytes()) try: now = local.stat() except OSError as exc: # gone or unreadable: the object is what was read, keep it - raise NetworkVolumeTransferError(f"{local} is gone after its upload: {exc}") from exc - if _same_version(stat, now): - return + raise NetworkVolumeTransferError( + f"{local} is gone after its upload: {exc}" + ) from exc + if not _same_version(stat, now): + self._delete_changed(local, key) + self._emit("advance", str(local), stat.st_size) + + def _delete_changed(self, local: Path, key: str) -> None: + """the file changed while it was read, so the object may hold either + version and is newer than the local edit; a re-run would skip it. + remove it so the re-run sends the file. not via _call: this cleanup + runs even after ctrl-c.""" try: - # not via self._retry: this cleanup runs even after ctrl-c - _with_retries(lambda: self.client.delete_object(Bucket=self.bucket, Key=key)) + self.client.delete_object(Bucket=self.bucket, Key=key) except Exception as exc: # noqa: BLE001 raise NetworkVolumeTransferError( f"{local} changed during upload and the remote copy could not be " f"removed ({exc}); re-run with --force" ) from exc - raise NetworkVolumeTransferError(f"{local} changed during upload", resumable=True) - - def _multipart_upload(self, up: "_Upload") -> None: - up.state_path = self._state_path(up.key, up.local) - if not self._resume_upload(up): - self._create_upload(up) - reused = sum(up.part_length(n) for n in up.parts) - if reused: - self._emit("reuse", up.label, reused) - self._send_missing_parts(up) - # the parts were read over the whole upload; if the file moved under - # them, they may mix two versions of it - if not _same_version(up.stat, up.local.stat()): - self._abort(up.key, up.upload_id) - _remove(up.state_path) - raise NetworkVolumeTransferError(f"{up.local} changed during upload", resumable=True) - self._complete_upload(up) - _remove(up.state_path) - - def _resume_upload(self, up: "_Upload") -> bool: - """pick up the saved upload for this file with its verified parts. - false when there is none to resume.""" - state = _read_json(up.state_path) - if not state or not state.get("upload_id"): - return False - if not up.matches(state): - # the local file changed since the upload began: its parts are stale - self._abort(up.key, state["upload_id"]) - _remove(up.state_path) - return False - try: - up.parts = self._verified_parts(up, state["upload_id"]) - except LookupError: - # the server no longer has this upload: start a new one - _remove(up.state_path) - return False - up.upload_id = state["upload_id"] - return True - - def _verified_parts(self, up: "_Upload", upload_id: str) -> Dict[int, str]: - """parts of the saved upload that match the local file.""" - try: - listed = self._list_parts(up.key, upload_id) - except Exception as exc: # noqa: BLE001 - code, _ = _error_code(exc) - if code == "NoSuchUpload" or _is_not_found(exc): - raise LookupError(upload_id) from exc - raise - return { - int(part["PartNumber"]): part.get("ETag") - for part in listed - if up.part_matches(part) - } - - def _create_upload(self, up: "_Upload") -> None: - created = self._retry( - lambda: self._request( - lambda: self.client.create_multipart_upload(Bucket=self.bucket, Key=up.key) - ) - ) - up.upload_id = created["UploadId"] - up.parts = {} - try: - _write_json_atomic(up.state_path, up.state(self.bucket)) - except OSError: - # without saved state no re-run can find this upload again - self._abort(up.key, up.upload_id) - raise - - def _send_missing_parts(self, up: "_Upload") -> None: - todo = [n for n in range(1, up.count + 1) if n not in up.parts] - for number, etag in self._pool([lambda n=n: self._send_part(up, n) for n in todo]): - up.parts[number] = etag - - def _send_part(self, up: "_Upload", number: int) -> Tuple[int, str]: - offset = (number - 1) * up.part_size - length = up.part_length(number) - response = self._retry( - lambda: self._request( - lambda: self.client.upload_part( - Bucket=self.bucket, - Key=up.key, - UploadId=up.upload_id, - PartNumber=number, - Body=_read_range(up.local, offset, length), - ) - ) + raise NetworkVolumeTransferError( + f"{local} changed during upload", resumable=True ) - self._emit("advance", up.label, length) - return number, response["ETag"] - - def _complete_upload(self, up: "_Upload") -> None: - parts = [{"PartNumber": n, "ETag": up.parts[n]} for n in sorted(up.parts)] - try: - self._retry( - lambda: self._request( - lambda: self.client.complete_multipart_upload( - Bucket=self.bucket, - Key=up.key, - UploadId=up.upload_id, - MultipartUpload={"Parts": parts}, - ) - ) - ) - except Exception as exc: # noqa: BLE001 - re-raised unless it completed - if not self._completed_anyway(exc, up.key, up.stat.st_size, _multipart_etag(parts)): - raise - - def _completed_anyway(self, exc: Exception, key: str, size: int, want_etag: str) -> bool: - """whether a failed complete took effect anyway. - - assembling a large file can outlive the proxy (524) while it still - finishes server-side, and a retry then finds the upload gone. only - those ambiguous failures qualify, and only when the object at the key - is this upload's: its size matches and its etag equals the multipart - etag of our parts. - """ - if not (_is_retryable(exc) or _is_not_found(exc)): - return False - head = self.head(key) - if head is None or int(head.get("ContentLength", -1)) != size: - return False - # without our multipart etag the object's identity cannot be confirmed - return bool(want_etag) and _strip_etag(head.get("ETag")) == want_etag def _abort(self, key: str, upload_id: str) -> None: - """discard an upload that can no longer be resumed (best effort).""" + """discard an upload that can no longer be resumed (best effort: a + failed abort only leaves orphaned parts behind).""" try: - self._request( - lambda: self.client.abort_multipart_upload( - Bucket=self.bucket, Key=key, UploadId=upload_id - ) + self.client.abort_multipart_upload( + Bucket=self.bucket, Key=key, UploadId=upload_id ) except Exception: # noqa: BLE001 - # a failed abort only leaves orphaned parts behind return - def _run_files(self, tasks: List[Callable[[], Any]], jobs: List[Tuple[Any, Any]]) -> None: - """run per-file tasks. every file is attempted; failures are reported - together, resumable only when every one of them is.""" - failures = [] - resumable = [] - - def guarded(index: int) -> None: - try: - tasks[index]() - except Exception as exc: # noqa: BLE001 - collected and re-raised - with self._lock: - failures.append(f"{jobs[index][0]}: {exc}") - resumable.append(is_resumable(exc)) - - self._pool([lambda i=i: guarded(i) for i in range(len(tasks))]) - if failures: - raise NetworkVolumeTransferError( - f"{len(failures)} file(s) failed:\n " + "\n ".join(sorted(failures)), - resumable=all(resumable), - ) - # -- download def download(self, remote_path: str, local_path: Any = None) -> Dict[str, int]: @@ -693,189 +441,276 @@ def download(self, remote_path: str, local_path: Any = None) -> Dict[str, int]: download into (a Path cannot carry the slash).""" _check_remote_path(remote_path) key = remote_path.lstrip("/") - head = None if key == "" or key.endswith("/") else self.head(key, directory_ok=True) - if head is not None: - self._download_object(key, head, local_path) - else: - self._download_prefix(key.rstrip("/"), remote_path, local_path) - return self.summary - - def _download_object(self, key: str, head: Dict[str, Any], local_path: Any) -> None: - self.download_file( - key, - int(head["ContentLength"]), - _strip_etag(head.get("ETag")), - _object_dest(key, local_path), - head.get("LastModified"), + head = ( + None + if key == "" or key.endswith("/") + else self.head(key, directory_ok=True) ) + if head is None: + self._download_folder(key.rstrip("/"), remote_path, local_path) + return self.summary + name = PurePosixPath(key).name + dest = Path(local_path) if local_path is not None else Path(name) + # a trailing slash names a folder to download into, as on upload + if dest.is_dir() or str(local_path or "").endswith(("/", os.sep)): + dest = dest / name + _Download(self, key, head["ContentLength"], head, dest).run() + return self.summary - def _download_prefix(self, prefix: str, remote_path: str, local_path: Any) -> None: - list_prefix = f"{prefix}/" if prefix else "" - objects = self._listed_files(list_prefix) + def _download_folder(self, prefix: str, remote_path: str, local_path: Any) -> None: + listed = f"{prefix}/" if prefix else "" + objects = self.list_files(listed) # a typo must not look like an empty folder; the whole volume may be empty if not objects and prefix: - raise FileNotFoundError(f"no such file or directory on the volume: {remote_path}") - root = Path(local_path) if local_path is not None else Path( - PurePosixPath(prefix).name or self.bucket + raise FileNotFoundError( + f"no such file or directory on the volume: {remote_path}" + ) + root = ( + Path(local_path) + if local_path is not None + else Path(PurePosixPath(prefix).name or self.bucket) ) - jobs = [(_safe_target(root, obj["Key"][len(list_prefix):]), obj) for obj in objects] + jobs = [(_safe_target(root, obj["Key"][len(listed) :]), obj) for obj in objects] _check_local_names([dest for dest, _ in jobs], root) self._run_files( [ - lambda d=d, o=o: self.download_file( - o["Key"], int(o["Size"]), _strip_etag(o.get("ETag")), d, o.get("LastModified") + ( + dest, + lambda d=dest, o=obj: _Download( + self, o["Key"], o["Size"], o, d + ).run(), ) - for d, o in jobs - ], - jobs, - ) - - def _listed_files(self, list_prefix: str) -> List[Dict[str, Any]]: - """the files under a prefix, without folder markers and in-progress - multipart parts.""" - return [ - obj - for obj in self.list_prefix(list_prefix) - if not obj["Key"].endswith("/") - and UPLOADS_DIR not in PurePosixPath(obj["Key"]).parts - ] - - def download_file( - self, key: str, size: int, etag: str, dest: Path, modified: Any = None - ) -> None: - job = _Download(key, size, etag, dest, modified) - if self._unchanged_local(job): - self._emit("skip", job.label, size) - return - self._emit("start", job.label, size) - dest.parent.mkdir(parents=True, exist_ok=True) - job.refuse_linked_work_files() - job.completed_chunks = job.load_or_init() - reused = sum(job.chunk_length(i) for i in job.completed_chunks) - if reused: - self._emit("reuse", job.label, reused) - todo = [i for i in range(job.count) if i not in job.completed_chunks] - self._pool([lambda i=i: self._fetch_chunk(job, i) for i in todo]) - self._finish_download(job) - - def _unchanged_local(self, job: "_Download") -> bool: - """downloads stamp the remote time, so an unchanged file is the same - size and not older than the remote copy.""" - if self.force or not job.dest.is_file(): - return False - stat = job.dest.stat() - if stat.st_size != job.size: - return False - return job.modified is None or _same_second_or_later( - stat.st_mtime, job.modified.timestamp() + for dest, obj in jobs + ] ) - def _fetch_chunk(self, job: "_Download", index: int) -> None: - start = index * DOWNLOAD_CHUNK_SIZE - length = job.chunk_length(index) - data = self._retry(lambda: self._request(lambda: self._get_range(job, start, length))) - _write_at(job.tmp, start, data) - with self._lock: - job.completed_chunks.add(index) - job.save_state() - self._emit("advance", job.label, length) - - def _get_range(self, job: "_Download", start: int, length: int) -> bytes: - end = start + length - 1 - response = self.client.get_object( - Bucket=self.bucket, Key=job.key, Range=f"bytes={start}-{end}" - ) - got = _strip_etag(response.get("ETag")) - if job.etag and got and got != job.etag: - raise NetworkVolumeTransferError( - f"{job.key} changed on the volume during download; " - "run the command again to restart it", - resumable=True, - ) - data = response["Body"].read() - if len(data) != length: - raise _Retryable(f"short read for {job.key} bytes {start}-{end}") - return data - - def _finish_download(self, job: "_Download") -> None: - if job.tmp.stat().st_size != job.size: - raise NetworkVolumeTransferError( - f"{job.dest}: size mismatch after download", resumable=True - ) - os.replace(job.tmp, job.dest) - if job.modified is not None: - # stamp the remote time so uploading the file back counts as unchanged - os.utime(job.dest, (time.time(), job.modified.timestamp())) - _remove(job.sidecar) - self._emit("done", job.label) - class _Upload: """one multipart upload of a local file, pinned to its size and mtime.""" - def __init__(self, local: Path, key: str, stat: os.stat_result, part_size: int): + def __init__( + self, transfer: _Transfer, local: Path, key: str, stat: os.stat_result + ): + self.t = transfer self.local = local self.key = key self.stat = stat - self.part_size = part_size - self.count = math.ceil(stat.st_size / part_size) - self.label = str(local) - self.state_path = Path() + self.part_size = part_size_for(stat.st_size) + self.count = math.ceil(stat.st_size / self.part_size) + ident = f"{transfer.endpoint}|{transfer.bucket}|{key}|{local.resolve()}" + self.state_path = transfer.state_dir / ( + hashlib.sha256(ident.encode()).hexdigest() + ".json" + ) self.upload_id = "" self.parts: Dict[int, str] = {} # part number -> etag - def part_length(self, number: int) -> int: - return min(self.part_size, self.stat.st_size - (number - 1) * self.part_size) - - def matches(self, state: Dict[str, Any]) -> bool: - """whether a saved state belongs to this version of the file.""" - return ( - state.get("file_size") == self.stat.st_size - and state.get("mtime_ns") == self.stat.st_mtime_ns - and state.get("part_size") == self.part_size - ) + def run(self) -> None: + if not self._resume(): + self._create() + reused = sum(self._length(n) for n in self.parts) + if reused: + self.t._emit("reuse", str(self.local), reused) + todo = [n for n in range(1, self.count + 1) if n not in self.parts] + self.parts.update(self.t._run([lambda n=n: self._send(n) for n in todo])) + # the parts were read over the whole upload; if the file moved under + # them, they may mix two versions of it + if not _same_version(self.stat, self.local.stat()): + self.t._abort(self.key, self.upload_id) + self.state_path.unlink(missing_ok=True) + raise NetworkVolumeTransferError( + f"{self.local} changed during upload", resumable=True + ) + self._complete() + self.state_path.unlink(missing_ok=True) - def part_matches(self, part: Dict[str, Any]) -> bool: - """whether a listed part holds this file's bytes. a part whose - content cannot be verified is uploaded again.""" - number = int(part["PartNumber"]) - expected = self.part_length(number) - if expected <= 0 or int(part.get("Size", -1)) != expected: - return False - etag = _strip_etag(part.get("ETag")) - offset = (number - 1) * self.part_size - return _is_md5_etag(etag) and _md5_hex(self.local, offset, expected) == etag.lower() + def _length(self, number: int) -> int: + return min(self.part_size, self.stat.st_size - (number - 1) * self.part_size) - def state(self, bucket: str) -> Dict[str, Any]: + def _state(self) -> Dict[str, Any]: return { "upload_id": self.upload_id, "part_size": self.part_size, "file_size": self.stat.st_size, "mtime_ns": self.stat.st_mtime_ns, - "bucket": bucket, + "bucket": self.t.bucket, "key": self.key, } + def _resume(self) -> bool: + """pick up the saved upload for this file with its verified parts; + false when there is none to resume.""" + state = _read_json(self.state_path) or {} + upload_id = state.get("upload_id") + if not upload_id: + return False + pinned = (self.part_size, self.stat.st_size, self.stat.st_mtime_ns) + if ( + state.get("part_size"), + state.get("file_size"), + state.get("mtime_ns"), + ) != pinned: + # the local file changed since the upload began: its parts are stale + self.t._abort(self.key, upload_id) + self.state_path.unlink(missing_ok=True) + return False + try: + listed = self.t._pages( + "list_parts", + "Parts", + "NextPartNumberMarker", + "PartNumberMarker", + Key=self.key, + UploadId=upload_id, + ) + except Exception as exc: # noqa: BLE001 + if _error(exc)[0] != "NoSuchUpload" and not _is_not_found(exc): + raise + # the server no longer has this upload: start a new one + self.state_path.unlink(missing_ok=True) + return False + self.upload_id = upload_id + self.parts = { + int(p["PartNumber"]): p["ETag"] for p in listed if self._part_matches(p) + } + return True + + def _part_matches(self, part: Dict[str, Any]) -> bool: + """whether a listed part holds this file's bytes. a part whose + content cannot be verified is uploaded again.""" + number = int(part["PartNumber"]) + length = self._length(number) + etag = _etag(part.get("ETag")) + if length <= 0 or int(part.get("Size", -1)) != length or not _is_md5(etag): + return False + return ( + _md5_hex(self.local, (number - 1) * self.part_size, length) == etag.lower() + ) + + def _create(self) -> None: + self.upload_id = self.t._call("create_multipart_upload", Key=self.key)[ + "UploadId" + ] + self.parts = {} + try: + _write_json(self.state_path, self._state()) + except OSError: + # without saved state no re-run can find this upload again + self.t._abort(self.key, self.upload_id) + raise + + def _send(self, number: int) -> Tuple[int, str]: + length = self._length(number) + with open(self.local, "rb") as handle: + handle.seek((number - 1) * self.part_size) + body = handle.read(length) + response = self.t._call( + "upload_part", + Key=self.key, + UploadId=self.upload_id, + PartNumber=number, + Body=body, + ) + self.t._emit("advance", str(self.local), length) + return number, response["ETag"] + + def _complete(self) -> None: + parts = [{"PartNumber": n, "ETag": self.parts[n]} for n in sorted(self.parts)] + try: + self.t._call( + "complete_multipart_upload", + Key=self.key, + UploadId=self.upload_id, + MultipartUpload={"Parts": parts}, + ) + except Exception as exc: # noqa: BLE001 - re-raised unless it completed + if not self._completed_anyway(exc, _multipart_etag(parts)): + raise + + def _completed_anyway(self, exc: Exception, want_etag: str) -> bool: + """whether a failed complete took effect anyway. assembling a large + file can outlive the proxy (524) while it still finishes server-side, + and a retry then finds the upload gone. only those ambiguous failures + qualify, and only when the object at the key is this upload's: its + size matches and its etag equals the multipart etag of our parts.""" + if not want_etag or not (is_resumable(exc) or _is_not_found(exc)): + return False + head = self.t.head(self.key) + return bool( + head + and int(head.get("ContentLength", -1)) == self.stat.st_size + and _etag(head.get("ETag")) == want_etag + ) + class _Download: """one file download into `.rpdownload`, with its finished chunks recorded in a sidecar json.""" - def __init__(self, key: str, size: int, etag: str, dest: Path, modified: Any): + def __init__( + self, transfer: _Transfer, key: str, size: Any, meta: Dict[str, Any], dest: Path + ): + self.t = transfer self.key = key - self.size = size - self.etag = etag + self.size = int(size) + # the etag is only compared for equality: a multipart object's etag is + # "-", never the md5 of the file + self.etag = _etag(meta.get("ETag")) + self.modified = meta.get("LastModified") self.dest = dest - self.modified = modified - self.label = str(dest) - self.tmp = dest.with_name(dest.name + PARTIAL_SUFFIX) + self.partial = dest.with_name(dest.name + PARTIAL_SUFFIX) self.sidecar = dest.with_name(dest.name + SIDECAR_SUFFIX) - self.count = math.ceil(size / DOWNLOAD_CHUNK_SIZE) - self.completed_chunks: set = set() + self.count = math.ceil(self.size / DOWNLOAD_CHUNK_SIZE) + self.done: set = set() + + def run(self) -> None: + label = str(self.dest) + if self._unchanged_local(): + self.t._emit("skip", label, self.size) + return + self.t._emit("start", label, self.size) + self.dest.parent.mkdir(parents=True, exist_ok=True) + self._refuse_linked_work_files() + self.done = self._load_or_init() + reused = sum(self._length(i) for i in self.done) + if reused: + self.t._emit("reuse", label, reused) + self.t._run( + [ + lambda i=i: self._fetch(i) + for i in range(self.count) + if i not in self.done + ] + ) + if self.partial.stat().st_size != self.size: + raise NetworkVolumeTransferError( + f"{self.dest}: size mismatch after download", resumable=True + ) + os.replace(self.partial, self.dest) + if self.modified is not None: + # stamp the remote time so uploading the file back counts as unchanged + os.utime(self.dest, (time.time(), self.modified.timestamp())) + self.sidecar.unlink(missing_ok=True) + self.t._emit("done", label) - def refuse_linked_work_files(self) -> None: - """refuse a partial, sidecar, or sidecar temp file that exists as - anything but a regular file: a symlink would aim the write elsewhere.""" + def _length(self, index: int) -> int: + return min(DOWNLOAD_CHUNK_SIZE, self.size - index * DOWNLOAD_CHUNK_SIZE) + + def _unchanged_local(self) -> bool: + """downloads stamp the remote time, so an unchanged file is the same + size and not older than the remote copy.""" + if ( + self.t.force + or not self.dest.is_file() + or self.dest.stat().st_size != self.size + ): + return False + return self.modified is None or _not_older( + self.dest.stat().st_mtime, self.modified.timestamp() + ) + + def _refuse_linked_work_files(self) -> None: + """refuse a working file that exists as anything but a regular file: + a symlink would aim the write elsewhere.""" for suffix in TRANSFER_SUFFIXES: path = self.dest.with_name(self.dest.name + suffix) if path.is_symlink() or (path.exists() and not path.is_file()): @@ -883,165 +718,128 @@ def refuse_linked_work_files(self) -> None: f"{path} exists and is not a regular file; remove it and rerun" ) - def chunk_length(self, index: int) -> int: - return min(DOWNLOAD_CHUNK_SIZE, self.size - index * DOWNLOAD_CHUNK_SIZE) + def _state(self) -> Dict[str, Any]: + return { + "key": self.key, + "size": self.size, + "etag": self.etag, + "chunk_size": DOWNLOAD_CHUNK_SIZE, + "done": sorted(self.done), + } - def load_or_init(self) -> set: + def _load_or_init(self) -> set: """the chunks an earlier run finished, or a fresh partial file.""" - state = _read_json(self.sidecar) - if self._resumes(state): - return {int(i) for i in state.get("done", []) if 0 <= int(i) < self.count} + state = _read_json(self.sidecar) or {} + saved = state.pop("done", []) + if ( + state == {k: v for k, v in self._state().items() if k != "done"} + and self.partial.is_file() + and self.partial.stat().st_size == self.size + ): + return {int(i) for i in saved if 0 <= int(i) < self.count} # drop the old sidecar first: a stop between emptying the partial and # saving fresh state must not leave it vouching for discarded chunks - _remove(self.sidecar) - with open(self.tmp, "wb") as handle: + self.sidecar.unlink(missing_ok=True) + with open(self.partial, "wb") as handle: handle.truncate(self.size) - self.save_state() + _write_json(self.sidecar, self._state()) return set() - def _resumes(self, state: Optional[Dict[str, Any]]) -> bool: - # the etag is only compared for equality. a multipart object's etag - # is "-", never the md5 of the file. - return bool( - state - and state.get("key") == self.key - and state.get("size") == self.size - and state.get("etag") == self.etag - and state.get("chunk_size") == DOWNLOAD_CHUNK_SIZE - and self.tmp.is_file() - and self.tmp.stat().st_size == self.size - ) - - def save_state(self) -> None: - _write_json_atomic( - self.sidecar, - { - "key": self.key, - "size": self.size, - "etag": self.etag, - "chunk_size": DOWNLOAD_CHUNK_SIZE, - "done": sorted(self.completed_chunks), - }, + def _fetch(self, index: int) -> None: + start = index * DOWNLOAD_CHUNK_SIZE + length = self._length(index) + data = self._get_range(start, length) + with open(self.partial, "r+b") as handle: + handle.seek(start) + handle.write(data) + handle.flush() + # on disk before the sidecar claims it + os.fsync(handle.fileno()) + with self.t._lock: + self.done.add(index) + _write_json(self.sidecar, self._state()) + self.t._emit("advance", str(self.dest), length) + + def _get_range(self, start: int, length: int) -> bytes: + """botocore retries the request but not the body read, so a dropped + or short body is fetched again (three tries).""" + end = start + length - 1 + for _ in range(3): + response = self.t._call( + "get_object", Key=self.key, Range=f"bytes={start}-{end}" + ) + if self.etag and _etag(response.get("ETag")) not in ("", self.etag): + raise NetworkVolumeTransferError( + f"{self.key} changed on the volume during download; " + "run the command again to restart it", + resumable=True, + ) + try: + data = response["Body"].read() + except Exception as exc: # noqa: BLE001 - a dropped stream is retried + if not is_resumable(exc): + raise + continue + if len(data) == length: + return data + raise NetworkVolumeTransferError( + f"could not read {self.key} bytes {start}-{end}", resumable=True ) -def _same_version(a: os.stat_result, b: os.stat_result) -> bool: - """whether two stats of a file show the same size and mtime.""" - return a.st_size == b.st_size and a.st_mtime_ns == b.st_mtime_ns - - -def _join_key(prefix: str, rel: str) -> str: - return f"{prefix}/{rel}" if prefix else rel - - -def _file_key(local: Path, remote_path: Optional[str]) -> str: - """the key for a single-file upload. it defaults to the local name, and - a remote path ending in "/" keeps the local name inside that folder.""" - if remote_path is None: - return local.name - key = remote_path.lstrip("/") - return key + local.name if remote_path.endswith("/") else key - - -def _object_dest(key: str, local_path: Any) -> Path: - """the local file for a single-object download. a trailing slash names - a folder to download into, as on upload.""" - name = PurePosixPath(key).name - dest = Path(local_path) if local_path is not None else Path(name) - if dest.is_dir() or str(local_path or "").endswith(("/", os.sep)): - return dest / name - return dest - - -def _write_at(path: Path, offset: int, data: bytes) -> None: - """write and sync data at offset, so it is on disk before the sidecar - claims it.""" - with open(path, "r+b") as handle: - handle.seek(offset) - handle.write(data) - handle.flush() - os.fsync(handle.fileno()) - - -def _fold(name: str) -> str: - """the key under which two spellings a filesystem may treat as one name - collide: nfc first (apfs compares names normalization-insensitively), - then unicode case folding.""" - return unicodedata.normalize("NFC", name).casefold() - - def _check_local_names(targets: List[Path], root: Path) -> None: - """refuse targets that would share a path on disk: two targets that are - one name on a case-insensitive or normalizing filesystem ("Model.bin" and + """refuse targets that would share a path on disk: two that are one name + on a case-insensitive or normalizing filesystem ("Model.bin" and "model.bin"), a file whose name is also a folder another file lives in - ("a" beside "A/b"), or a file whose partial, sidecar, or sidecar temp - name is also a target or a directory a target lives in ("x" beside - "x.rpdownload" or "x.rpdownload/y"). that can refuse a listing a - case-sensitive filesystem would take, which is the safe direction.""" - files, dirs = _index_names(targets, root) - _check_file_folder_clash(files, dirs) - _check_transfer_names(targets, root, files, dirs) - - -def _index_names(targets: List[Path], root: Path) -> Tuple[Dict[str, Path], Dict[str, Path]]: - """targets and the folders they live in, keyed by folded name.""" - files: Dict[str, Path] = {} - dirs: Dict[str, Path] = {} - for target in targets: - rel = target.relative_to(root) - key = _fold(str(rel)) - if key in files: + ("a" beside "A/b"), or a file whose partial, sidecar or sidecar temp name + is also a target or a folder one lives in ("x" beside "x.rpdownload/y"). + names are compared nfc-normalized (apfs ignores normalization) and case + folded, which can refuse a listing a case-sensitive filesystem would + take: the safe direction.""" + + def fold(path: PurePosixPath) -> str: + return unicodedata.normalize("NFC", str(path)).casefold() + + rels = [PurePosixPath(target.relative_to(root).as_posix()) for target in targets] + files: Dict[str, PurePosixPath] = {} + dirs: Dict[str, PurePosixPath] = {} + for rel in rels: + if fold(rel) in files: raise NetworkVolumeTransferError( - f"cannot download both {files[key]} and {target}: they would be one " + f"cannot download both {files[fold(rel)]} and {rel}: they would be one " "file on a case-insensitive or normalizing filesystem" ) - files[key] = target - dirs.update({_fold(str(parent)): target for parent in rel.parents if str(parent) != "."}) - return files, dirs - - -def _check_file_folder_clash(files: Dict[str, Path], dirs: Dict[str, Path]) -> None: - clash = next((key for key in files if key in dirs), None) - if clash is not None: - raise NetworkVolumeTransferError( - f"cannot download both {files[clash]} and {dirs[clash]}: " - "the first would have to be a file and a folder" - ) - - -def _check_transfer_names( - targets: List[Path], root: Path, files: Dict[str, Path], dirs: Dict[str, Path] -) -> None: - """refuse a target whose download working files would land on another - target or a folder one lives in.""" - for target in targets: - rel = target.relative_to(root) + files[fold(rel)] = rel + dirs.update({fold(parent): rel for parent in rel.parents if str(parent) != "."}) + for rel in rels: + if fold(rel) in dirs: + raise NetworkVolumeTransferError( + f"cannot download both {rel} and {dirs[fold(rel)]}: " + "the first would have to be a file and a folder" + ) for suffix in TRANSFER_SUFFIXES: - twin = _fold(str(rel.with_name(rel.name + suffix))) + twin = fold(rel.with_name(rel.name + suffix)) other = files.get(twin) or dirs.get(twin) if other is not None: raise NetworkVolumeTransferError( - f"cannot download both {target} and {other}: " + f"cannot download both {rel} and {other}: " f"{rel.name + suffix} is used for the first's partial download" ) def _safe_target(root: Path, rel: str) -> Path: - """local path for a remote relative key, refusing anything outside root.""" - parts = PurePosixPath(rel).parts + """local path for a remote relative key, refusing anything outside root. + "." and empty segments are refused too: "a/./b" would alias "a/b".""" + segments = rel.split("/") if ( - not rel - or rel.startswith("/") + rel.startswith("/") or "\\" in rel - # PurePosixPath drops "." and empty segments; "a/./b" would alias "a/b" - or any(seg in ("..", ".", "") for seg in rel.split("/")) - or any(part in ("..", "") for part in parts) - or (parts and ":" in parts[0]) + or any(seg in ("..", ".", "") for seg in segments) + or ":" in segments[0] ): raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") - target = root.joinpath(*parts) - resolved_root = root.resolve() - if resolved_root not in target.resolve().parents: + target = root.joinpath(*segments) + if root.resolve() not in target.resolve().parents: raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") return target @@ -1051,11 +849,17 @@ def _safe_target(root: Path, rel: str) -> Path: def _resolve_volume(reference: str) -> Tuple[str, str]: """the id and datacenter of the network volume with this id or name.""" - from runpod.api.rest import run_rest_paginated_request_async # pylint: disable=import-outside-toplevel + from runpod.api.rest import ( + run_rest_paginated_request_async, + ) # pylint: disable=import-outside-toplevel from runpod.apps.context import block # pylint: disable=import-outside-toplevel - from runpod.apps.utils.lookup import find_by_id_or_name # pylint: disable=import-outside-toplevel + from runpod.apps.utils.lookup import ( + find_by_id_or_name, + ) # pylint: disable=import-outside-toplevel - volumes = block(run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes")) + volumes = block( + run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes") + ) record = find_by_id_or_name( volumes, reference, noun="network volumes", error=NetworkVolumeTransferError ) @@ -1069,31 +873,20 @@ def _resolve_volume(reference: str) -> Tuple[str, str]: return record["id"], str(datacenter) -def _check_endpoint_resolves(endpoint: str) -> None: - """fail fast on a datacenter whose s3 endpoint does not resolve, before - any transfer starts. once a transfer is running, a failed lookup is a - network drop and is retried like any other. behind an https proxy (which - botocore honours from the environment) the proxy resolves the name, so - the local lookup is skipped.""" - from urllib.parse import urlparse # pylint: disable=import-outside-toplevel - from urllib.request import getproxies, proxy_bypass # pylint: disable=import-outside-toplevel - - parsed = urlparse(endpoint) - host = parsed.hostname - if not host: - raise NetworkVolumeTransferError(f"invalid s3 endpoint {endpoint!r}") - if getproxies().get(parsed.scheme or "https") and not proxy_bypass(host): - return - try: - socket.getaddrinfo(host, 443) - except socket.gaierror as exc: - raise NetworkVolumeTransferError( - f"the s3 endpoint {host} does not resolve: the volume's datacenter " - "may not support the s3 api, or check your network connection" - ) from exc - +def _make_transfer( + volume: str, + access_key: Optional[str], + secret_key: Optional[str], + endpoint_url: Optional[str], + concurrency: int, + progress: Optional[Progress], + force: bool = False, +) -> _Transfer: + """a transfer to the network volume with this id or name. backs + NetworkVolume.upload and NetworkVolume.download.""" + import boto3 # pylint: disable=import-outside-toplevel + from botocore.config import Config # pylint: disable=import-outside-toplevel -def _credentials(access_key: Optional[str], secret_key: Optional[str]) -> Tuple[str, str]: access_key = access_key or os.environ.get(ACCESS_KEY_ENV) secret_key = secret_key or os.environ.get(SECRET_KEY_ENV) if not access_key or not secret_key: @@ -1101,56 +894,28 @@ def _credentials(access_key: Optional[str], secret_key: Optional[str]) -> Tuple[ f"s3 api keys required: set {ACCESS_KEY_ENV} and {SECRET_KEY_ENV} " f"(create keys at {S3_KEYS_URL})" ) - return access_key, secret_key - - -def _client_config(concurrency: int) -> Any: - from botocore.config import Config # pylint: disable=import-outside-toplevel - - return Config( + volume_id, datacenter = _resolve_volume(volume) + endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" + concurrency = max(1, int(concurrency)) + config = Config( s3={"addressing_style": "path"}, request_checksum_calculation="when_required", response_checksum_validation="when_required", - # _with_retries owns retries; botocore's would multiply them and sleep - # while holding a concurrency slot - retries={"total_max_attempts": 1, "mode": "standard"}, + # total_max_attempts counts the first try (max_attempts would not) + retries={"mode": "standard", "total_max_attempts": 10}, connect_timeout=60, # a socket read silent this long is a stalled connection (runpod's # proxy answers a slow request with 524 after about 100s) read_timeout=300, - # every call holds one of `concurrency` slots; headroom for - # connections still being returned to the pool + # headroom over the request bound for connections being returned max_pool_connections=max(10, 2 * concurrency), ) - - -def _make_transfer( - volume: str, - access_key: Optional[str], - secret_key: Optional[str], - endpoint_url: Optional[str], - concurrency: int, - progress: Optional[Progress], - force: bool = False, -) -> _Transfer: - """a transfer to the network volume with this id or name. backs - NetworkVolume.upload and NetworkVolume.download.""" - import boto3 # pylint: disable=import-outside-toplevel - - access_key, secret_key = _credentials(access_key, secret_key) - volume_id, datacenter = _resolve_volume(volume) - endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" - if not endpoint_url: - # an explicit endpoint (a local stand-in, a custom route) is the - # caller's; only the derived datacenter endpoint is checked - _check_endpoint_resolves(endpoint) - concurrency = max(1, int(concurrency)) client = boto3.session.Session().client( "s3", region_name=datacenter, endpoint_url=endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, - config=_client_config(concurrency), + config=config, ) return _Transfer(client, volume_id, endpoint, concurrency, progress, force=force) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 7e8f1f00e..4a9d5d8ef 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -15,11 +15,6 @@ import runpod.network_volume_transfer as nvt import runpod.error -# the real helpers, before the autouse fixture stubs them -REAL_WAIT = nvt._wait -REAL_ENDPOINT_CHECK = nvt._check_endpoint_resolves - - def _client_error(code, status, op="Op"): return ClientError( {"Error": {"Code": code}, "ResponseMetadata": {"HTTPStatusCode": status}}, op @@ -156,10 +151,6 @@ def ops(self, op): @pytest.fixture(autouse=True) def small_sizes(monkeypatch): - monkeypatch.setattr(nvt, "_sleep", lambda _: None) - monkeypatch.setattr(nvt, "_wait", lambda delay, stop=None: None) - # transfers built through _make_transfer must not hit dns in tests - monkeypatch.setattr(nvt, "_check_endpoint_resolves", lambda endpoint: None) monkeypatch.setattr(nvt, "SINGLE_PUT_MAX", 10) monkeypatch.setattr(nvt, "DEFAULT_PART_SIZE", 8) monkeypatch.setattr(nvt, "DOWNLOAD_CHUNK_SIZE", 8) @@ -326,36 +317,6 @@ def test_complete_failure_without_object_keeps_state(tmp_path): assert len(list((tmp_path / "state").glob("*.json"))) == 1 -def test_transient_errors_are_retried(tmp_path): - fake = FakeS3() - src = tmp_path / "big.bin" - src.write_bytes(DATA) - failed = set() - - def flaky(op, kw): - if op == "upload_part" and kw["PartNumber"] not in failed: - failed.add(kw["PartNumber"]) - return _client_error("524", 524) - return None - - fake.fail = flaky - _transfer(fake, tmp_path).upload(src, "big.bin") - assert len(fake.ops("upload_part")) == 10 - assert fake.objects["big.bin"][0] == DATA - - -def test_retry_gives_up_after_attempts(): - calls = [] - - def always_503(): - calls.append(1) - raise _client_error("ServiceUnavailable", 503) - - with pytest.raises(ClientError): - nvt._with_retries(always_503, attempts=3) - assert len(calls) == 3 - - def test_folder_upload_skips_unchanged(tmp_path): fake = FakeS3() root = tmp_path / "data" @@ -519,6 +480,30 @@ def short_once(**kwargs): assert (tmp_path / "f.bin").read_bytes() == DATA +def test_download_dropped_stream_is_fetched_again(tmp_path): + from botocore.exceptions import IncompleteReadError + + fake = FakeS3() + fake.put("f.bin", DATA) + original = fake.get_object + dropped = [] + + class Dropped: + def read(self): + raise IncompleteReadError(actual_bytes=1, expected_bytes=8) + + def drop_once(**kwargs): + response = original(**kwargs) + if not dropped: + dropped.append(1) + return {"Body": Dropped(), "ETag": response["ETag"]} + return response + + fake.get_object = drop_once + _transfer(fake, tmp_path).download("f.bin", tmp_path / "f.bin") + assert (tmp_path / "f.bin").read_bytes() == DATA + + def test_download_file_into_existing_dir_and_skip(tmp_path): fake = FakeS3() fake.put("models/w.bin", b"weights") @@ -661,8 +646,8 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): assert transfer.client.meta.endpoint_url == "https://s3api-eu-ro-1.runpod.io/" assert transfer.client.meta.region_name == "EU-RO-1" assert transfer.client.meta.config.request_checksum_calculation == "when_required" - # _with_retries owns retries - assert transfer.client.meta.config.retries["total_max_attempts"] == 1 + # botocore owns retries + assert transfer.client.meta.config.retries == {"mode": "standard", "total_max_attempts": 10} credentials = transfer.client._request_signer._credentials assert credentials.access_key == "user_x" assert transfer.bucket == "vol123" @@ -772,14 +757,14 @@ def test_cli_ctrl_c_exits_130_with_resume_hint(): assert "interrupted; re-run the same command to resume" in result.output -def test_cli_unresolvable_endpoint_has_no_resume_hint(): +def test_cli_final_error_has_no_resume_hint(): from runpod.rp_cli.main import cli - err = nvt.NetworkVolumeTransferError("the s3 endpoint s3api-x.invalid does not resolve") + err = nvt.NetworkVolumeTransferError("s3 api keys required") with patch.object(runpod.NetworkVolume, "download", side_effect=err): result = CliRunner().invoke(cli, ["volume", "download", "vol", "a/b"]) assert result.exit_code == 1 - assert "does not resolve" in result.output + assert "keys required" in result.output assert "re-run" not in result.output @@ -841,7 +826,7 @@ def test_unsavable_state_aborts_the_upload(tmp_path, monkeypatch): def no_disk(path, data): raise OSError("disk full") - monkeypatch.setattr(nvt, "_write_json_atomic", no_disk) + monkeypatch.setattr(nvt, "_write_json", no_disk) with pytest.raises(OSError): _transfer(fake, tmp_path).upload(src, "big.bin") assert len(fake.ops("abort_multipart_upload")) == 1 @@ -868,7 +853,7 @@ def test_download_ignores_sidecar_of_another_key(tmp_path): { "key": "a.bin", "size": len(DATA), - "etag": nvt._strip_etag(fake.object_etag("b.bin")), + "etag": nvt._etag(fake.object_etag("b.bin")), "chunk_size": nvt.DOWNLOAD_CHUNK_SIZE, "done": [0, 1, 2, 3, 4], } @@ -924,26 +909,6 @@ def test_ambiguous_complete_without_verifiable_etag_fails(tmp_path): _transfer(fake, tmp_path).upload(src, "big.bin") -def test_single_put_detects_change_between_retries(tmp_path): - fake = FakeS3() - src = tmp_path / "small.txt" - src.write_bytes(b"original") - attempts = [] - - def change_then_fail(op, kw): - if op == "put_object" and not attempts: - attempts.append(1) - src.write_bytes(b"CHANGED!") - later = src.stat().st_mtime + 60 - os.utime(src, (later, later)) - return _client_error("ServiceUnavailable", 503) - return None - - fake.fail = change_then_fail - with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload"): - _transfer(fake, tmp_path).upload(src, "small.txt") - - def test_all_requests_share_the_concurrency_limit(tmp_path): import threading import time as _time @@ -972,61 +937,6 @@ def tracked_call(op, **kwargs): assert active["peak"] <= 2 -@pytest.fixture -def endpoint_check(monkeypatch): - """the real endpoint check with dns and proxies stubbed (the autouse - fixture replaces the check itself).""" - import socket as _socket - import urllib.request - - monkeypatch.setattr(nvt, "_check_endpoint_resolves", REAL_ENDPOINT_CHECK) - for name in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "NO_PROXY", "no_proxy"): - monkeypatch.delenv(name, raising=False) - # proxies from the environment only: macos and windows also read system settings - monkeypatch.setattr(urllib.request, "getproxies", urllib.request.getproxies_environment) - monkeypatch.setattr(urllib.request, "proxy_bypass", urllib.request.proxy_bypass_environment) - state = {"error": None} - - def fake_getaddrinfo(host, port, *args, **kwargs): - if state["error"] is not None: - raise state["error"] - return [(_socket.AF_INET, _socket.SOCK_STREAM, 6, "", ("127.0.0.1", port))] - - monkeypatch.setattr(_socket, "getaddrinfo", fake_getaddrinfo) - return state - - -def test_unresolvable_endpoint_fails_before_the_transfer(endpoint_check): - import socket as _socket - - endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") - with pytest.raises(nvt.NetworkVolumeTransferError, match="does not resolve"): - nvt._check_endpoint_resolves("https://s3api-nope-1.runpod.io/") - endpoint_check["error"] = None - nvt._check_endpoint_resolves("https://s3api-eu-ro-1.runpod.io/") - - -def test_endpoint_check_skipped_behind_a_proxy(endpoint_check, monkeypatch): - import socket as _socket - - endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") - monkeypatch.setenv("HTTPS_PROXY", "http://proxy.example:3128") - nvt._check_endpoint_resolves("https://s3api-eu-ro-1.runpod.io/") - - -def test_explicit_endpoint_skips_the_dns_check(endpoint_check, monkeypatch): - import socket as _socket - - monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") - monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") - monkeypatch.setattr(nvt, "_resolve_volume", lambda volume: ("vol123", "NOPE-1")) - endpoint_check["error"] = _socket.gaierror(8, "nodename nor servname provided") - transfer = nvt._make_transfer("vol123", None, None, "http://localhost:9000", 4, None) - assert transfer.client.meta.endpoint_url == "http://localhost:9000" - with pytest.raises(nvt.NetworkVolumeTransferError, match="does not resolve"): - nvt._make_transfer("vol123", None, None, None, 4, None) - - def test_dns_failure_mid_transfer_is_a_network_drop(): import socket as _socket @@ -1035,7 +945,6 @@ def test_dns_failure_mid_transfer_is_a_network_drop(): exc = EndpointConnectionError( endpoint_url="https://s3api-eu-ro-1.runpod.io/", error=_socket.gaierror(8, "nodename nor servname") ) - assert nvt._is_retryable(exc) assert nvt.is_resumable(exc) @@ -1101,53 +1010,29 @@ def test_download_refuses_file_that_is_also_a_folder(tmp_path): _transfer(fake, tmp_path).download("d/", tmp_path / "out") -def test_retry_backoff_ends_when_stopped(monkeypatch): - import threading - import time as _time - - monkeypatch.setattr(nvt, "_wait", REAL_WAIT) - # the real sleep too, so a revert to an uninterruptible sleep is caught - monkeypatch.setattr(nvt, "_sleep", _time.sleep) - stop = threading.Event() - calls = [] - - def fail(): - calls.append(1) - if len(calls) == 1: - stop.set() # ctrl-c arrives during the first attempt - raise ConnectionError("down") - - started = _time.monotonic() - with pytest.raises(ConnectionError): - nvt._with_retries(fail, attempts=3, stop=stop) - # the backoffs (1s, 2s plus jitter) are cut short - assert _time.monotonic() - started < 0.5 - - def test_single_put_change_with_older_mtime_is_resent_on_rerun(tmp_path): fake = FakeS3() src = tmp_path / "small.txt" src.write_bytes(b"original") old = src.stat().st_mtime - 3600 os.utime(src, (old, old)) - attempts = [] - - def edit_then_fail(op, kw): - if op == "put_object" and not attempts: - attempts.append(1) - # a same-size edit whose mtime is older than the object the retry - # writes, so the skip rule alone would keep the object - src.write_bytes(b"CHANGED!") - past = src.stat().st_mtime - 1800 - os.utime(src, (past, past)) - return _client_error("ServiceUnavailable", 503) - return None + original = fake.put_object - fake.fail = edit_then_fail - with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload"): + def put_then_edit(**kwargs): + result = original(**kwargs) + # a same-size edit whose mtime is older than the object just written, + # so the skip rule alone would keep the object + src.write_bytes(b"CHANGED!") + past = src.stat().st_mtime - 1800 + os.utime(src, (past, past)) + return result + + fake.put_object = put_then_edit + with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload") as info: _transfer(fake, tmp_path).upload(src, "small.txt") + assert info.value.resumable assert "small.txt" not in fake.objects - fake.fail = None + fake.put_object = original _transfer(fake, tmp_path).upload(src, "small.txt") assert fake.objects["small.txt"][0] == b"CHANGED!" @@ -1262,22 +1147,22 @@ def interrupt_first(**kwargs): def test_discarded_chunks_drop_the_sidecar_first(tmp_path): dest = tmp_path / "f.bin" - job = nvt._Download("f.bin", len(DATA), "e", dest, None) + job = nvt._Download(_transfer(FakeS3(), tmp_path), "f.bin", len(DATA), {"ETag": "e"}, dest) # a sidecar that vouches for every chunk beside a partial of the wrong size job.sidecar.write_text(json.dumps( {"key": "f.bin", "size": len(DATA), "etag": "e", "chunk_size": nvt.DOWNLOAD_CHUNK_SIZE, "done": list(range(job.count))} )) - job.tmp.write_bytes(b"short") + job.partial.write_bytes(b"short") sidecar_seen = [] real_open = open def watch_open(path, *args, **kwargs): - if path == job.tmp: + if path == job.partial: sidecar_seen.append(job.sidecar.exists()) return real_open(path, *args, **kwargs) with patch("builtins.open", watch_open): - assert job.load_or_init() == set() + assert job._load_or_init() == set() assert sidecar_seen == [False] From c29e0d2155a480409e8a94d41e562bfc54102897 Mon Sep 17 00:00:00 2001 From: Justin Date: Mon, 26 Oct 2026 20:48:57 -0400 Subject: [PATCH 17/23] fix: retry the 524 runpod's proxy returns on slow operations botocore's standard mode retries 500/502/503/504 only, so the switch to it stopped retrying the 524 a slow complete or upload gets. a needs-retry hook ahead of botocore's adds it. --- runpod/network_volume_transfer.py | 16 ++++++++++- tests/test_network_volume_transfer.py | 39 +++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 1 deletion(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 2f13cc900..af1e8c662 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -17,6 +17,7 @@ import json import math import os +import random import threading import time import unicodedata @@ -37,6 +38,9 @@ # every file a download writes next to its target TRANSFER_SUFFIXES = (PARTIAL_SUFFIX, SIDECAR_SUFFIX, SIDECAR_SUFFIX + ".tmp") +# tries per request, the first included +MAX_ATTEMPTS = 10 + ACCESS_KEY_ENV = "RUNPOD_S3_ACCESS_KEY" SECRET_KEY_ENV = "RUNPOD_S3_SECRET_KEY" S3_KEYS_URL = "https://console.runpod.io/user/settings" @@ -902,7 +906,7 @@ def _make_transfer( request_checksum_calculation="when_required", response_checksum_validation="when_required", # total_max_attempts counts the first try (max_attempts would not) - retries={"mode": "standard", "total_max_attempts": 10}, + retries={"mode": "standard", "total_max_attempts": MAX_ATTEMPTS}, connect_timeout=60, # a socket read silent this long is a stalled connection (runpod's # proxy answers a slow request with 524 after about 100s) @@ -918,4 +922,14 @@ def _make_transfer( aws_secret_access_key=secret_key, config=config, ) + client.meta.events.register_first("needs-retry.s3", _retry_proxy_timeout) return _Transfer(client, volume_id, endpoint, concurrency, progress, force=force) + + +def _retry_proxy_timeout(response=None, attempts=1, **_kwargs) -> Optional[float]: + """the backoff before retrying a 524, which runpod's proxy returns when an + operation outlives it. botocore's standard mode retries 500/502/503/504 + only; returning None leaves every other response to it.""" + if response is None or response[0].status_code != 524 or attempts >= MAX_ATTEMPTS: + return None + return random.uniform(0, min(20.0, 2.0**attempts)) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 4a9d5d8ef..7102066d4 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -1180,3 +1180,42 @@ def test_download_refuses_symlinked_work_files(tmp_path, suffix): _transfer(fake, tmp_path).download("x", tmp_path / "x") assert not nvt.is_resumable(info.value) assert other.read_bytes() == b"keep me" + + +def test_client_retries_the_proxy_524(monkeypatch): + import http.server + import threading + + hits = [] + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): # noqa: N802 - http.server naming + hits.append(self.path) + if len(hits) <= 2: + self.send_response(524) + self.send_header("Content-Length", "0") + self.end_headers() + return + body = b"" + self.send_response(200) + self.send_header("Content-Type", "application/xml") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, *args): + pass + + server = http.server.HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") + monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") + monkeypatch.setattr(nvt, "_resolve_volume", lambda volume: ("vol123", "EU-RO-1")) + monkeypatch.setattr(nvt.random, "uniform", lambda a, b: 0) + endpoint = f"http://127.0.0.1:{server.server_port}" + transfer = nvt._make_transfer("vol123", None, None, endpoint, 1, None) + transfer.client.list_objects_v2(Bucket="vol123") + finally: + server.shutdown() + assert len(hits) == 3 From a8431ea29bd3e497517ce8d0047583d715003249 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 27 Oct 2026 14:26:57 -0400 Subject: [PATCH 18/23] feat(cli): prompt for login, volume and s3 keys; explain transfer failures rp volume upload/download now asks to log in when no api key is set, lets you pick a volume when VOLUME is left out, and asks for s3 keys (offering to save them in ~/.runpod/config.toml) when none are set. in a script each of these fails with a message naming the fix. rejected keys and a datacenter without the s3 api get their own non-resumable errors. also from review: a request and its body or part read share one slot, cleanup takes a slot too, retries are decided by one stop-aware hook (capped backoff, 524 included, no shared quota), an edit during a multipart complete removes the object, symlinked download subfolders are refused, state writes sync their folder, and an opt-in live e2e test round-trips a temporary volume. --- runpod/cli/groups/config/functions.py | 64 +++- runpod/network_volume_transfer.py | 341 +++++++++++++----- runpod/rp_cli/volume.py | 148 ++++++-- .../e2e/test_network_volume_transfer_live.py | 74 ++++ tests/test_network_volume_transfer.py | 340 ++++++++++++++++- 5 files changed, 824 insertions(+), 143 deletions(-) create mode 100644 tests/e2e/test_network_volume_transfer_live.py diff --git a/runpod/cli/groups/config/functions.py b/runpod/cli/groups/config/functions.py index 72e9b738c..4e7d3e513 100644 --- a/runpod/cli/groups/config/functions.py +++ b/runpod/cli/groups/config/functions.py @@ -29,29 +29,57 @@ def set_credentials(api_key: str, profile: str = "default", overwrite=False) -> [default] api_key = "RUNPOD_API_KEY" """ - cred_dir = os.path.dirname(CREDENTIAL_FILE) - os.makedirs(cred_dir, exist_ok=True) - Path(CREDENTIAL_FILE).touch(exist_ok=True) + config = _load_config() + if not overwrite and profile in config: + raise ValueError( + "Profile already exists. Use set_credentials(overwrite=True) to update." + ) + _set_profile_values(config, profile, {"api_key": api_key}) + _write_config(config) + + +def set_s3_keys(access_key: str, secret_key: str, profile: str = "default") -> None: + """ + Saves S3 API keys for network volume transfers in ~/.runpod/config.toml, + next to the profile's api_key. + + --- File Structure --- + + [default] + api_key = "RUNPOD_API_KEY" + s3_access_key = "..." + s3_secret_key = "..." + """ + config = _load_config() + _set_profile_values( + config, profile, {"s3_access_key": access_key, "s3_secret_key": secret_key} + ) + _write_config(config) + +def _load_config() -> tomlkit.TOMLDocument: + """the config file as an editable document, created empty if missing.""" + os.makedirs(os.path.dirname(CREDENTIAL_FILE), exist_ok=True) + Path(CREDENTIAL_FILE).touch(exist_ok=True) with open(CREDENTIAL_FILE, "r", encoding="UTF-8") as cred_file: - try: - content = cred_file.read() - config = ( - tomlkit.parse(content) - if content.strip() - else tomlkit.document() - ) - except tomlkit.exceptions.ParseError as exc: - raise ValueError("~/.runpod/config.toml is not a valid TOML file.") from exc + content = cred_file.read() + try: + return tomlkit.parse(content) if content.strip() else tomlkit.document() + except tomlkit.exceptions.ParseError as exc: + raise ValueError("~/.runpod/config.toml is not a valid TOML file.") from exc - if not overwrite: - if profile in config: - raise ValueError( - "Profile already exists. Use set_credentials(overwrite=True) to update." - ) - config[profile] = {"api_key": api_key} +def _set_profile_values(config: tomlkit.TOMLDocument, profile: str, values: dict) -> None: + """set values in a profile, keeping its other keys.""" + if profile in config: + config[profile].update(values) + else: + config[profile] = values + +def _write_config(config: tomlkit.TOMLDocument) -> None: + """replace the config file atomically; mkstemp keeps it owner-only.""" + cred_dir = os.path.dirname(CREDENTIAL_FILE) fd, tmp_path = tempfile.mkstemp(dir=cred_dir, suffix=".toml") try: with os.fdopen(fd, "w", encoding="UTF-8") as tmp_file: diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index af1e8c662..fbd811510 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -9,8 +9,8 @@ only the parts the server does not already hold. downloads fetch 64 MiB ranges into `.rpdownload`, recording finished ranges in a sidecar json, so a re-run fetches only what is missing. folder transfers skip -files that are already present and unchanged. botocore retries failed -requests. +files that are already present and unchanged. each client retries failed +requests through `_Transfer._retry_delay`. """ import hashlib @@ -18,12 +18,15 @@ import math import os import random +import socket import threading import time import unicodedata from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path, PurePosixPath -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar + +T = TypeVar("T") MIB = 1024 * 1024 SINGLE_PUT_MAX = 64 * MIB @@ -44,6 +47,10 @@ ACCESS_KEY_ENV = "RUNPOD_S3_ACCESS_KEY" SECRET_KEY_ENV = "RUNPOD_S3_SECRET_KEY" S3_KEYS_URL = "https://console.runpod.io/user/settings" +S3_DOCS_URL = "https://docs.runpod.io/storage/s3-api" +NOT_LOGGED_IN = "not logged in: run `rp login` or set RUNPOD_API_KEY" +# error codes s3 gives for keys it does not accept +_REJECTED_KEY_CODES = ("InvalidAccessKeyId", "SignatureDoesNotMatch", "AccessDenied") # progress(event, path, nbytes). events: "start" (nbytes = file size), # "advance" (nbytes sent or received), "reuse" (nbytes an earlier run already @@ -71,13 +78,39 @@ def part_size_for(size: int) -> int: return part_size -def _error(exc: BaseException) -> Tuple[str, int]: - """the s3 error code and http status of a botocore ClientError.""" - response = getattr(exc, "response", None) or {} +def _error_fields(response: Dict[str, Any]) -> Tuple[str, int]: + """the s3 error code and http status in a parsed botocore response.""" code = str(response.get("Error", {}).get("Code", "")) return code, int(response.get("ResponseMetadata", {}).get("HTTPStatusCode") or 0) +def _error(exc: BaseException) -> Tuple[str, int]: + """the s3 error code and http status of a botocore ClientError.""" + return _error_fields(getattr(exc, "response", None) or {}) + + +def _transient(code: str, status: int) -> bool: + """a server answer that may pass on its own: throttling and 5xx, + including the 524 runpod's proxy returns when an operation outlives it.""" + return ( + status == 429 + or status >= 500 + or code in ("SlowDown", "RequestTimeout", "Throttling") + ) + + +def _keys_rejected(exc: BaseException) -> Optional[NetworkVolumeTransferError]: + """the error to raise when s3 refused the keys, or None.""" + code, status = _error(exc) + if status not in (401, 403) and code not in _REJECTED_KEY_CODES: + return None + return NetworkVolumeTransferError( + f"the volume's s3 api rejected the keys ({status or code}): check that " + f"{ACCESS_KEY_ENV} and {SECRET_KEY_ENV} are current and belong to this " + f"account (create keys at {S3_KEYS_URL})" + ) + + def _is_not_found(exc: BaseException) -> bool: code, status = _error(exc) return status == 404 or code in ("404", "NoSuchKey", "NotFound") @@ -110,11 +143,7 @@ def is_resumable(exc: BaseException) -> bool: # NoSuchUpload: the upload vanished (a complete whose 524 hid that it # finished); 416: a range that no longer fits an object that changed code, status = _error(exc) - return ( - code in ("NoSuchUpload", "SlowDown", "RequestTimeout") - or status in (416, 429) - or status >= 500 - ) + return code == "NoSuchUpload" or status == 416 or _transient(code, status) def _read_json(path: Path) -> Optional[Dict[str, Any]]: @@ -135,6 +164,22 @@ def _write_json(path: Path, data: Dict[str, Any]) -> None: handle.flush() os.fsync(handle.fileno()) os.replace(tmp, path) + _sync_dir(path.parent) + + +def _sync_dir(path: Path) -> None: + """make a rename in path survive a power loss. best effort: windows + cannot open a directory.""" + try: + fd = os.open(path, os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + except OSError: + pass + finally: + os.close(fd) def _md5_hex(path: Path, offset: int, length: int) -> str: @@ -221,11 +266,48 @@ def __init__( # set on ctrl-c so queued work stops instead of finishing files in flight self._stop = threading.Event() - def _call(self, method: str, **kwargs: Any) -> Any: + def _request(self, work: Callable[[], T]) -> T: + """run work, one request plus its body or part read, holding a slot, + so `concurrency` also bounds the bytes in flight.""" with self._slots: if self._stop.is_set(): raise NetworkVolumeTransferError("interrupted", resumable=True) - return getattr(self.client, method)(Bucket=self.bucket, **kwargs) + try: + return work() + except Exception as exc: + rejected = _keys_rejected(exc) + if rejected is None: + raise + raise rejected from exc + + def _client_call(self, method: str, **kwargs: Any) -> Any: + return getattr(self.client, method)(Bucket=self.bucket, **kwargs) + + def _call(self, method: str, **kwargs: Any) -> Any: + return self._request(lambda: self._client_call(method, **kwargs)) + + def _cleanup(self, method: str, **kwargs: Any) -> Any: + """a cleanup request (delete, abort). it takes a slot but skips the + stop check, so it still runs after ctrl-c.""" + with self._slots: + return self._client_call(method, **kwargs) + + def _retry_delay( + self, response=None, caught_exception=None, attempts=1, **_kwargs + ) -> Optional[float]: + """botocore's needs-retry hook for this transfer's client, whose own + retries are off: botocore's standard mode skips 524 and shares one + retry quota across a whole run. waits here, so ctrl-c ends the wait, + and returns 0 to retry or None not to.""" + if caught_exception is not None: + failed = is_resumable(caught_exception) + else: + failed = response is not None and _transient(*_error_fields(response[1])) + if not failed or attempts >= MAX_ATTEMPTS: + return None + if self._stop.wait(random.uniform(0, min(20.0, 2.0**attempts))): + raise NetworkVolumeTransferError("interrupted", resumable=True) + return 0 def _run(self, tasks: List[Callable[[], Any]]) -> List[Any]: """run tasks in parallel; the first failure cancels the ones not started.""" @@ -400,7 +482,14 @@ def upload_file(self, local: Path, key: str) -> None: self._emit("done", str(local)) def _single_put(self, local: Path, key: str, stat: os.stat_result) -> None: - self._call("put_object", Key=key, Body=local.read_bytes()) + self._request( + lambda: self._client_call("put_object", Key=key, Body=local.read_bytes()) + ) + self._keep_if_unchanged(local, key, stat) + self._emit("advance", str(local), stat.st_size) + + def _keep_if_unchanged(self, local: Path, key: str, stat: os.stat_result) -> None: + """after an upload, check the file did not change while it was read.""" try: now = local.stat() except OSError as exc: @@ -410,15 +499,13 @@ def _single_put(self, local: Path, key: str, stat: os.stat_result) -> None: ) from exc if not _same_version(stat, now): self._delete_changed(local, key) - self._emit("advance", str(local), stat.st_size) def _delete_changed(self, local: Path, key: str) -> None: """the file changed while it was read, so the object may hold either version and is newer than the local edit; a re-run would skip it. - remove it so the re-run sends the file. not via _call: this cleanup - runs even after ctrl-c.""" + remove it so the re-run sends the file.""" try: - self.client.delete_object(Bucket=self.bucket, Key=key) + self._cleanup("delete_object", Key=key) except Exception as exc: # noqa: BLE001 raise NetworkVolumeTransferError( f"{local} changed during upload and the remote copy could not be " @@ -432,9 +519,7 @@ def _abort(self, key: str, upload_id: str) -> None: """discard an upload that can no longer be resumed (best effort: a failed abort only leaves orphaned parts behind).""" try: - self.client.abort_multipart_upload( - Bucket=self.bucket, Key=key, UploadId=upload_id - ) + self._cleanup("abort_multipart_upload", Key=key, UploadId=upload_id) except Exception: # noqa: BLE001 return @@ -526,6 +611,8 @@ def run(self) -> None: ) self._complete() self.state_path.unlink(missing_ok=True) + # completing can take long enough for an edit to land meanwhile + self.t._keep_if_unchanged(self.local, self.key, self.stat) def _length(self, number: int) -> int: return min(self.part_size, self.stat.st_size - (number - 1) * self.part_size) @@ -603,19 +690,21 @@ def _create(self) -> None: raise def _send(self, number: int) -> Tuple[int, str]: - length = self._length(number) + response = self.t._request(lambda: self._read_and_send(number)) + self.t._emit("advance", str(self.local), self._length(number)) + return number, response["ETag"] + + def _read_and_send(self, number: int) -> Dict[str, Any]: with open(self.local, "rb") as handle: handle.seek((number - 1) * self.part_size) - body = handle.read(length) - response = self.t._call( + body = handle.read(self._length(number)) + return self.t._client_call( "upload_part", Key=self.key, UploadId=self.upload_id, PartNumber=number, Body=body, ) - self.t._emit("advance", str(self.local), length) - return number, response["ETag"] def _complete(self) -> None: parts = [{"PartNumber": n, "ETag": self.parts[n]} for n in sorted(self.parts)] @@ -767,28 +856,32 @@ def _fetch(self, index: int) -> None: def _get_range(self, start: int, length: int) -> bytes: """botocore retries the request but not the body read, so a dropped or short body is fetched again (three tries).""" - end = start + length - 1 for _ in range(3): - response = self.t._call( - "get_object", Key=self.key, Range=f"bytes={start}-{end}" - ) - if self.etag and _etag(response.get("ETag")) not in ("", self.etag): - raise NetworkVolumeTransferError( - f"{self.key} changed on the volume during download; " - "run the command again to restart it", - resumable=True, - ) - try: - data = response["Body"].read() - except Exception as exc: # noqa: BLE001 - a dropped stream is retried - if not is_resumable(exc): - raise - continue - if len(data) == length: + data = self.t._request(lambda: self._read_range(start, length)) + if data is not None and len(data) == length: return data raise NetworkVolumeTransferError( - f"could not read {self.key} bytes {start}-{end}", resumable=True + f"could not read {self.key} bytes {start}-{start + length - 1}", + resumable=True, + ) + + def _read_range(self, start: int, length: int) -> Optional[bytes]: + """one GET and its body read; None when the body dropped.""" + response = self.t._client_call( + "get_object", Key=self.key, Range=f"bytes={start}-{start + length - 1}" ) + if self.etag and _etag(response.get("ETag")) not in ("", self.etag): + raise NetworkVolumeTransferError( + f"{self.key} changed on the volume during download; " + "run the command again to restart it", + resumable=True, + ) + try: + return response["Body"].read() + except Exception as exc: # noqa: BLE001 - a dropped stream is retried + if not is_resumable(exc): + raise + return None def _check_local_names(targets: List[Path], root: Path) -> None: @@ -845,30 +938,53 @@ def _safe_target(root: Path, rel: str) -> Path: target = root.joinpath(*segments) if root.resolve() not in target.resolve().parents: raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") + # a symlinked folder would let two remote paths share one local file, + # which the name collision check cannot see + for depth in range(1, len(segments)): + if root.joinpath(*segments[:depth]).is_symlink(): + raise NetworkVolumeTransferError( + f"{root.joinpath(*segments[:depth])} is a symlink; " + "download into a folder without symlinked subfolders" + ) return target # ---------------------------------------------------------------- client -def _resolve_volume(reference: str) -> Tuple[str, str]: - """the id and datacenter of the network volume with this id or name.""" - from runpod.api.rest import ( +def _list_volumes() -> List[Dict[str, Any]]: + """every network volume of the account.""" + from runpod import error # pylint: disable=import-outside-toplevel + from runpod.api.rest import ( # pylint: disable=import-outside-toplevel run_rest_paginated_request_async, - ) # pylint: disable=import-outside-toplevel + ) from runpod.apps.context import block # pylint: disable=import-outside-toplevel - from runpod.apps.utils.lookup import ( - find_by_id_or_name, - ) # pylint: disable=import-outside-toplevel - volumes = block( - run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes") + try: + return block( + run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes") + ) + except error.AuthenticationError as exc: + raise NetworkVolumeTransferError(NOT_LOGGED_IN) from exc + + +def _resolve_volume(reference: str) -> Tuple[str, str]: + """the id and datacenter of the network volume with this id or name.""" + from runpod.apps.utils.lookup import ( # pylint: disable=import-outside-toplevel + find_by_id_or_name, ) + record = find_by_id_or_name( - volumes, reference, noun="network volumes", error=NetworkVolumeTransferError + _list_volumes(), + reference, + noun="network volumes", + error=NetworkVolumeTransferError, ) if record is None: - raise NetworkVolumeTransferError(f"network volume {reference!r} not found") + raise NetworkVolumeTransferError( + f"no network volume named {reference!r}: pass a volume id or name " + "(see https://console.runpod.io/user/storage)" + ) datacenter = record.get("dataCenter") or record.get("dataCenterId") if not datacenter: raise NetworkVolumeTransferError( @@ -877,6 +993,81 @@ def _resolve_volume(reference: str) -> Tuple[str, str]: return record["id"], str(datacenter) +def saved_s3_keys() -> Tuple[Optional[str], Optional[str]]: + """the s3 keys saved in ~/.runpod/config.toml for the current profile.""" + import runpod # pylint: disable=import-outside-toplevel,cyclic-import + from runpod.cli.groups.config.functions import ( # pylint: disable=import-outside-toplevel + get_credentials, + ) + + saved = get_credentials(runpod.profile) or {} + return saved.get("s3_access_key"), saved.get("s3_secret_key") + + +def find_s3_keys( + access_key: Optional[str] = None, secret_key: Optional[str] = None +) -> Optional[Tuple[str, str]]: + """the first complete pair of s3 keys from the arguments, the + environment, or the saved config; None when there is none. a pair never + mixes sources.""" + pairs = ( + (access_key, secret_key), + (os.environ.get(ACCESS_KEY_ENV), os.environ.get(SECRET_KEY_ENV)), + saved_s3_keys(), + ) + return next(((a, s) for a, s in pairs if a and s), None) + + +def _s3_keys(access_key: Optional[str], secret_key: Optional[str]) -> Tuple[str, str]: + keys = find_s3_keys(access_key, secret_key) + if keys is None: + raise NetworkVolumeTransferError( + f"s3 api keys required: set {ACCESS_KEY_ENV} and {SECRET_KEY_ENV} " + f"or pass access_key and secret_key (create keys at {S3_KEYS_URL})" + ) + return keys + + +def _check_endpoint(endpoint: str, datacenter: str) -> None: + """fail fast on a datacenter without the s3 api: its endpoint does not + resolve, and retrying would only delay the error. behind an https proxy + the proxy resolves the name, so the check is skipped.""" + from urllib.parse import urlparse # pylint: disable=import-outside-toplevel + from urllib.request import ( # pylint: disable=import-outside-toplevel + getproxies, + proxy_bypass, + ) + + host = urlparse(endpoint).hostname or "" + if getproxies().get("https") and not proxy_bypass(host): + return + try: + socket.getaddrinfo(host, 443) + except socket.gaierror as exc: + raise NetworkVolumeTransferError( + f"{datacenter} does not support the s3 api yet ({S3_DOCS_URL}), " + f"or {host} is unreachable from this network" + ) from exc + + +def _client_config(concurrency: int) -> Any: + from botocore.config import Config # pylint: disable=import-outside-toplevel + + return Config( + s3={"addressing_style": "path"}, + request_checksum_calculation="when_required", + response_checksum_validation="when_required", + # one try: _Transfer._retry_delay makes every retry decision + retries={"mode": "standard", "total_max_attempts": 1}, + connect_timeout=60, + # a socket read silent this long is a stalled connection (runpod's + # proxy answers a slow request with 524 after about 100s) + read_timeout=300, + # headroom over the request bound for connections being returned + max_pool_connections=max(10, 2 * concurrency), + ) + + def _make_transfer( volume: str, access_key: Optional[str], @@ -889,47 +1080,21 @@ def _make_transfer( """a transfer to the network volume with this id or name. backs NetworkVolume.upload and NetworkVolume.download.""" import boto3 # pylint: disable=import-outside-toplevel - from botocore.config import Config # pylint: disable=import-outside-toplevel - access_key = access_key or os.environ.get(ACCESS_KEY_ENV) - secret_key = secret_key or os.environ.get(SECRET_KEY_ENV) - if not access_key or not secret_key: - raise NetworkVolumeTransferError( - f"s3 api keys required: set {ACCESS_KEY_ENV} and {SECRET_KEY_ENV} " - f"(create keys at {S3_KEYS_URL})" - ) volume_id, datacenter = _resolve_volume(volume) + access_key, secret_key = _s3_keys(access_key, secret_key) endpoint = endpoint_url or f"https://s3api-{datacenter.lower()}.runpod.io/" + if not endpoint_url: + _check_endpoint(endpoint, datacenter) concurrency = max(1, int(concurrency)) - config = Config( - s3={"addressing_style": "path"}, - request_checksum_calculation="when_required", - response_checksum_validation="when_required", - # total_max_attempts counts the first try (max_attempts would not) - retries={"mode": "standard", "total_max_attempts": MAX_ATTEMPTS}, - connect_timeout=60, - # a socket read silent this long is a stalled connection (runpod's - # proxy answers a slow request with 524 after about 100s) - read_timeout=300, - # headroom over the request bound for connections being returned - max_pool_connections=max(10, 2 * concurrency), - ) client = boto3.session.Session().client( "s3", region_name=datacenter, endpoint_url=endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, - config=config, + config=_client_config(concurrency), ) - client.meta.events.register_first("needs-retry.s3", _retry_proxy_timeout) - return _Transfer(client, volume_id, endpoint, concurrency, progress, force=force) - - -def _retry_proxy_timeout(response=None, attempts=1, **_kwargs) -> Optional[float]: - """the backoff before retrying a 524, which runpod's proxy returns when an - operation outlives it. botocore's standard mode retries 500/502/503/504 - only; returning None leaves every other response to it.""" - if response is None or response[0].status_code != 524 or attempts >= MAX_ATTEMPTS: - return None - return random.uniform(0, min(20.0, 2.0**attempts)) + transfer = _Transfer(client, volume_id, endpoint, concurrency, progress, force=force) + client.meta.events.register_first("needs-retry.s3", transfer._retry_delay) + return transfer diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index 0b1d1e20f..66413fcb2 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -8,9 +8,9 @@ def volume(): """Upload and download files on network volumes (resumable). - Uses the S3-compatible API. Set RUNPOD_S3_ACCESS_KEY and - RUNPOD_S3_SECRET_KEY (create keys at console.runpod.io/user/settings). - Re-run an interrupted command to resume it. Parts of an upload that is + Uses the S3-compatible API, which needs S3 API keys (separate from your + API key): set RUNPOD_S3_ACCESS_KEY and RUNPOD_S3_SECRET_KEY, pass the + --s3-* flags, or enter them when asked. Re-run an interrupted command to resume it. Parts of an upload that is never re-run stay in the volume's hidden .s3compat_uploads folder and count toward its size. """ @@ -113,39 +113,141 @@ def _run_transfer(action_label, transfer_method, source, target, opts) -> None: ) -def _volume(reference): +def _interactive() -> bool: + return sys.stdin.isatty() and sys.stdout.isatty() + + +def _ensure_logged_in(ctx) -> None: + """offer the browser login when no api key is set; the volume lookup + needs one.""" + import runpod + from runpod import network_volume_transfer as nvt + from runpod.cli.groups.config.functions import get_credentials + from runpod.rp_cli.main import login + + if runpod.api_key: + return + if not _interactive() or not click.confirm( + "you are not logged in to runpod. log in now?", default=True + ): + raise click.ClickException(nvt.NOT_LOGGED_IN) + ctx.invoke(login, no_open=False, api_key_opt=None, profile=runpod.profile) + runpod.api_key = (get_credentials(runpod.profile) or {}).get("api_key") + if not runpod.api_key: + raise click.ClickException(nvt.NOT_LOGGED_IN) + + +def _pick_volume() -> str: + """let the user choose one of their network volumes by number.""" + from runpod import network_volume_transfer as nvt + + if not _interactive(): + raise click.UsageError("pass VOLUME first (a network volume id or name)") + try: + volumes = nvt._list_volumes() + except nvt.NetworkVolumeTransferError as exc: + raise click.ClickException(str(exc)) from exc + if not volumes: + raise click.ClickException( + "no network volumes yet: create one at https://console.runpod.io/user/storage" + ) + for number, vol in enumerate(volumes, 1): + datacenter = vol.get("dataCenterId") or vol.get("dataCenter") + click.echo( + f" {number}. {vol['name']} {vol['id']} {vol.get('size')} GB {datacenter}" + ) + choice = click.prompt("volume", type=click.IntRange(1, len(volumes))) + return volumes[choice - 1]["id"] + + +def _missing_keys_message() -> str: + from runpod import network_volume_transfer as nvt + + return ( + "s3 api keys required (separate from your api key): set " + f"{nvt.ACCESS_KEY_ENV} and {nvt.SECRET_KEY_ENV} or pass --s3-access-key " + f"and --s3-secret-key (create keys at {nvt.S3_KEYS_URL})" + ) + + +def _prompt_s3_keys(): + """ask for the s3 keys, offering to save them for next time.""" + import runpod + from runpod import network_volume_transfer as nvt + from runpod.cli.groups.config.functions import set_s3_keys + + if not _interactive(): + raise click.ClickException(_missing_keys_message()) + click.echo( + "network volume transfers use s3 api keys, separate from your api key.\n" + f"create a pair at {nvt.S3_KEYS_URL} (s3 api keys)." + ) + access = click.prompt("s3 access key") + secret = click.prompt("s3 secret key", hide_input=True) + if click.confirm("save them to ~/.runpod/config.toml?", default=True): + set_s3_keys(access, secret, profile=runpod.profile) + return access, secret + + +def _ensure_s3_keys(opts) -> None: + """fill in the s3 keys from the flags, the environment, or the saved + config, prompting for them when none are set.""" + from runpod import network_volume_transfer as nvt + + keys = nvt.find_s3_keys(opts["s3_access_key"], opts["s3_secret_key"]) + opts["s3_access_key"], opts["s3_secret_key"] = keys or _prompt_s3_keys() + + +def _split_args(args, second: str, third: str): + """[VOLUME] SECOND [THIRD]: a lone argument is SECOND, and the volume is + picked interactively.""" + if not 1 <= len(args) <= 3: + raise click.UsageError(f"expected [VOLUME] {second} [{third}]") + if len(args) == 1: + return None, args[0], None + return args[0], args[1], args[2] if len(args) == 3 else None + + +def _transfer(ctx, action_label, method_name, reference, source, target, opts) -> None: + """the checks every transfer needs, in order: login, volume, s3 keys.""" from runpod.apps import NetworkVolume - return NetworkVolume(reference, create=False) + _ensure_logged_in(ctx) + reference = reference or _pick_volume() + _ensure_s3_keys(opts) + method = getattr(NetworkVolume(reference, create=False), method_name) + _run_transfer(action_label, method, source, target, opts) @volume.command(name="upload") -@click.argument("volume_ref", metavar="VOLUME") -@click.argument("local_path", type=click.Path(exists=True)) -@click.argument("remote_path", required=False) +@click.argument("args", nargs=-1, required=True, metavar="[VOLUME] LOCAL_PATH [REMOTE_PATH]") @_transfer_options -def volume_upload(volume_ref, local_path, remote_path, **opts): +@click.pass_context +def volume_upload(ctx, args, **opts): """Upload a file or folder to a network volume. - VOLUME is the volume's id or name. REMOTE_PATH defaults to the local - name. A folder's contents go under REMOTE_PATH/. Files whose remote copy - has the same size and is not older are skipped; --force sends them - anyway. + VOLUME is the volume's id or name; leave it out to pick one. REMOTE_PATH + defaults to the local name. A folder's contents go under REMOTE_PATH/. + Files whose remote copy has the same size and is not older are skipped; + --force sends them anyway. """ - _run_transfer("uploaded", _volume(volume_ref).upload, local_path, remote_path, opts) + reference, local_path, remote_path = _split_args(args, "LOCAL_PATH", "REMOTE_PATH") + if not Path(local_path).exists(): + raise click.UsageError(f"local path {local_path} does not exist") + _transfer(ctx, "uploaded", "upload", reference, local_path, remote_path, opts) @volume.command(name="download") -@click.argument("volume_ref", metavar="VOLUME") -@click.argument("remote_path") -@click.argument("local_path", required=False) +@click.argument("args", nargs=-1, required=True, metavar="[VOLUME] REMOTE_PATH [LOCAL_PATH]") @_transfer_options -def volume_download(volume_ref, remote_path, local_path, **opts): +@click.pass_context +def volume_download(ctx, args, **opts): """Download a file or folder from a network volume. - VOLUME is the volume's id or name. A REMOTE_PATH ending in / is a - folder. LOCAL_PATH defaults to the remote name. Local files of the same - size and not older than the remote copy are skipped; --force fetches - them anyway. + VOLUME is the volume's id or name; leave it out to pick one. A + REMOTE_PATH ending in / is a folder. LOCAL_PATH defaults to the remote + name. Local files of the same size and not older than the remote copy + are skipped; --force fetches them anyway. """ - _run_transfer("downloaded", _volume(volume_ref).download, remote_path, local_path, opts) + reference, remote_path, local_path = _split_args(args, "REMOTE_PATH", "LOCAL_PATH") + _transfer(ctx, "downloaded", "download", reference, remote_path, local_path, opts) diff --git a/tests/e2e/test_network_volume_transfer_live.py b/tests/e2e/test_network_volume_transfer_live.py new file mode 100644 index 000000000..ae4c6f56e --- /dev/null +++ b/tests/e2e/test_network_volume_transfer_live.py @@ -0,0 +1,74 @@ +"""live round trip through a network volume's s3 api. + +creates a 10 GB volume (billed while it exists) and deletes it afterwards, so +it only runs when asked for: set RUNPOD_E2E_VOLUMES=1 together with +RUNPOD_API_KEY, RUNPOD_S3_ACCESS_KEY and RUNPOD_S3_SECRET_KEY. +""" + +import hashlib +import os +import uuid + +import pytest + +import runpod +from runpod.api.rest import run_rest_request + +DATACENTER = os.environ.get("RUNPOD_E2E_VOLUME_DATACENTER", "EU-RO-1") +REQUIRED = ("RUNPOD_API_KEY", "RUNPOD_S3_ACCESS_KEY", "RUNPOD_S3_SECRET_KEY") + +pytestmark = pytest.mark.skipif( + os.environ.get("RUNPOD_E2E_VOLUMES") != "1" + or not all(os.environ.get(name) for name in REQUIRED), + reason="set RUNPOD_E2E_VOLUMES=1 and " + ", ".join(REQUIRED), +) + + +@pytest.fixture +def volume(): + created = run_rest_request( + "POST", + "/v2/network-volumes", + json={ + "name": f"e2e-transfer-{uuid.uuid4().hex[:8]}", + "size": 10, + "dataCenter": DATACENTER, + }, + ) + try: + yield runpod.NetworkVolume(created["id"], create=False) + finally: + run_rest_request("DELETE", f"/v2/network-volumes/{created['id']}") + + +def _sha256(path): + digest = hashlib.sha256() + with open(path, "rb") as handle: + for block in iter(lambda: handle.read(1 << 20), b""): + digest.update(block) + return digest.hexdigest() + + +def test_round_trip_resumes_and_skips(volume, tmp_path): + big = tmp_path / "up" / "big.bin" + big.parent.mkdir() + # over the 64 MiB single-request limit, so the upload is multipart and the + # download is chunked + big.write_bytes(os.urandom(70 * 1024 * 1024)) + folder = tmp_path / "up" / "folder" + (folder / "sub").mkdir(parents=True) + (folder / "a.txt").write_text("a") + (folder / "sub" / "b.txt").write_text("b") + + assert volume.upload(big, "e2e/big.bin")["files_transferred"] == 1 + assert volume.upload(folder, "e2e/folder")["files_transferred"] == 2 + + down = tmp_path / "down" + volume.download("e2e/big.bin", down / "big.bin") + volume.download("e2e/folder/", str(down / "folder")) + assert _sha256(down / "big.bin") == _sha256(big) + assert (down / "folder" / "sub" / "b.txt").read_text() == "b" + + # a rerun finds everything unchanged in both directions + assert volume.upload(folder, "e2e/folder")["files_skipped"] == 2 + assert volume.download("e2e/folder/", str(down / "folder"))["files_skipped"] == 2 diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 7102066d4..4cf58ed8c 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -149,6 +149,16 @@ def ops(self, op): return [kwargs for name, kwargs in self.calls if name == op] +@pytest.fixture(autouse=True) +def credentials(monkeypatch): + """a logged-in user with s3 keys in the environment and nothing saved; + tests that need otherwise change it.""" + monkeypatch.setattr(runpod, "api_key", "rpa_test") + monkeypatch.setattr(nvt, "saved_s3_keys", lambda: (None, None)) + monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_env") + monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_env") + + @pytest.fixture(autouse=True) def small_sizes(monkeypatch): monkeypatch.setattr(nvt, "SINGLE_PUT_MAX", 10) @@ -207,7 +217,7 @@ def test_multipart_resume_uploads_only_missing_and_bad_parts(tmp_path): src.write_bytes(DATA) fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) + _client_error("InvalidArgument", 400) if op == "upload_part" and kw["PartNumber"] == 5 else None ) @@ -236,7 +246,7 @@ def test_multipart_resume_reuploads_parts_without_md5_etag(tmp_path): src = tmp_path / "big.bin" src.write_bytes(DATA) fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) + _client_error("InvalidArgument", 400) if op == "upload_part" and kw["PartNumber"] == 5 else None ) @@ -255,7 +265,7 @@ def test_changed_file_aborts_stale_upload(tmp_path): src = tmp_path / "big.bin" src.write_bytes(DATA) fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) if op == "upload_part" and kw["PartNumber"] == 2 else None + _client_error("InvalidArgument", 400) if op == "upload_part" and kw["PartNumber"] == 2 else None ) with pytest.raises(ClientError): _transfer(fake, tmp_path).upload(src, "big.bin") @@ -274,7 +284,7 @@ def test_missing_server_upload_starts_fresh(tmp_path): src = tmp_path / "big.bin" src.write_bytes(DATA) fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) if op == "upload_part" and kw["PartNumber"] == 2 else None + _client_error("InvalidArgument", 400) if op == "upload_part" and kw["PartNumber"] == 2 else None ) with pytest.raises(ClientError): _transfer(fake, tmp_path).upload(src, "big.bin") @@ -310,7 +320,7 @@ def test_complete_failure_without_object_keeps_state(tmp_path): src = tmp_path / "big.bin" src.write_bytes(DATA) fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) if op == "complete_multipart_upload" else None + _client_error("InvalidArgument", 400) if op == "complete_multipart_upload" else None ) with pytest.raises(ClientError): _transfer(fake, tmp_path).upload(src, "big.bin") @@ -342,7 +352,7 @@ def test_folder_upload_default_prefix_and_failure_report(tmp_path): (root / "ok.txt").write_bytes(b"1") (root / "bad.txt").write_bytes(b"2") fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) if kw.get("Key") == "data/bad.txt" else None + _client_error("InvalidArgument", 400) if kw.get("Key") == "data/bad.txt" else None ) with pytest.raises(nvt.NetworkVolumeTransferError, match="1 file"): _transfer(fake, tmp_path).upload(root, None) @@ -357,7 +367,7 @@ def test_folder_failure_is_resumable_only_when_every_failure_is(tmp_path): (root / "denied.txt").write_bytes(b"2") errors = { "data/cut.txt": _client_error("ServiceUnavailable", 503), - "data/denied.txt": _client_error("AccessDenied", 403), + "data/denied.txt": _client_error("InvalidArgument", 400), } fake.fail = lambda op, kw: errors.get(kw.get("Key")) if op == "put_object" else None with pytest.raises(nvt.NetworkVolumeTransferError, match="2 file") as mixed: @@ -392,7 +402,7 @@ def test_download_resumes_from_sidecar(tmp_path): fake.put("big.bin", DATA) dest = tmp_path / "out" / "big.bin" fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) + _client_error("InvalidArgument", 400) if op == "get_object" and kw["Range"] == "bytes=32-39" else None ) @@ -418,7 +428,7 @@ def test_download_restarts_when_remote_changed(tmp_path): fake.put("big.bin", DATA) dest = tmp_path / "big.bin" fake.fail = lambda op, kw: ( - _client_error("AccessDenied", 403) if kw.get("Range") == "bytes=16-23" else None + _client_error("InvalidArgument", 400) if kw.get("Range") == "bytes=16-23" else None ) with pytest.raises(ClientError): _transfer(fake, tmp_path).download("big.bin", dest) @@ -590,7 +600,7 @@ def test_folder_failure_resumable_only_if_a_file_was(tmp_path): root = tmp_path / "data" root.mkdir() (root / "a.txt").write_bytes(b"1") - fake.fail = lambda op, kw: _client_error("AccessDenied", 403) if op == "put_object" else None + fake.fail = lambda op, kw: _client_error("InvalidArgument", 400) if op == "put_object" else None with pytest.raises(nvt.NetworkVolumeTransferError) as denied: _transfer(fake, tmp_path).upload(root, None) assert denied.value.resumable is False @@ -629,16 +639,27 @@ def test_make_transfer_requires_s3_keys(monkeypatch): monkeypatch.delenv(nvt.SECRET_KEY_ENV, raising=False) monkeypatch.setenv("AWS_ACCESS_KEY_ID", "aws") monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "aws") + monkeypatch.setattr(nvt, "_resolve_volume", lambda volume: ("vol", "EU-RO-1")) with pytest.raises(nvt.NetworkVolumeTransferError, match="RUNPOD_S3_ACCESS_KEY"): nvt._make_transfer("vol", None, None, None, 4, None) +def test_saved_s3_keys_are_the_last_fallback(monkeypatch): + monkeypatch.delenv(nvt.ACCESS_KEY_ENV) + monkeypatch.setattr(nvt, "saved_s3_keys", lambda: ("user_saved", "rps_saved")) + assert nvt._s3_keys(None, None) == ("user_saved", "rps_saved") + # a pair never mixes sources: a lone flag does not complete a pair + assert nvt._s3_keys("user_flag", None) == ("user_saved", "rps_saved") + assert nvt._s3_keys("user_flag", "rps_flag") == ("user_flag", "rps_flag") + + def test_make_transfer_builds_datacenter_client(monkeypatch): # a stale local aws profile must not leak into the session for name in ("AWS_PROFILE", "AWS_DEFAULT_PROFILE", "AWS_CONFIG_FILE", "AWS_SHARED_CREDENTIALS_FILE"): monkeypatch.delenv(name, raising=False) monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") + monkeypatch.setattr(nvt, "_check_endpoint", lambda endpoint, datacenter: None) listing, calls = _list_volumes(_VOLUME) with listing: transfer = nvt._make_transfer("models", None, None, None, 4, None) @@ -646,8 +667,8 @@ def test_make_transfer_builds_datacenter_client(monkeypatch): assert transfer.client.meta.endpoint_url == "https://s3api-eu-ro-1.runpod.io/" assert transfer.client.meta.region_name == "EU-RO-1" assert transfer.client.meta.config.request_checksum_calculation == "when_required" - # botocore owns retries - assert transfer.client.meta.config.retries == {"mode": "standard", "total_max_attempts": 10} + # _retry_delay makes every retry decision + assert transfer.client.meta.config.retries == {"mode": "standard", "total_max_attempts": 1} credentials = transfer.client._request_signer._credentials assert credentials.access_key == "user_x" assert transfer.bucket == "vol123" @@ -668,7 +689,7 @@ def test_resolve_volume_by_id(): @pytest.mark.parametrize( "volumes, match", [ - ([], "not found"), + ([], "no network volume named"), ([{**_VOLUME, "id": "a"}, {**_VOLUME, "id": "b"}], "reference by id"), ([{"id": "vol123", "name": "models"}], "datacenter"), ], @@ -989,7 +1010,7 @@ def put_then_edit(**kwargs): return result fake.put_object = put_then_edit - fake.fail = lambda op, kw: _client_error("AccessDenied", 403) if op == "delete_object" else None + fake.fail = lambda op, kw: _client_error("InvalidArgument", 400) if op == "delete_object" else None with pytest.raises(nvt.NetworkVolumeTransferError, match="--force") as info: _transfer(fake, tmp_path).upload(src, "small.txt") assert not nvt.is_resumable(info.value) @@ -1219,3 +1240,294 @@ def log_message(self, *args): finally: server.shutdown() assert len(hits) == 3 + + +# ------------------------------------------------- prompts and failure cases + + +def _cli(): + from runpod.rp_cli.main import cli + + return cli + + +@pytest.fixture +def interactive(monkeypatch): + monkeypatch.setattr("runpod.rp_cli.volume._interactive", lambda: True) + + +def _fake_transfer(calls): + """a NetworkVolume.upload/download stand-in that records its call.""" + + def method(volume, source, target, **kwargs): + calls.append((volume.reference, source, target, kwargs)) + return {"files_transferred": 1, "files_skipped": 0, "bytes_transferred": 1} + + return method + + +def test_cli_not_logged_in_says_how_to_log_in(monkeypatch, tmp_path): + monkeypatch.setattr(runpod, "api_key", None) + result = CliRunner().invoke(_cli(), ["volume", "upload", "vol", str(tmp_path)]) + assert result.exit_code == 1 + assert "not logged in: run `rp login` or set RUNPOD_API_KEY" in result.output + + +def test_cli_offers_login_and_continues(monkeypatch, tmp_path, interactive): + import click as _click + + logins = [] + monkeypatch.setattr(runpod, "api_key", None) + monkeypatch.setattr( + "runpod.rp_cli.main.login", + _click.Command("login", callback=lambda **kw: logins.append(kw)), + ) + monkeypatch.setattr( + "runpod.cli.groups.config.functions.get_credentials", + lambda profile="default": {"api_key": "rpa_new"}, + ) + calls = [] + with patch.object(runpod.NetworkVolume, "upload", autospec=True, side_effect=_fake_transfer(calls)): + result = CliRunner().invoke(_cli(), ["volume", "upload", "vol", str(tmp_path)], input="y\n") + assert result.exit_code == 0, result.output + assert "log in now?" in result.output + assert len(logins) == 1 and runpod.api_key == "rpa_new" and calls + + +def test_cli_declined_login_stops(monkeypatch, tmp_path, interactive): + monkeypatch.setattr(runpod, "api_key", None) + result = CliRunner().invoke(_cli(), ["volume", "upload", "vol", str(tmp_path)], input="n\n") + assert result.exit_code == 1 + assert "not logged in" in result.output + + +def test_cli_picks_a_volume_when_left_out(monkeypatch, tmp_path, interactive): + volumes = [ + {"id": "aaa", "name": "first", "size": 10, "dataCenterId": "EU-RO-1"}, + {"id": "bbb", "name": "second", "size": 20, "dataCenterId": "US-KS-2"}, + ] + monkeypatch.setattr(nvt, "_list_volumes", lambda: volumes) + calls = [] + with patch.object(runpod.NetworkVolume, "upload", autospec=True, side_effect=_fake_transfer(calls)): + result = CliRunner().invoke(_cli(), ["volume", "upload", str(tmp_path)], input="2\n") + assert result.exit_code == 0, result.output + assert "2. second bbb 20 GB US-KS-2" in result.output + assert calls[0][0] == "bbb" and calls[0][1] == str(tmp_path) and calls[0][2] is None + + +def test_cli_without_volume_in_a_script_asks_for_it(tmp_path): + result = CliRunner().invoke(_cli(), ["volume", "download", "outputs/"]) + assert result.exit_code == 2 + assert "pass VOLUME first" in result.output + + +def test_cli_missing_keys_in_a_script_names_every_option(monkeypatch, tmp_path): + monkeypatch.delenv(nvt.ACCESS_KEY_ENV) + result = CliRunner().invoke(_cli(), ["volume", "upload", "vol", str(tmp_path)]) + assert result.exit_code == 1 + for hint in ("separate from your api key", "RUNPOD_S3_ACCESS_KEY", "--s3-access-key", nvt.S3_KEYS_URL): + assert hint in result.output + + +@pytest.mark.parametrize("save", [True, False]) +def test_cli_prompts_for_keys_and_offers_to_save(monkeypatch, tmp_path, interactive, save): + monkeypatch.delenv(nvt.ACCESS_KEY_ENV) + saved = [] + monkeypatch.setattr( + "runpod.cli.groups.config.functions.set_s3_keys", + lambda access, secret, profile="default": saved.append((access, secret)), + ) + calls = [] + with patch.object(runpod.NetworkVolume, "upload", autospec=True, side_effect=_fake_transfer(calls)): + result = CliRunner().invoke( + _cli(), + ["volume", "upload", "vol", str(tmp_path)], + input=f"user_typed\nrps_typed\n{'y' if save else 'n'}\n", + ) + assert result.exit_code == 0, result.output + assert calls[0][3]["access_key"] == "user_typed" and calls[0][3]["secret_key"] == "rps_typed" + assert saved == ([("user_typed", "rps_typed")] if save else []) + assert "rps_typed" not in result.output + + +def test_set_s3_keys_keeps_the_api_key_and_login_keeps_the_s3_keys(monkeypatch, tmp_path): + from runpod.cli.groups.config import functions + + monkeypatch.setattr(functions, "CREDENTIAL_FILE", str(tmp_path / "config.toml")) + functions.set_credentials("rpa_one", overwrite=True) + functions.set_s3_keys("user_x", "rps_x") + functions.set_credentials("rpa_two", overwrite=True) + assert functions.get_credentials() == { + "api_key": "rpa_two", + "s3_access_key": "user_x", + "s3_secret_key": "rps_x", + } + assert oct(os.stat(tmp_path / "config.toml").st_mode & 0o777) == "0o600" + + +@pytest.mark.parametrize("code, status", [("403", 403), ("InvalidAccessKeyId", 403), ("SignatureDoesNotMatch", 403)]) +def test_rejected_keys_say_so_and_do_not_resume(tmp_path, code, status): + fake = FakeS3() + fake.fail = lambda op, kw: _client_error(code, status, "HeadObject") + src = tmp_path / "f.txt" + src.write_text("x") + with pytest.raises(nvt.NetworkVolumeTransferError, match="rejected the keys") as info: + _transfer(fake, tmp_path).upload(src, None) + assert not nvt.is_resumable(info.value) + + +def test_datacenter_without_the_s3_api_fails_fast(monkeypatch): + import socket as _socket + + for name in ("HTTPS_PROXY", "https_proxy", "NO_PROXY", "no_proxy"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr("urllib.request.getproxies", lambda: {}) + + def no_such_host(host, port, *args, **kwargs): + raise _socket.gaierror(8, "nodename nor servname provided") + + monkeypatch.setattr(_socket, "getaddrinfo", no_such_host) + with pytest.raises(nvt.NetworkVolumeTransferError, match="EU-SE-1 does not support the s3 api yet") as info: + nvt._check_endpoint("https://s3api-eu-se-1.runpod.io/", "EU-SE-1") + assert not nvt.is_resumable(info.value) + # behind a proxy the proxy resolves the name + monkeypatch.setattr("urllib.request.getproxies", lambda: {"https": "http://proxy:3128"}) + monkeypatch.setattr("urllib.request.proxy_bypass", lambda host: False) + nvt._check_endpoint("https://s3api-eu-se-1.runpod.io/", "EU-SE-1") + + +def test_explicit_endpoint_skips_the_datacenter_check(monkeypatch): + checked = [] + monkeypatch.setattr(nvt, "_resolve_volume", lambda volume: ("vol123", "EU-SE-1")) + monkeypatch.setattr(nvt, "_check_endpoint", lambda endpoint, dc: checked.append(dc)) + nvt._make_transfer("vol123", None, None, "http://localhost:9000", 1, None) + assert checked == [] + nvt._make_transfer("vol123", None, None, None, 1, None) + assert checked == ["EU-SE-1"] + + +def _response(status, code=""): + parsed = {"ResponseMetadata": {"HTTPStatusCode": status}} + if code: + parsed["Error"] = {"Code": code} + return (type("R", (), {"status_code": status})(), parsed) + + +def test_retry_delay_decides_every_retry(monkeypatch, tmp_path): + from botocore.exceptions import EndpointConnectionError + + monkeypatch.setattr(nvt.random, "uniform", lambda a, b: 0) + transfer = _transfer(FakeS3(), tmp_path) + assert transfer._retry_delay(response=_response(524, "InternalError")) == 0 + assert transfer._retry_delay(response=_response(503)) == 0 + assert transfer._retry_delay(response=_response(429, "SlowDown")) == 0 + assert transfer._retry_delay(caught_exception=EndpointConnectionError(endpoint_url="x")) == 0 + assert transfer._retry_delay(response=_response(200)) is None + assert transfer._retry_delay(response=_response(404, "NoSuchUpload")) is None + assert transfer._retry_delay(response=_response(403, "AccessDenied")) is None + assert transfer._retry_delay(response=_response(503), attempts=nvt.MAX_ATTEMPTS) is None + + +def test_ctrl_c_ends_a_retry_wait(tmp_path): + transfer = _transfer(FakeS3(), tmp_path) + transfer._stop.set() + with pytest.raises(nvt.NetworkVolumeTransferError, match="interrupted") as info: + transfer._retry_delay(response=_response(503)) + assert info.value.resumable + # a successful response is never turned into an error + assert transfer._retry_delay(response=_response(200)) is None + + +def test_body_reads_hold_a_request_slot(tmp_path): + fake = FakeS3() + fake.put("big.bin", DATA) + transfer = _transfer(fake, tmp_path, concurrency=1) + held = [] + original_get = fake.get_object + + def get_object(**kwargs): + response = original_get(**kwargs) + body = response["Body"] + + class Body: + def read(self): + held.append(transfer._slots._value == 0) + return body.read() + + return {**response, "Body": Body()} + + fake.get_object = get_object + transfer.download("big.bin", tmp_path / "out.bin") + assert held and all(held) + + +def test_part_reads_hold_a_request_slot(tmp_path, monkeypatch): + src = tmp_path / "big.bin" + src.write_bytes(DATA) + transfer = _transfer(FakeS3(), tmp_path, concurrency=1) + held = [] + real_open = open + + def watch_open(path, mode="r", *args, **kwargs): + if str(path) == str(src) and "b" in mode: + held.append(transfer._slots._value == 0) + return real_open(path, mode, *args, **kwargs) + + monkeypatch.setattr("builtins.open", watch_open) + transfer.upload(src, None) + assert held and all(held) + + +def test_cleanup_takes_a_slot_even_after_ctrl_c(tmp_path): + fake = FakeS3() + transfer = _transfer(fake, tmp_path, concurrency=1) + held = [] + original_abort = fake.abort_multipart_upload + fake.abort_multipart_upload = lambda **kw: (held.append(transfer._slots._value == 0), original_abort(**kw)) + transfer._stop.set() + transfer._abort("k", "u") + assert held == [True] + + +def test_edit_during_complete_removes_the_object(tmp_path): + fake = FakeS3() + src = tmp_path / "big.bin" + src.write_bytes(DATA) + + def edit_during_complete(op, kw): + if op == "complete_multipart_upload": + src.write_bytes(DATA[::-1]) + later = src.stat().st_mtime + 60 + os.utime(src, (later, later)) + return None + + fake.fail = edit_during_complete + with pytest.raises(nvt.NetworkVolumeTransferError, match="changed during upload") as info: + _transfer(fake, tmp_path).upload(src, "big.bin") + assert info.value.resumable + assert "big.bin" not in fake.objects + + +def test_download_refuses_a_symlinked_subfolder(tmp_path): + fake = FakeS3() + fake.put("d/alias/x", b"1") + fake.put("d/real/x", b"2") + out = tmp_path / "out" + (out / "real").mkdir(parents=True) + try: + (out / "alias").symlink_to("real") + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + with pytest.raises(nvt.NetworkVolumeTransferError, match="is a symlink"): + _transfer(fake, tmp_path).download("d/", out) + assert not (out / "real" / "x").exists() + + +def test_state_writes_sync_the_folder(tmp_path, monkeypatch): + real_sync_dir = nvt._sync_dir + synced = [] + monkeypatch.setattr(nvt, "_sync_dir", synced.append) + nvt._write_json(tmp_path / "s" / "state.json", {"a": 1}) + assert synced == [tmp_path / "s"] + real_sync_dir(tmp_path / "s") + real_sync_dir(tmp_path / "missing") # best effort: no error From 675087f3598a48cd74e4b1cae789f900d0c9adec Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 27 Oct 2026 14:36:05 -0400 Subject: [PATCH 19/23] fix: read saved s3 keys only as a fallback, and say how the picker works the help now says the volume list appears when only a path is given, and the live e2e test is named for what it checks. --- runpod/network_volume_transfer.py | 13 +++++++------ runpod/rp_cli/volume.py | 8 ++++---- tests/e2e/test_network_volume_transfer_live.py | 2 +- tests/test_network_volume_transfer.py | 11 +++++++++++ 4 files changed, 23 insertions(+), 11 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index fbd811510..c6b82be79 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -1010,12 +1010,13 @@ def find_s3_keys( """the first complete pair of s3 keys from the arguments, the environment, or the saved config; None when there is none. a pair never mixes sources.""" - pairs = ( - (access_key, secret_key), - (os.environ.get(ACCESS_KEY_ENV), os.environ.get(SECRET_KEY_ENV)), - saved_s3_keys(), - ) - return next(((a, s) for a, s in pairs if a and s), None) + if access_key and secret_key: + return access_key, secret_key + from_env = (os.environ.get(ACCESS_KEY_ENV), os.environ.get(SECRET_KEY_ENV)) + if all(from_env): + return from_env + saved = saved_s3_keys() + return saved if all(saved) else None def _s3_keys(access_key: Optional[str], secret_key: Optional[str]) -> Tuple[str, str]: diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index 66413fcb2..531d68375 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -226,8 +226,8 @@ def _transfer(ctx, action_label, method_name, reference, source, target, opts) - def volume_upload(ctx, args, **opts): """Upload a file or folder to a network volume. - VOLUME is the volume's id or name; leave it out to pick one. REMOTE_PATH - defaults to the local name. A folder's contents go under REMOTE_PATH/. + VOLUME is the volume's id or name. Give only LOCAL_PATH to pick the + volume from a list. REMOTE_PATH defaults to the local name. A folder's contents go under REMOTE_PATH/. Files whose remote copy has the same size and is not older are skipped; --force sends them anyway. """ @@ -244,8 +244,8 @@ def volume_upload(ctx, args, **opts): def volume_download(ctx, args, **opts): """Download a file or folder from a network volume. - VOLUME is the volume's id or name; leave it out to pick one. A - REMOTE_PATH ending in / is a folder. LOCAL_PATH defaults to the remote + VOLUME is the volume's id or name. Give only REMOTE_PATH to pick the + volume from a list. A REMOTE_PATH ending in / is a folder. LOCAL_PATH defaults to the remote name. Local files of the same size and not older than the remote copy are skipped; --force fetches them anyway. """ diff --git a/tests/e2e/test_network_volume_transfer_live.py b/tests/e2e/test_network_volume_transfer_live.py index ae4c6f56e..c90bdbb87 100644 --- a/tests/e2e/test_network_volume_transfer_live.py +++ b/tests/e2e/test_network_volume_transfer_live.py @@ -49,7 +49,7 @@ def _sha256(path): return digest.hexdigest() -def test_round_trip_resumes_and_skips(volume, tmp_path): +def test_round_trip_then_rerun_skips(volume, tmp_path): big = tmp_path / "up" / "big.bin" big.parent.mkdir() # over the 64 MiB single-request limit, so the upload is multipart and the diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 4cf58ed8c..6534ca65b 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -1531,3 +1531,14 @@ def test_state_writes_sync_the_folder(tmp_path, monkeypatch): assert synced == [tmp_path / "s"] real_sync_dir(tmp_path / "s") real_sync_dir(tmp_path / "missing") # best effort: no error + + +def test_s3_keys_from_env_skip_the_saved_config(monkeypatch): + monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") + monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") + + def unexpected(): + raise AssertionError("the saved config was read") + + monkeypatch.setattr(nvt, "saved_s3_keys", unexpected) + assert nvt.find_s3_keys() == ("user_x", "rps_x") From cd56170fc46f0b85ff2032b91378dcb63159c5c3 Mon Sep 17 00:00:00 2001 From: Justin Date: Tue, 27 Oct 2026 14:46:58 -0400 Subject: [PATCH 20/23] fix: stop an upload mid-part on ctrl-c, and tell a rejected api key apart a part body now raises once ctrl-c sets the stop flag, so the cli exits in a fraction of a second instead of finishing every 64 mib part in flight (23 s live). an api key the api rejects gets its own message instead of "not logged in". --- runpod/network_volume_transfer.py | 35 +++++++++++++++++++++--- tests/test_network_volume_transfer.py | 38 ++++++++++++++++++++++++--- 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index c6b82be79..4bc5bcd86 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -14,6 +14,7 @@ """ import hashlib +import io import json import math import os @@ -49,6 +50,10 @@ S3_KEYS_URL = "https://console.runpod.io/user/settings" S3_DOCS_URL = "https://docs.runpod.io/storage/s3-api" NOT_LOGGED_IN = "not logged in: run `rp login` or set RUNPOD_API_KEY" +API_KEY_REJECTED = ( + "the runpod api rejected the api key (401): run `rp login` again or check " + "RUNPOD_API_KEY" +) # error codes s3 gives for keys it does not accept _REJECTED_KEY_CODES = ("InvalidAccessKeyId", "SignatureDoesNotMatch", "AccessDenied") @@ -233,6 +238,19 @@ def _check_remote_path(remote_path: Optional[str]) -> None: } +class _StoppableBody(io.BytesIO): + """bytes read by the http client in blocks; a read after stop raises.""" + + def __init__(self, data: bytes, stop: threading.Event): + super().__init__(data) + self._stop = stop + + def read(self, size: Optional[int] = -1) -> bytes: + if self._stop.is_set(): + raise NetworkVolumeTransferError("interrupted", resumable=True) + return super().read(size) + + class _Transfer: """one bucket on one endpoint. _make_transfer builds this.""" @@ -280,6 +298,11 @@ def _request(self, work: Callable[[], T]) -> T: raise raise rejected from exc + def _body(self, data: bytes) -> "_StoppableBody": + """a request body that ends the upload once ctrl-c sets the stop flag, + instead of sending the rest of a part that can be 64 MiB or more.""" + return _StoppableBody(data, self._stop) + def _client_call(self, method: str, **kwargs: Any) -> Any: return getattr(self.client, method)(Bucket=self.bucket, **kwargs) @@ -483,7 +506,9 @@ def upload_file(self, local: Path, key: str) -> None: def _single_put(self, local: Path, key: str, stat: os.stat_result) -> None: self._request( - lambda: self._client_call("put_object", Key=key, Body=local.read_bytes()) + lambda: self._client_call( + "put_object", Key=key, Body=self._body(local.read_bytes()) + ) ) self._keep_if_unchanged(local, key, stat) self._emit("advance", str(local), stat.st_size) @@ -703,7 +728,7 @@ def _read_and_send(self, number: int) -> Dict[str, Any]: Key=self.key, UploadId=self.upload_id, PartNumber=number, - Body=body, + Body=self.t._body(body), ) def _complete(self) -> None: @@ -965,7 +990,11 @@ def _list_volumes() -> List[Dict[str, Any]]: run_rest_paginated_request_async("/v2/network-volumes", "networkVolumes") ) except error.AuthenticationError as exc: - raise NetworkVolumeTransferError(NOT_LOGGED_IN) from exc + import runpod # pylint: disable=import-outside-toplevel,cyclic-import + + if not runpod.api_key: + raise NetworkVolumeTransferError(NOT_LOGGED_IN) from exc + raise NetworkVolumeTransferError(API_KEY_REJECTED) from exc def _resolve_volume(reference: str) -> Tuple[str, str]: diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 6534ca65b..f2eb12898 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -21,6 +21,11 @@ def _client_error(code, status, op="Op"): ) +def _payload(body): + """the bytes of a request body, read the way botocore reads a file.""" + return body.read() if hasattr(body, "read") else bytes(body) + + class FakeS3: """the subset of the s3 api the transfer module uses.""" @@ -56,8 +61,9 @@ def object_etag(self, key): def put_object(self, Bucket, Key, Body): self._call("put_object", Key=Key) - self.put(Key, bytes(Body)) - return {"ETag": self._etag(bytes(Body))} + data = _payload(Body) + self.put(Key, data) + return {"ETag": self._etag(data)} def head_object(self, Bucket, Key): self._call("head_object", Key=Key) @@ -108,8 +114,9 @@ def _upload(self, upload_id): def upload_part(self, Bucket, Key, UploadId, PartNumber, Body): self._call("upload_part", PartNumber=PartNumber) - self._upload(UploadId)["parts"][PartNumber] = bytes(Body) - return {"ETag": self._etag(bytes(Body))} + data = _payload(Body) + self._upload(UploadId)["parts"][PartNumber] = data + return {"ETag": self._etag(data)} def list_parts(self, Bucket, Key, UploadId, PartNumberMarker=0): self._call("list_parts", UploadId=UploadId, PartNumberMarker=PartNumberMarker) @@ -1542,3 +1549,26 @@ def unexpected(): monkeypatch.setattr(nvt, "saved_s3_keys", unexpected) assert nvt.find_s3_keys() == ("user_x", "rps_x") + + +def test_upload_body_stops_mid_part(tmp_path): + transfer = _transfer(FakeS3(), tmp_path) + body = transfer._body(b"x" * 100) + assert body.read(10) == b"x" * 10 + transfer._stop.set() + with pytest.raises(nvt.NetworkVolumeTransferError, match="interrupted") as info: + body.read(10) + assert nvt.is_resumable(info.value) + + +@pytest.mark.parametrize( + "api_key, message", [(None, "not logged in"), ("rpa_bad", "rejected the api key")] +) +def test_api_key_missing_vs_rejected(monkeypatch, api_key, message): + async def unauthorized(*args, **kwargs): + raise runpod.error.AuthenticationError("Unauthorized request") + + monkeypatch.setattr(runpod, "api_key", api_key) + monkeypatch.setattr("runpod.api.rest.run_rest_paginated_request_async", unauthorized) + with pytest.raises(nvt.NetworkVolumeTransferError, match=message): + nvt._resolve_volume("vol") From 369264c348c0e7643f093d6a7482970ef89e8091 Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 28 Oct 2026 15:21:08 -0400 Subject: [PATCH 21/23] fix(volume): validate before key prompts and recover rejected saved keys --- docs/cli/references/command_line_interface.md | 5 +- runpod/apps/volume.py | 7 +- runpod/network_volume_transfer.py | 15 ++- runpod/rp_cli/volume.py | 47 ++++++--- tests/test_network_volume_transfer.py | 99 ++++++++++++++++++- 5 files changed, 151 insertions(+), 22 deletions(-) diff --git a/docs/cli/references/command_line_interface.md b/docs/cli/references/command_line_interface.md index 9ac4e857c..435713ea0 100644 --- a/docs/cli/references/command_line_interface.md +++ b/docs/cli/references/command_line_interface.md @@ -75,7 +75,10 @@ rp volume download VOLUME outputs/ ./outputs Transfers use the volume's S3-compatible API, so they need S3 API keys (separate from the API key): set `RUNPOD_S3_ACCESS_KEY` and -`RUNPOD_S3_SECRET_KEY`, or pass `--s3-access-key` / `--s3-secret-key`. The +`RUNPOD_S3_SECRET_KEY`, pass `--s3-access-key` / `--s3-secret-key`, or save +`s3_access_key` / `s3_secret_key` in `~/.runpod/config.toml`. The CLI prompts +for missing keys and offers to save them. You can omit `VOLUME` in an +interactive session to pick a volume from a list. The volume must be in a [datacenter with the S3 API](https://docs.runpod.io/storage/s3-api). `VOLUME` is the volume's ID or name, looked up with your API key. diff --git a/runpod/apps/volume.py b/runpod/apps/volume.py index 4635c9df6..65b8dd3a0 100644 --- a/runpod/apps/volume.py +++ b/runpod/apps/volume.py @@ -106,8 +106,11 @@ def upload( skipped unless `force`, and an interrupted multipart upload resumes on re-run. - s3 keys come from the arguments or RUNPOD_S3_ACCESS_KEY and - RUNPOD_S3_SECRET_KEY. `progress(event, path, nbytes)` receives + s3 keys come from the arguments, RUNPOD_S3_ACCESS_KEY and + RUNPOD_S3_SECRET_KEY, or the saved s3_access_key/s3_secret_key in + ~/.runpod/config.toml. the CLI prompts for missing keys and can pick + a volume from a list when VOLUME is omitted. + `progress(event, path, nbytes)` receives "start", "advance", "reuse", "skip", and "done" events. returns {files_transferred, files_skipped, bytes_transferred}. diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index 4bc5bcd86..df685b357 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -52,7 +52,7 @@ NOT_LOGGED_IN = "not logged in: run `rp login` or set RUNPOD_API_KEY" API_KEY_REJECTED = ( "the runpod api rejected the api key (401): run `rp login` again or check " - "RUNPOD_API_KEY" + "RUNPOD_API_KEY or ~/.runpod/config.toml" ) # error codes s3 gives for keys it does not accept _REJECTED_KEY_CODES = ("InvalidAccessKeyId", "SignatureDoesNotMatch", "AccessDenied") @@ -71,6 +71,10 @@ def __init__(self, message: str, resumable: bool = False): self.resumable = resumable +class RejectedS3Keys(NetworkVolumeTransferError): + """The S3 endpoint rejected the supplied credential pair.""" + + def part_size_for(size: int) -> int: """part size for a multipart upload of `size` bytes.""" part_size = DEFAULT_PART_SIZE @@ -109,10 +113,11 @@ def _keys_rejected(exc: BaseException) -> Optional[NetworkVolumeTransferError]: code, status = _error(exc) if status not in (401, 403) and code not in _REJECTED_KEY_CODES: return None - return NetworkVolumeTransferError( - f"the volume's s3 api rejected the keys ({status or code}): check that " - f"{ACCESS_KEY_ENV} and {SECRET_KEY_ENV} are current and belong to this " - f"account (create keys at {S3_KEYS_URL})" + return RejectedS3Keys( + f"the volume's s3 api rejected the keys ({status or code}): check " + f"--s3-access-key/--s3-secret-key, {ACCESS_KEY_ENV}/{SECRET_KEY_ENV}, " + "or s3_access_key/s3_secret_key in ~/.runpod/config.toml " + f"(create keys at {S3_KEYS_URL}): {exc}" ) diff --git a/runpod/rp_cli/volume.py b/runpod/rp_cli/volume.py index 531d68375..96811c14c 100644 --- a/runpod/rp_cli/volume.py +++ b/runpod/rp_cli/volume.py @@ -1,3 +1,4 @@ +import os import sys from pathlib import Path @@ -10,7 +11,8 @@ def volume(): Uses the S3-compatible API, which needs S3 API keys (separate from your API key): set RUNPOD_S3_ACCESS_KEY and RUNPOD_S3_SECRET_KEY, pass the - --s3-* flags, or enter them when asked. Re-run an interrupted command to resume it. Parts of an upload that is + --s3-* flags, saved config, or enter them when asked. Re-run an interrupted + command to resume it. Parts of an upload that is never re-run stay in the volume's hidden .s3compat_uploads folder and count toward its size. """ @@ -83,16 +85,26 @@ def _invoke_transfer(transfer_method, source, target, opts, progress): from runpod import network_volume_transfer as nvt try: - return transfer_method( - source, - target, - access_key=opts["s3_access_key"], - secret_key=opts["s3_secret_key"], - endpoint_url=opts["endpoint_url"], - concurrency=opts["concurrency"], - progress=progress, - force=opts["force"], - ) + for attempt in range(2): + try: + return transfer_method( + source, + target, + access_key=opts["s3_access_key"], + secret_key=opts["s3_secret_key"], + endpoint_url=opts["endpoint_url"], + concurrency=opts["concurrency"], + progress=progress, + force=opts["force"], + ) + except nvt.RejectedS3Keys: + if not opts.get("_saved_keys") or attempt or not _interactive(): + raise + if not click.confirm( + "saved s3 keys were rejected. enter a new pair?", default=True + ): + raise + opts["s3_access_key"], opts["s3_secret_key"] = _prompt_s3_keys() except KeyboardInterrupt: click.echo("interrupted; re-run the same command to resume", err=True) sys.exit(130) @@ -194,7 +206,10 @@ def _ensure_s3_keys(opts) -> None: config, prompting for them when none are set.""" from runpod import network_volume_transfer as nvt + explicit = bool(opts["s3_access_key"] and opts["s3_secret_key"]) + env = bool(os.environ.get(nvt.ACCESS_KEY_ENV) and os.environ.get(nvt.SECRET_KEY_ENV)) keys = nvt.find_s3_keys(opts["s3_access_key"], opts["s3_secret_key"]) + opts["_saved_keys"] = bool(keys and not explicit and not env) opts["s3_access_key"], opts["s3_secret_key"] = keys or _prompt_s3_keys() @@ -211,9 +226,16 @@ def _split_args(args, second: str, third: str): def _transfer(ctx, action_label, method_name, reference, source, target, opts) -> None: """the checks every transfer needs, in order: login, volume, s3 keys.""" from runpod.apps import NetworkVolume + from runpod import network_volume_transfer as nvt _ensure_logged_in(ctx) reference = reference or _pick_volume() + try: + _, datacenter = nvt._resolve_volume(reference) + if not opts["endpoint_url"]: + nvt._check_endpoint(f"https://s3api-{datacenter.lower()}.runpod.io/", datacenter) + except nvt.NetworkVolumeTransferError as exc: + raise click.ClickException(str(exc)) from exc _ensure_s3_keys(opts) method = getattr(NetworkVolume(reference, create=False), method_name) _run_transfer(action_label, method, source, target, opts) @@ -227,7 +249,8 @@ def volume_upload(ctx, args, **opts): """Upload a file or folder to a network volume. VOLUME is the volume's id or name. Give only LOCAL_PATH to pick the - volume from a list. REMOTE_PATH defaults to the local name. A folder's contents go under REMOTE_PATH/. + volume from a list. REMOTE_PATH defaults to the local name. A folder's + contents go under REMOTE_PATH/. Files whose remote copy has the same size and is not older are skipped; --force sends them anyway. """ diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index f2eb12898..6e765e307 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -8,6 +8,7 @@ from unittest.mock import patch import pytest +import click from botocore.exceptions import ClientError, ReadTimeoutError from click.testing import CliRunner @@ -742,6 +743,7 @@ def test_no_module_level_transfer_exports(): # ------------------------------------------------------------------- cli +@pytest.mark.usefixtures("cli_preflight") def test_cli_upload_prints_summary(tmp_path): from runpod.rp_cli.main import cli @@ -766,6 +768,7 @@ def fake_upload(volume, source, target, **kwargs): assert "uploaded 1 file(s), 1 bytes; skipped 0 unchanged" in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_download_failure_says_resume(): from runpod.rp_cli.main import cli @@ -776,6 +779,7 @@ def test_cli_download_failure_says_resume(): assert "re-run the same command to resume" in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_ctrl_c_exits_130_with_resume_hint(): from runpod.rp_cli.main import cli @@ -785,6 +789,7 @@ def test_cli_ctrl_c_exits_130_with_resume_hint(): assert "interrupted; re-run the same command to resume" in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_final_error_has_no_resume_hint(): from runpod.rp_cli.main import cli @@ -796,6 +801,7 @@ def test_cli_final_error_has_no_resume_hint(): assert "re-run" not in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_missing_path_has_no_resume_hint(): from runpod.rp_cli.main import cli @@ -1159,7 +1165,7 @@ def interrupt_first(**kwargs): if first: raise KeyboardInterrupt # every other part waits for the interrupt to land, so what runs - # after it is deterministic: only parts already in flight finish + # after it is deterministic: parts in flight stop at their next read transfer._stop.wait(5) return original(**kwargs) @@ -1235,7 +1241,8 @@ def log_message(self, *args): pass server = http.server.HTTPServer(("127.0.0.1", 0), Handler) - threading.Thread(target=server.serve_forever, daemon=True).start() + server_thread = threading.Thread(target=server.serve_forever, daemon=True) + server_thread.start() try: monkeypatch.setenv(nvt.ACCESS_KEY_ENV, "user_x") monkeypatch.setenv(nvt.SECRET_KEY_ENV, "rps_x") @@ -1246,6 +1253,8 @@ def log_message(self, *args): transfer.client.list_objects_v2(Bucket="vol123") finally: server.shutdown() + server.server_close() + server_thread.join() assert len(hits) == 3 @@ -1263,6 +1272,13 @@ def interactive(monkeypatch): monkeypatch.setattr("runpod.rp_cli.volume._interactive", lambda: True) +@pytest.fixture +def cli_preflight(monkeypatch): + """Keep CLI presentation tests independent of the account and DNS.""" + monkeypatch.setattr(nvt, "_resolve_volume", lambda reference: (reference, "EU-RO-1")) + monkeypatch.setattr(nvt, "_check_endpoint", lambda endpoint, datacenter: None) + + def _fake_transfer(calls): """a NetworkVolume.upload/download stand-in that records its call.""" @@ -1280,6 +1296,7 @@ def test_cli_not_logged_in_says_how_to_log_in(monkeypatch, tmp_path): assert "not logged in: run `rp login` or set RUNPOD_API_KEY" in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_offers_login_and_continues(monkeypatch, tmp_path, interactive): import click as _click @@ -1308,6 +1325,7 @@ def test_cli_declined_login_stops(monkeypatch, tmp_path, interactive): assert "not logged in" in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_picks_a_volume_when_left_out(monkeypatch, tmp_path, interactive): volumes = [ {"id": "aaa", "name": "first", "size": 10, "dataCenterId": "EU-RO-1"}, @@ -1328,6 +1346,7 @@ def test_cli_without_volume_in_a_script_asks_for_it(tmp_path): assert "pass VOLUME first" in result.output +@pytest.mark.usefixtures("cli_preflight") def test_cli_missing_keys_in_a_script_names_every_option(monkeypatch, tmp_path): monkeypatch.delenv(nvt.ACCESS_KEY_ENV) result = CliRunner().invoke(_cli(), ["volume", "upload", "vol", str(tmp_path)]) @@ -1337,6 +1356,7 @@ def test_cli_missing_keys_in_a_script_names_every_option(monkeypatch, tmp_path): @pytest.mark.parametrize("save", [True, False]) +@pytest.mark.usefixtures("cli_preflight") def test_cli_prompts_for_keys_and_offers_to_save(monkeypatch, tmp_path, interactive, save): monkeypatch.delenv(nvt.ACCESS_KEY_ENV) saved = [] @@ -1357,6 +1377,53 @@ def test_cli_prompts_for_keys_and_offers_to_save(monkeypatch, tmp_path, interact assert "rps_typed" not in result.output +def test_cli_checks_volume_before_asking_for_keys(monkeypatch, tmp_path, interactive): + def missing_volume(reference): + raise nvt.NetworkVolumeTransferError("no network volume named 'missing'") + + monkeypatch.setattr(nvt, "_resolve_volume", missing_volume) + monkeypatch.setattr("runpod.rp_cli.volume._prompt_s3_keys", lambda: pytest.fail("prompted for keys")) + result = CliRunner().invoke(_cli(), ["volume", "upload", "missing", str(tmp_path)]) + assert result.exit_code == 1 + assert "no network volume named" in result.output + + +def test_cli_checks_endpoint_before_asking_for_keys(monkeypatch, tmp_path, interactive): + monkeypatch.setattr(nvt, "_resolve_volume", lambda reference: (reference, "EU-SE-1")) + + def unsupported_endpoint(endpoint, datacenter): + raise nvt.NetworkVolumeTransferError("does not support the s3 api") + + monkeypatch.setattr(nvt, "_check_endpoint", unsupported_endpoint) + monkeypatch.setattr("runpod.rp_cli.volume._prompt_s3_keys", lambda: pytest.fail("prompted for keys")) + result = CliRunner().invoke(_cli(), ["volume", "upload", "vol", str(tmp_path)]) + assert result.exit_code == 1 + assert "does not support the s3 api" in result.output + + +def test_cli_offers_to_replace_rejected_saved_keys(monkeypatch, interactive): + from runpod.rp_cli import volume as volume_cli + + calls = [] + + def transfer(source, target, **kwargs): + calls.append(kwargs["access_key"]) + if len(calls) == 1: + raise nvt.RejectedS3Keys("rejected the keys") + return {"files_transferred": 1} + + monkeypatch.setattr(volume_cli, "_prompt_s3_keys", lambda: ("new", "secret")) + opts = {"s3_access_key": "saved", "s3_secret_key": "old", "_saved_keys": True, + "endpoint_url": None, "concurrency": 1, "force": False} + result = CliRunner().invoke( + click.Command("test", callback=lambda: volume_cli._invoke_transfer( + transfer, "src", "dst", opts, None + )), input="y\n" + ) + assert result.exit_code == 0, result.output + assert calls == ["saved", "new"] + + def test_set_s3_keys_keeps_the_api_key_and_login_keeps_the_s3_keys(monkeypatch, tmp_path): from runpod.cli.groups.config import functions @@ -1381,6 +1448,8 @@ def test_rejected_keys_say_so_and_do_not_resume(tmp_path, code, status): with pytest.raises(nvt.NetworkVolumeTransferError, match="rejected the keys") as info: _transfer(fake, tmp_path).upload(src, None) assert not nvt.is_resumable(info.value) + assert "--s3-access-key" in str(info.value) + assert "~/.runpod/config.toml" in str(info.value) def test_datacenter_without_the_s3_api_fails_fast(monkeypatch): @@ -1561,6 +1630,32 @@ def test_upload_body_stops_mid_part(tmp_path): assert nvt.is_resumable(info.value) +@pytest.mark.parametrize("multipart", [False, True]) +def test_upload_call_sites_stop_before_sending_body(tmp_path, monkeypatch, multipart): + fake = FakeS3() + transfer = _transfer(fake, tmp_path) + src = tmp_path / "f.bin" + src.write_bytes(b"x" * (20 if multipart else 1)) + if multipart: + monkeypatch.setattr(nvt, "SINGLE_PUT_MAX", 8) + original = fake.upload_part + def interrupted(**kwargs): + transfer._stop.set() + return original(**kwargs) + fake.upload_part = interrupted + else: + original = fake.put_object + def interrupted(**kwargs): + transfer._stop.set() + return original(**kwargs) + fake.put_object = interrupted + with pytest.raises(nvt.NetworkVolumeTransferError, match="interrupted") as info: + transfer.upload(src, "f.bin") + assert nvt.is_resumable(info.value) + assert "f.bin" not in fake.objects + assert not any(upload["parts"] for upload in fake.uploads.values()) + + @pytest.mark.parametrize( "api_key, message", [(None, "not logged in"), ("rpa_bad", "rejected the api key")] ) From 721ef957a49cb700e18c5a7feba2fc37e58dcc5e Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 28 Oct 2026 15:26:44 -0400 Subject: [PATCH 22/23] test(volume): exercise saved key replacement through CLI --- tests/test_network_volume_transfer.py | 28 +++++++++++++-------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/tests/test_network_volume_transfer.py b/tests/test_network_volume_transfer.py index 6e765e307..2f67df45c 100644 --- a/tests/test_network_volume_transfer.py +++ b/tests/test_network_volume_transfer.py @@ -8,7 +8,6 @@ from unittest.mock import patch import pytest -import click from botocore.exceptions import ClientError, ReadTimeoutError from click.testing import CliRunner @@ -1401,25 +1400,24 @@ def unsupported_endpoint(endpoint, datacenter): assert "does not support the s3 api" in result.output -def test_cli_offers_to_replace_rejected_saved_keys(monkeypatch, interactive): - from runpod.rp_cli import volume as volume_cli - +@pytest.mark.usefixtures("cli_preflight") +def test_cli_offers_to_replace_rejected_saved_keys(monkeypatch, interactive, tmp_path): calls = [] - def transfer(source, target, **kwargs): + def transfer(volume, source, target, **kwargs): calls.append(kwargs["access_key"]) if len(calls) == 1: raise nvt.RejectedS3Keys("rejected the keys") - return {"files_transferred": 1} - - monkeypatch.setattr(volume_cli, "_prompt_s3_keys", lambda: ("new", "secret")) - opts = {"s3_access_key": "saved", "s3_secret_key": "old", "_saved_keys": True, - "endpoint_url": None, "concurrency": 1, "force": False} - result = CliRunner().invoke( - click.Command("test", callback=lambda: volume_cli._invoke_transfer( - transfer, "src", "dst", opts, None - )), input="y\n" - ) + return {"files_transferred": 1, "files_skipped": 0, "bytes_transferred": 1} + + monkeypatch.delenv(nvt.ACCESS_KEY_ENV, raising=False) + monkeypatch.delenv(nvt.SECRET_KEY_ENV, raising=False) + monkeypatch.setattr(nvt, "saved_s3_keys", lambda: ("saved", "old")) + monkeypatch.setattr("runpod.rp_cli.volume._prompt_s3_keys", lambda: ("new", "secret")) + with patch.object(runpod.NetworkVolume, "upload", autospec=True, side_effect=transfer): + result = CliRunner().invoke( + _cli(), ["volume", "upload", "vol", str(tmp_path)], input="y\n" + ) assert result.exit_code == 0, result.output assert calls == ["saved", "new"] From 13069a85e83763069d2aae7a3b0d1fde35b81560 Mon Sep 17 00:00:00 2001 From: Justin Date: Wed, 28 Oct 2026 16:28:05 -0400 Subject: [PATCH 23/23] refactor(volume): extract transfer filesystem helpers --- runpod/_volume_transfer_utils.py | 160 ++++++++++++++++++++++++++++++ runpod/network_volume_transfer.py | 158 ++++------------------------- 2 files changed, 181 insertions(+), 137 deletions(-) create mode 100644 runpod/_volume_transfer_utils.py diff --git a/runpod/_volume_transfer_utils.py b/runpod/_volume_transfer_utils.py new file mode 100644 index 000000000..f146d8424 --- /dev/null +++ b/runpod/_volume_transfer_utils.py @@ -0,0 +1,160 @@ +"""Filesystem and checksum helpers for resumable transfers. + +This private module keeps disk state and path safety details out of the S3 +request orchestration. The public error type remains in network_volume_transfer. +""" + +import hashlib +import json +import os +import unicodedata +from pathlib import Path, PurePosixPath +from typing import Any, Callable, Dict, List, Optional, Tuple + +MIB = 1024 * 1024 + + +def _read_json(path: Path) -> Optional[Dict[str, Any]]: + try: + data = json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + return data if isinstance(data, dict) else None + + +def _write_json(path: Path, data: Dict[str, Any], sync_dir: Callable[[Path], None]) -> None: + """replace path atomically, synced, so a crash leaves the old state or + the new one.""" + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(path.name + ".tmp") + with open(tmp, "w", encoding="utf-8") as handle: + json.dump(data, handle) + handle.flush() + os.fsync(handle.fileno()) + os.replace(tmp, path) + sync_dir(path.parent) + + +def _sync_dir(path: Path) -> None: + """make a rename in path survive a power loss. best effort: windows + cannot open a directory.""" + try: + fd = os.open(path, os.O_RDONLY) + except OSError: + return + try: + os.fsync(fd) + except OSError: + pass + finally: + os.close(fd) + + +def _md5_hex(path: Path, offset: int, length: int) -> str: + digest = hashlib.md5() # noqa: S324 - etag comparison, not security + with open(path, "rb") as handle: + handle.seek(offset) + while length > 0 and (block := handle.read(min(8 * MIB, length))): + digest.update(block) + length -= len(block) + return digest.hexdigest() + + +def _etag(value: Any) -> str: + return str(value or "").strip('"') + + +def _is_md5(etag: str) -> bool: + return len(etag) == 32 and all(c in "0123456789abcdef" for c in etag.lower()) + + +def _multipart_etag(parts: List[Dict[str, Any]]) -> str: + """the etag s3 (and runpod, verified) gives an object assembled from + parts: md5 of the parts' binary md5s, then "-". "" when a part + etag is not an md5.""" + etags = [_etag(part["ETag"]) for part in parts] + if not all(_is_md5(etag) for etag in etags): + return "" + digest = hashlib.md5(b"".join(bytes.fromhex(etag) for etag in etags)) # noqa: S324 + return f"{digest.hexdigest()}-{len(parts)}" + + +def _same_version(a: os.stat_result, b: os.stat_result) -> bool: + return a.st_size == b.st_size and a.st_mtime_ns == b.st_mtime_ns + + +def _not_older(a: float, b: float) -> bool: + """a >= b in whole seconds; s3 dates carry no finer precision.""" + return int(a) >= int(b) + + +def _check_remote_path(remote_path: Optional[str]) -> None: + if remote_path and ".." in remote_path.replace("\\", "/").split("/"): + raise ValueError(f"remote path {remote_path!r} must not contain '..'") + + +def _check_local_names( + targets: List[Path], root: Path, suffixes: Tuple[str, ...], error_type: type[Exception] +) -> None: + """refuse targets that would share a path on disk: two that are one name + on a case-insensitive or normalizing filesystem ("Model.bin" and + "model.bin"), a file whose name is also a folder another file lives in + ("a" beside "A/b"), or a file whose partial, sidecar or sidecar temp name + is also a target or a folder one lives in ("x" beside "x.rpdownload/y"). + names are compared nfc-normalized (apfs ignores normalization) and case + folded, which can refuse a listing a case-sensitive filesystem would + take: the safe direction.""" + + def fold(path: PurePosixPath) -> str: + return unicodedata.normalize("NFC", str(path)).casefold() + + rels = [PurePosixPath(target.relative_to(root).as_posix()) for target in targets] + files: Dict[str, PurePosixPath] = {} + dirs: Dict[str, PurePosixPath] = {} + for rel in rels: + if fold(rel) in files: + raise error_type( + f"cannot download both {files[fold(rel)]} and {rel}: they would be one " + "file on a case-insensitive or normalizing filesystem" + ) + files[fold(rel)] = rel + dirs.update({fold(parent): rel for parent in rel.parents if str(parent) != "."}) + for rel in rels: + if fold(rel) in dirs: + raise error_type( + f"cannot download both {rel} and {dirs[fold(rel)]}: " + "the first would have to be a file and a folder" + ) + for suffix in suffixes: + twin = fold(rel.with_name(rel.name + suffix)) + other = files.get(twin) or dirs.get(twin) + if other is not None: + raise error_type( + f"cannot download both {rel} and {other}: " + f"{rel.name + suffix} is used for the first's partial download" + ) + + +def _safe_target(root: Path, rel: str, error_type: type[Exception]) -> Path: + """local path for a remote relative key, refusing anything outside root. + "." and empty segments are refused too: "a/./b" would alias "a/b".""" + segments = rel.split("/") + if ( + rel.startswith("/") + or "\\" in rel + or any(seg in ("..", ".", "") for seg in segments) + or ":" in segments[0] + ): + raise error_type(f"refusing unsafe remote path: {rel!r}") + target = root.joinpath(*segments) + if root.resolve() not in target.resolve().parents: + raise error_type(f"refusing unsafe remote path: {rel!r}") + # a symlinked folder would let two remote paths share one local file, + # which the name collision check cannot see + for depth in range(1, len(segments)): + if root.joinpath(*segments[:depth]).is_symlink(): + raise error_type( + f"{root.joinpath(*segments[:depth])} is a symlink; " + "download into a folder without symlinked subfolders" + ) + return target diff --git a/runpod/network_volume_transfer.py b/runpod/network_volume_transfer.py index df685b357..69e628108 100644 --- a/runpod/network_volume_transfer.py +++ b/runpod/network_volume_transfer.py @@ -11,22 +11,38 @@ json, so a re-run fetches only what is missing. folder transfers skip files that are already present and unchanged. each client retries failed requests through `_Transfer._retry_delay`. + +filesystem state, checksums, and destination path checks live in the private +`_volume_transfer_utils` module; this module coordinates requests and resume. """ import hashlib import io -import json import math import os import random import socket import threading import time -import unicodedata from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path, PurePosixPath from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar +from ._volume_transfer_utils import ( + _check_local_names as _check_download_names, + _check_remote_path, + _etag, + _is_md5, + _md5_hex, + _multipart_etag, + _not_older, + _read_json, + _safe_target as _download_target, + _same_version, + _sync_dir, + _write_json as _persist_json, +) + T = TypeVar("T") MIB = 1024 * 1024 @@ -156,83 +172,8 @@ def is_resumable(exc: BaseException) -> bool: return code == "NoSuchUpload" or status == 416 or _transient(code, status) -def _read_json(path: Path) -> Optional[Dict[str, Any]]: - try: - data = json.loads(path.read_text(encoding="utf-8")) - except (OSError, ValueError): - return None - return data if isinstance(data, dict) else None - - def _write_json(path: Path, data: Dict[str, Any]) -> None: - """replace path atomically, synced, so a crash leaves the old state or - the new one.""" - path.parent.mkdir(parents=True, exist_ok=True) - tmp = path.with_name(path.name + ".tmp") - with open(tmp, "w", encoding="utf-8") as handle: - json.dump(data, handle) - handle.flush() - os.fsync(handle.fileno()) - os.replace(tmp, path) - _sync_dir(path.parent) - - -def _sync_dir(path: Path) -> None: - """make a rename in path survive a power loss. best effort: windows - cannot open a directory.""" - try: - fd = os.open(path, os.O_RDONLY) - except OSError: - return - try: - os.fsync(fd) - except OSError: - pass - finally: - os.close(fd) - - -def _md5_hex(path: Path, offset: int, length: int) -> str: - digest = hashlib.md5() # noqa: S324 - etag comparison, not security - with open(path, "rb") as handle: - handle.seek(offset) - while length > 0 and (block := handle.read(min(8 * MIB, length))): - digest.update(block) - length -= len(block) - return digest.hexdigest() - - -def _etag(value: Any) -> str: - return str(value or "").strip('"') - - -def _is_md5(etag: str) -> bool: - return len(etag) == 32 and all(c in "0123456789abcdef" for c in etag.lower()) - - -def _multipart_etag(parts: List[Dict[str, Any]]) -> str: - """the etag s3 (and runpod, verified) gives an object assembled from - parts: md5 of the parts' binary md5s, then "-". "" when a part - etag is not an md5.""" - etags = [_etag(part["ETag"]) for part in parts] - if not all(_is_md5(etag) for etag in etags): - return "" - digest = hashlib.md5(b"".join(bytes.fromhex(etag) for etag in etags)) # noqa: S324 - return f"{digest.hexdigest()}-{len(parts)}" - - -def _same_version(a: os.stat_result, b: os.stat_result) -> bool: - return a.st_size == b.st_size and a.st_mtime_ns == b.st_mtime_ns - - -def _not_older(a: float, b: float) -> bool: - """a >= b in whole seconds; s3 dates carry no finer precision.""" - return int(a) >= int(b) - - -def _check_remote_path(remote_path: Optional[str]) -> None: - if remote_path and ".." in remote_path.replace("\\", "/").split("/"): - raise ValueError(f"remote path {remote_path!r} must not contain '..'") + _persist_json(path, data, _sync_dir) # the summary counter each progress event advances @@ -915,68 +856,11 @@ def _read_range(self, start: int, length: int) -> Optional[bytes]: def _check_local_names(targets: List[Path], root: Path) -> None: - """refuse targets that would share a path on disk: two that are one name - on a case-insensitive or normalizing filesystem ("Model.bin" and - "model.bin"), a file whose name is also a folder another file lives in - ("a" beside "A/b"), or a file whose partial, sidecar or sidecar temp name - is also a target or a folder one lives in ("x" beside "x.rpdownload/y"). - names are compared nfc-normalized (apfs ignores normalization) and case - folded, which can refuse a listing a case-sensitive filesystem would - take: the safe direction.""" - - def fold(path: PurePosixPath) -> str: - return unicodedata.normalize("NFC", str(path)).casefold() - - rels = [PurePosixPath(target.relative_to(root).as_posix()) for target in targets] - files: Dict[str, PurePosixPath] = {} - dirs: Dict[str, PurePosixPath] = {} - for rel in rels: - if fold(rel) in files: - raise NetworkVolumeTransferError( - f"cannot download both {files[fold(rel)]} and {rel}: they would be one " - "file on a case-insensitive or normalizing filesystem" - ) - files[fold(rel)] = rel - dirs.update({fold(parent): rel for parent in rel.parents if str(parent) != "."}) - for rel in rels: - if fold(rel) in dirs: - raise NetworkVolumeTransferError( - f"cannot download both {rel} and {dirs[fold(rel)]}: " - "the first would have to be a file and a folder" - ) - for suffix in TRANSFER_SUFFIXES: - twin = fold(rel.with_name(rel.name + suffix)) - other = files.get(twin) or dirs.get(twin) - if other is not None: - raise NetworkVolumeTransferError( - f"cannot download both {rel} and {other}: " - f"{rel.name + suffix} is used for the first's partial download" - ) + _check_download_names(targets, root, TRANSFER_SUFFIXES, NetworkVolumeTransferError) def _safe_target(root: Path, rel: str) -> Path: - """local path for a remote relative key, refusing anything outside root. - "." and empty segments are refused too: "a/./b" would alias "a/b".""" - segments = rel.split("/") - if ( - rel.startswith("/") - or "\\" in rel - or any(seg in ("..", ".", "") for seg in segments) - or ":" in segments[0] - ): - raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") - target = root.joinpath(*segments) - if root.resolve() not in target.resolve().parents: - raise NetworkVolumeTransferError(f"refusing unsafe remote path: {rel!r}") - # a symlinked folder would let two remote paths share one local file, - # which the name collision check cannot see - for depth in range(1, len(segments)): - if root.joinpath(*segments[:depth]).is_symlink(): - raise NetworkVolumeTransferError( - f"{root.joinpath(*segments[:depth])} is a symlink; " - "download into a folder without symlinked subfolders" - ) - return target + return _download_target(root, rel, NetworkVolumeTransferError) # ---------------------------------------------------------------- client