From 292f1ea6067108e0df35936f23c431e2b7966765 Mon Sep 17 00:00:00 2001 From: Zachery Aaron Shores-Chmielewski Date: Sun, 26 Jul 2026 13:05:38 +0400 Subject: [PATCH] Prepare remote sweep streaming and preflight fixes --- configs/full_airfrans_incumbent_70gb.toml | 14 +- configs/remote_full_70gb.toml | 16 +- .../deeponet_branch_trunk.toml | 6 +- .../film_fourier_inr.toml | 6 +- ...shgraphnet_or_point_transformer_local.toml | 6 +- .../nerf_cfd_multires.toml | 6 +- .../point_context_perceiver.toml | 6 +- .../remote_model_zoo_7gb/raster_fno_unet.toml | 6 +- .../siren_conditioned_inr.toml | 6 +- pyproject.toml | 2 +- src/airfrans_frontier/cli.py | 10 +- src/airfrans_frontier/raw/process.py | 57 +- src/airfrans_frontier/raw/public.py | 600 ++++++++- src/airfrans_frontier/remote/artifacts.py | 24 +- src/airfrans_frontier/remote/cleanup.py | 166 +++ src/airfrans_frontier/remote/cli.py | 166 ++- src/airfrans_frontier/remote/collection.py | 180 +++ src/airfrans_frontier/remote/launch_group.py | 344 +++++ src/airfrans_frontier/remote/selection.py | 120 ++ src/airfrans_frontier/remote/skypilot.py | 113 +- src/airfrans_frontier/remote/smoke.py | 34 +- src/airfrans_frontier/remote/vast.py | 105 +- src/airfrans_frontier/training/config.py | 36 +- src/airfrans_frontier/training/hf_upload.py | 198 ++- src/airfrans_frontier/training/loop.py | 750 ++++++++++- .../training/streaming_data.py | 1132 +++++++++++++++++ tests/test_hf_upload.py | 55 +- tests/test_preflight_polish.py | 279 ++++ tests/test_public_data.py | 96 +- tests/test_remote_run.py | 32 +- tests/test_remote_smoke.py | 51 + tests/test_streaming_data.py | 372 ++++++ tests/test_training_loop.py | 35 + uv.lock | 349 ++--- 34 files changed, 5012 insertions(+), 366 deletions(-) create mode 100644 src/airfrans_frontier/remote/cleanup.py create mode 100644 src/airfrans_frontier/remote/collection.py create mode 100644 src/airfrans_frontier/remote/launch_group.py create mode 100644 src/airfrans_frontier/remote/selection.py create mode 100644 src/airfrans_frontier/training/streaming_data.py create mode 100644 tests/test_preflight_polish.py create mode 100644 tests/test_remote_smoke.py create mode 100644 tests/test_streaming_data.py diff --git a/configs/full_airfrans_incumbent_70gb.toml b/configs/full_airfrans_incumbent_70gb.toml index 05e3a1e..0ae02c7 100644 --- a/configs/full_airfrans_incumbent_70gb.toml +++ b/configs/full_airfrans_incumbent_70gb.toml @@ -4,17 +4,25 @@ seed = 20260723 artifact_dir = "artifacts/current_run/training_runs" [data] -root = "data/processed/full" +root = "artifacts/data_cache/airfrans_streaming_processed/processed/full" train_cases = 900 val_cases = 50 test_cases = 50 points_per_case = 999999999 batch_size = 4096 -source = "huggingface" +source = "public_zip_streaming" +public_source_url = "https://data.isir.upmc.fr/extrality/NeurIPS_2022/OF_dataset.zip" hf_repo_id = "zacheryasc/airfrans-processed" hf_repo_type = "dataset" hf_path_prefix = "processed/full" -cache_dir = "artifacts/data_cache/airfrans_processed" +cache_dir = "artifacts/data_cache/airfrans_streaming_processed/processed/full" +streaming_scratch_dir = "artifacts/data_cache/airfrans_streaming_processed/raw_scratch" +streaming_cache_max_bytes = 68719476736 +streaming_cache_high_water_bytes = 51539607552 +streaming_cache_low_water_bytes = 34359738368 +streaming_queue_max_cases = 2 +streaming_upload_processed = true +streaming_upload_batch_size = 16 [model] type = "film_fourier_inr" diff --git a/configs/remote_full_70gb.toml b/configs/remote_full_70gb.toml index f027ab6..b805b38 100644 --- a/configs/remote_full_70gb.toml +++ b/configs/remote_full_70gb.toml @@ -6,7 +6,7 @@ max_attempts = 5 [provider] kind = "vastai" -disk_gb = 512 +disk_gb = 192 max_price_per_hour = 0.80 image = "vastai/base:0.0.2" @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -31,6 +31,7 @@ workdir = "." exclude = [ "/artifacts", "/data/raw", + "/data/processed", "/.venv", "/notebooks", "__pycache__", @@ -40,13 +41,12 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] validation_command = """ -uv run --no-dev airfrans-frontier prepare-public-hf --repo-id zacheryasc/airfrans-processed --path-in-repo processed/full --work-dir artifacts/public_airfrans --output-dir artifacts/data_cache/airfrans_processed/processed/full --min-cases 1000 -uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/full_airfrans_incumbent_70gb.toml'); assert c.data.source == 'huggingface'; assert c.data.train_cases == 900; assert c.data.val_cases == 50; assert c.data.test_cases == 50; print('data_source=' + c.data.source + ' repo=' + str(c.data.hf_repo_id) + ' split=' + str((c.data.train_cases, c.data.val_cases, c.data.test_cases)))" +uv run --no-dev python -c "from airfrans_frontier.training.config import load_training_config; c=load_training_config('configs/full_airfrans_incumbent_70gb.toml'); assert c.data.source == 'public_zip_streaming'; assert c.data.train_cases == 900; assert c.data.val_cases == 50; assert c.data.test_cases == 50; assert c.data.streaming_cache_high_water_bytes < c.data.streaming_cache_max_bytes; assert c.data.streaming_cache_low_water_bytes < c.data.streaming_cache_high_water_bytes; print('data_source=' + c.data.source + ' public_url=' + str(c.data.public_source_url) + ' split=' + str((c.data.train_cases, c.data.val_cases, c.data.test_cases)) + ' cache_high_water=' + str(c.data.streaming_cache_high_water_bytes))" """ [job] @@ -79,8 +79,12 @@ required = [ "artifact_manifest.json", "checksums.txt", "verification_report.json", + "streaming_events.jsonl", + "streaming_state.json", + "streaming_summary.json", + "processed_upload_manifest.json", ] [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml b/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml index 40abc67..5460202 100644 --- a/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml +++ b/configs/remote_model_zoo_7gb/deeponet_branch_trunk.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/film_fourier_inr.toml b/configs/remote_model_zoo_7gb/film_fourier_inr.toml index 3e76811..8bef928 100644 --- a/configs/remote_model_zoo_7gb/film_fourier_inr.toml +++ b/configs/remote_model_zoo_7gb/film_fourier_inr.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml b/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml index f54e7b1..a986c54 100644 --- a/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml +++ b/configs/remote_model_zoo_7gb/meshgraphnet_or_point_transformer_local.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml b/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml index 8118085..5537eaf 100644 --- a/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml +++ b/configs/remote_model_zoo_7gb/nerf_cfd_multires.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/point_context_perceiver.toml b/configs/remote_model_zoo_7gb/point_context_perceiver.toml index 26e3055..3c0328b 100644 --- a/configs/remote_model_zoo_7gb/point_context_perceiver.toml +++ b/configs/remote_model_zoo_7gb/point_context_perceiver.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/raster_fno_unet.toml b/configs/remote_model_zoo_7gb/raster_fno_unet.toml index af814ef..b640358 100644 --- a/configs/remote_model_zoo_7gb/raster_fno_unet.toml +++ b/configs/remote_model_zoo_7gb/raster_fno_unet.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml b/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml index 3d6a474..05392e4 100644 --- a/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml +++ b/configs/remote_model_zoo_7gb/siren_conditioned_inr.toml @@ -21,7 +21,7 @@ min_down_mbps = 100 min_up_mbps = 25 require_verified = true blocked_geos = ["CN"] -blacklist_hosts = [59017] +blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031] drop_cheap_frac = 0.30 image_size_gb = 5.0 base_url = "https://cloud.vast.ai" @@ -41,7 +41,7 @@ exclude = [ [bootstrap] command = """ uv sync --no-dev -uv run --no-dev python -c "import torch; print('torch_cuda_available=' + str(torch.cuda.is_available())); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(torch.cuda.device_count()))" +uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); count=torch.cuda.device_count(); print('torch_cuda_available=' + str(ok)); print('torch_cuda_version=' + str(torch.version.cuda)); print('torch_device_count=' + str(count)); print('torch_device_name=' + (torch.cuda.get_device_name(0) if ok and count else 'none')); assert ok, 'torch CUDA unavailable'" """ [data] @@ -86,4 +86,4 @@ required = [ [cleanup] on_success = "sky_down" -on_failure = "collect_then_keep" +on_failure = "sky_down" diff --git a/pyproject.toml b/pyproject.toml index f40b930..8227deb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ requires-python = ">=3.11" dependencies = [ "huggingface-hub>=0.36.0", "numpy>=2.4.0", - "torch>=2.8.0", + "torch>=2.7.1,<2.8.0", "wandb>=0.23.0", ] diff --git a/src/airfrans_frontier/cli.py b/src/airfrans_frontier/cli.py index 84c9bb4..cab11a9 100644 --- a/src/airfrans_frontier/cli.py +++ b/src/airfrans_frontier/cli.py @@ -65,6 +65,10 @@ def build_parser() -> argparse.ArgumentParser: def main(argv: list[str] | None = None) -> int: parser = build_parser() args = parser.parse_args(argv) + if args.command in {"process-raw", "prepare-public-hf", "train", "model-sanity"}: + from airfrans_frontier.runtime import remove_pythonpath_entries + + remove_pythonpath_entries() if args.command == "inspect-raw": if args.sample_limit < 0: @@ -147,9 +151,6 @@ def main(argv: list[str] | None = None) -> int: return 0 if args.command == "train": - from airfrans_frontier.runtime import remove_pythonpath_entries - - remove_pythonpath_entries() from airfrans_frontier.training.loop import train_from_config_path try: @@ -166,9 +167,6 @@ def main(argv: list[str] | None = None) -> int: if args.steps <= 0: print("error: --steps must be positive", file=sys.stderr) return 1 - from airfrans_frontier.runtime import remove_pythonpath_entries - - remove_pythonpath_entries() from airfrans_frontier.training.sanity import MODEL_FAMILIES, run_model_sanity families = tuple(args.families) if args.families else MODEL_FAMILIES diff --git a/src/airfrans_frontier/raw/process.py b/src/airfrans_frontier/raw/process.py index bcce47a..f5eb637 100644 --- a/src/airfrans_frontier/raw/process.py +++ b/src/airfrans_frontier/raw/process.py @@ -92,30 +92,49 @@ def process_raw_dataset( started = time.perf_counter() total_points = 0 for index, case_dir in enumerate(case_dirs, start=1): - target_path = out_root / f"{case_dir.name}.npz" - if target_path.exists() and not force: - with np.load(target_path, allow_pickle=False) as npz: - points = int(npz["features"].shape[0]) - records.append({"case_id": case_dir.name, "path": str(target_path), "points": points, "skipped_existing": True}) - total_points += points - continue - metadata, features, targets = process_raw_case(case_dir) - _atomic_save_npz( - target_path, - features=features, - targets=targets, - feature_names=FEATURE_NAMES, - target_names=TARGET_NAMES, - metadata=json.dumps(_metadata_json(metadata), sort_keys=True), - ) - points = int(features.shape[0]) + record, points = process_raw_case_to_npz(case_dir, out_root, force=force) total_points += points - records.append({"case_id": case_dir.name, "path": str(target_path), "points": points, "metadata": _metadata_json(metadata)}) + records.append(record) if progress_every is not None and progress_every > 0 and (index % progress_every == 0 or index == len(case_dirs)): print(f"processed_airfrans_cases={index}/{len(case_dirs)} total_points={total_points}", flush=True) + return write_processing_manifest(out_root, raw_root, records=records, total_points=total_points, started=started) + + +def process_raw_case_to_npz(case_dir: str | Path, output_dir: str | Path, *, force: bool = False) -> tuple[dict[str, object], int]: + case_path = Path(case_dir).expanduser() + out_root = Path(output_dir).expanduser() + out_root.mkdir(parents=True, exist_ok=True) + target_path = out_root / f"{case_path.name}.npz" + if target_path.exists() and not force: + with np.load(target_path, allow_pickle=False) as npz: + points = int(npz["features"].shape[0]) + return {"case_id": case_path.name, "path": str(target_path), "points": points, "skipped_existing": True}, points + + metadata, features, targets = process_raw_case(case_path) + _atomic_save_npz( + target_path, + features=features, + targets=targets, + feature_names=FEATURE_NAMES, + target_names=TARGET_NAMES, + metadata=json.dumps(_metadata_json(metadata), sort_keys=True), + ) + points = int(features.shape[0]) + return {"case_id": case_path.name, "path": str(target_path), "points": points, "metadata": _metadata_json(metadata)}, points + + +def write_processing_manifest( + output_dir: str | Path, + raw_dir: str | Path, + *, + records: list[dict[str, object]], + total_points: int, + started: float, +) -> ProcessingResult: + out_root = Path(output_dir).expanduser() manifest = { - "raw_dir": str(raw_root), + "raw_dir": str(raw_dir), "output_dir": str(out_root), "case_count": len(records), "total_points": total_points, diff --git a/src/airfrans_frontier/raw/public.py b/src/airfrans_frontier/raw/public.py index c54fb67..d9693f9 100644 --- a/src/airfrans_frontier/raw/public.py +++ b/src/airfrans_frontier/raw/public.py @@ -3,12 +3,16 @@ from __future__ import annotations import json import os import shutil +import struct import time import urllib.error import urllib.request +import urllib.parse +import zlib import zipfile -from pathlib import Path -from typing import Any +from dataclasses import dataclass +from pathlib import Path, PurePosixPath +from typing import Any, Protocol from airfrans_frontier.training.data_sources import publish_processed_dataset @@ -16,6 +20,94 @@ PUBLIC_OF_DATASET_URL = "https://data.isir.upmc.fr/extrality/NeurIPS_2022/OF_dat DEFAULT_PUBLIC_WORK_DIR = Path("artifacts/public_airfrans") DEFAULT_PUBLIC_OUTPUT_DIR = Path("artifacts/data_cache/airfrans_processed/processed/full") +_EOCD_SIGNATURE = b"PK\x05\x06" +_ZIP64_EOCD_LOCATOR_SIGNATURE = 0x07064B50 +_ZIP64_EOCD_SIGNATURE = 0x06064B50 +_CENTRAL_DIRECTORY_SIGNATURE = 0x02014B50 +_LOCAL_FILE_HEADER_SIGNATURE = 0x04034B50 +_ZIP64_EXTRA_ID = 0x0001 +_ZIP64_LIMIT_16 = 0xFFFF +_ZIP64_LIMIT_32 = 0xFFFFFFFF + + +class RangeReader(Protocol): + size: int + bytes_read: int + + def read_range(self, start: int, length: int) -> bytes: ... + + +@dataclass(frozen=True) +class RemoteZipMember: + filename: str + flag_bits: int + compress_type: int + compress_size: int + file_size: int + header_offset: int + + @property + def is_dir(self) -> bool: + return self.filename.endswith("/") + + +@dataclass(frozen=True) +class StreamingZipProcessingResult: + processing: object + source_bytes: int + ranged_bytes_read: int + + +class PathRangeReader: + def __init__(self, path: str | Path) -> None: + self.path = Path(path).expanduser() + self.size = self.path.stat().st_size + self.bytes_read = 0 + + def read_range(self, start: int, length: int) -> bytes: + _validate_range(start, length, self.size) + if length == 0: + return b"" + with self.path.open("rb") as handle: + handle.seek(start) + data = handle.read(length) + if len(data) != length: + raise RuntimeError(f"Local range read returned {len(data)} bytes; expected {length}") + self.bytes_read += len(data) + return data + + +class HttpRangeReader: + def __init__(self, url: str) -> None: + self.url = url + size = _remote_content_length(url) + if size is None: + raise RuntimeError(f"Could not determine remote content length for range streaming: {url}") + self.size = size + self.bytes_read = 0 + + def read_range(self, start: int, length: int) -> bytes: + _validate_range(start, length, self.size) + if length == 0: + return b"" + end = start + length - 1 + request = urllib.request.Request(self.url, headers={"Range": f"bytes={start}-{end}"}) + try: + with urllib.request.urlopen(request, timeout=60) as response: + status = getattr(response, "status", None) + data = response.read() + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + raise RuntimeError(f"HTTP range read failed for {self.url} bytes={start}-{end}: {exc.code} {body}") from exc + except OSError as exc: + raise RuntimeError(f"HTTP range read failed for {self.url} bytes={start}-{end}: {exc}") from exc + if status != 206: + raise RuntimeError(f"Server did not honor HTTP Range for {self.url}: status={status}") + if len(data) != length: + raise RuntimeError(f"HTTP range read returned {len(data)} bytes; expected {length}") + self.bytes_read += len(data) + return data + def ensure_public_airfrans_processed_hf( *, repo_id: str, @@ -50,14 +142,17 @@ def ensure_public_airfrans_processed_hf( work_root.mkdir(parents=True, exist_ok=True) output_root.mkdir(parents=True, exist_ok=True) - archive_path = work_root / "OF_dataset.zip" - download = download_file(source_url, archive_path) - extract_root = work_root / "raw" - raw_root = extract_of_dataset(archive_path, extract_root, min_cases=min_cases) - from airfrans_frontier.raw.process import process_raw_dataset - - print(f"process_airfrans_raw raw_root={raw_root} output_dir={output_root}", flush=True) - processed = process_raw_dataset(raw_root, output_root, force=force, progress_every=25) + scratch_root = work_root / "streaming_raw" + print(f"range_stream_process_airfrans_zip source={source_url} output_dir={output_root}", flush=True) + streamed = process_of_dataset_url_streaming( + source_url, + output_root, + scratch_dir=scratch_root, + min_cases=min_cases, + force=force, + progress_every=25, + ) + processed = streamed.processing if processed.case_count < min_cases: raise RuntimeError(f"Processed only {processed.case_count} cases from public AirfRANS archive; expected at least {min_cases}") print(f"publish_airfrans_processed_hf repo={repo_id} path_in_repo={prefix}", flush=True) @@ -81,14 +176,16 @@ def ensure_public_airfrans_processed_hf( "repo_url": f"https://huggingface.co/datasets/{repo_id}", "path_in_repo": prefix, "source_url": source_url, - "archive_path": str(archive_path), - "archive_bytes": archive_path.stat().st_size, - "raw_root": str(raw_root), + "streaming": True, + "streaming_mode": "zip_range", + "streaming_scratch_dir": str(scratch_root), + "source_bytes": streamed.source_bytes, + "ranged_bytes_read": streamed.ranged_bytes_read, "output_dir": str(output_root), "processed_case_count": processed.case_count, "processed_total_points": processed.total_points, "processed_manifest_path": str(processed.manifest_path), - "download": download, + "download": {"url": source_url, "mode": "zip_range", "source_bytes": streamed.source_bytes, "ranged_bytes_read": streamed.ranged_bytes_read}, "publish": publish, "elapsed_seconds": time.time() - started, **final, @@ -145,6 +242,375 @@ def download_file(url: str, destination: str | Path, *, chunk_size: int = 16 * 1 return {"url": url, "path": str(path), "bytes": final_size, "resumed": resumed, "skipped": False} + +def process_of_dataset_url_streaming( + source_url: str, + output_dir: str | Path, + *, + scratch_dir: str | Path, + min_cases: int = 1000, + force: bool = False, + progress_every: int | None = None, +) -> StreamingZipProcessingResult: + if min_cases <= 0: + raise ValueError("min_cases must be positive") + reader = _range_reader_for(source_url) + members = _read_zip_central_directory(reader) + processing = _process_remote_zip_members( + reader, + members, + output_dir, + scratch_dir=scratch_dir, + raw_dir_label=f"{source_url}!OF_dataset", + min_cases=min_cases, + force=force, + progress_every=progress_every, + ) + print( + f"range_stream_airfrans_bytes_read={reader.bytes_read} range_stream_airfrans_source_bytes={reader.size}", + flush=True, + ) + return StreamingZipProcessingResult( + processing=processing, + source_bytes=reader.size, + ranged_bytes_read=reader.bytes_read, + ) + + +def _range_reader_for(source_url: str) -> RangeReader: + parsed = urllib.parse.urlparse(source_url) + if parsed.scheme in {"http", "https"}: + return HttpRangeReader(source_url) + if parsed.scheme == "file": + return PathRangeReader(Path(urllib.request.url2pathname(parsed.path))) + if not parsed.scheme: + return PathRangeReader(source_url) + raise RuntimeError(f"Unsupported AirfRANS streaming URL scheme: {parsed.scheme}") + + +def _read_zip_central_directory(reader: RangeReader) -> list[RemoteZipMember]: + tail_size = min(reader.size, 1024 * 1024) + tail_start = reader.size - tail_size + tail = reader.read_range(tail_start, tail_size) + eocd_index = tail.rfind(_EOCD_SIGNATURE) + if eocd_index < 0: + raise RuntimeError("ZIP end-of-central-directory record not found") + eocd_offset = tail_start + eocd_index + eocd = tail[eocd_index : eocd_index + 22] + if len(eocd) < 22: + raise RuntimeError("Truncated ZIP end-of-central-directory record") + ( + _signature, + _disk_number, + _central_disk, + disk_entries, + total_entries, + central_size, + central_offset, + _comment_length, + ) = struct.unpack(" tuple[int, int, int]: + locator_offset = eocd_offset - 20 + if locator_offset < 0: + raise RuntimeError("ZIP64 end-of-central-directory locator is missing") + locator = reader.read_range(locator_offset, 20) + signature, _disk_with_record, zip64_eocd_offset, _disk_count = struct.unpack(" list[RemoteZipMember]: + members: list[RemoteZipMember] = [] + offset = 0 + while offset < len(central): + if offset + 46 > len(central): + raise RuntimeError("Truncated ZIP central directory entry") + fields = struct.unpack_from(" len(central): + raise RuntimeError("Truncated ZIP central directory variable fields") + filename_bytes = central[name_start:extra_start] + encoding = "utf-8" if flag_bits & 0x800 else "cp437" + filename = filename_bytes.decode(encoding, errors="replace") + extra = central[extra_start:comment_start] + file_size, compress_size, header_offset = _apply_zip64_extra( + extra, + file_size=file_size, + compress_size=compress_size, + header_offset=header_offset, + ) + members.append( + RemoteZipMember( + filename=filename, + flag_bits=flag_bits, + compress_type=compress_type, + compress_size=compress_size, + file_size=file_size, + header_offset=header_offset, + ) + ) + offset = next_offset + if expected_entries not in (0, len(members)): + raise RuntimeError(f"ZIP central directory entry count mismatch: parsed={len(members)} expected={expected_entries}") + return members + + +def _apply_zip64_extra(extra: bytes, *, file_size: int, compress_size: int, header_offset: int) -> tuple[int, int, int]: + values_needed = [ + file_size == _ZIP64_LIMIT_32, + compress_size == _ZIP64_LIMIT_32, + header_offset == _ZIP64_LIMIT_32, + ] + if not any(values_needed): + return file_size, compress_size, header_offset + offset = 0 + while offset + 4 <= len(extra): + header_id, data_size = struct.unpack_from(" len(extra): + raise RuntimeError("Truncated ZIP extra field") + if header_id == _ZIP64_EXTRA_ID: + cursor = data_start + resolved = [file_size, compress_size, header_offset] + for index, needed in enumerate(values_needed): + if needed: + if cursor + 8 > data_end: + raise RuntimeError("Truncated ZIP64 extra field") + resolved[index] = struct.unpack_from(" 0 and (index % progress_every == 0 or index == len(case_names)): + print(f"range_streamed_airfrans_cases={index}/{len(case_names)} total_points={total_points}", flush=True) + + try: + scratch_root.rmdir() + except OSError: + pass + return write_processing_manifest( + out_root, + raw_dir_label, + records=records, + total_points=total_points, + started=started, + ) + + +def _remote_archive_case_members(members: list[RemoteZipMember]) -> dict[str, list[tuple[RemoteZipMember, PurePosixPath]]]: + cases: dict[str, list[tuple[RemoteZipMember, PurePosixPath]]] = {} + for member in members: + parsed = _case_member_parts_from_name(member.filename) + if parsed is None: + continue + case_name, relative = parsed + cases.setdefault(case_name, []).append((member, relative)) + return cases + + +def _extract_remote_case_members( + reader: RangeReader, + members: list[tuple[RemoteZipMember, PurePosixPath]], + root: Path, +) -> None: + resolved_root = root.resolve() + for member, relative in members: + target = _safe_relative_target(root, relative, resolved_root=resolved_root) + if member.is_dir: + target.mkdir(parents=True, exist_ok=True) + continue + target.parent.mkdir(parents=True, exist_ok=True) + payload = _read_remote_member_payload(reader, member) + target.write_bytes(payload) + + +def _read_remote_member_payload(reader: RangeReader, member: RemoteZipMember) -> bytes: + if member.flag_bits & 0x1: + raise RuntimeError(f"Encrypted ZIP member is unsupported: {member.filename}") + local_header = reader.read_range(member.header_offset, 30) + ( + signature, + _version_needed, + _flag_bits, + _compress_type, + _mod_time, + _mod_date, + _crc, + _compress_size, + _file_size, + filename_length, + extra_length, + ) = struct.unpack(" None: + if start < 0 or length < 0 or start + length > size: + raise RuntimeError(f"Invalid range start={start} length={length} size={size}") + + +def process_of_dataset_archive_streaming( + archive_path: str | Path, + output_dir: str | Path, + *, + scratch_dir: str | Path, + min_cases: int = 1000, + force: bool = False, + progress_every: int | None = None, +): + if min_cases <= 0: + raise ValueError("min_cases must be positive") + archive = Path(archive_path).expanduser() + out_root = Path(output_dir).expanduser() + scratch_root = Path(scratch_dir).expanduser() + out_root.mkdir(parents=True, exist_ok=True) + if scratch_root.exists(): + shutil.rmtree(scratch_root) + scratch_root.mkdir(parents=True, exist_ok=True) + + from airfrans_frontier.raw.process import process_raw_case_to_npz, write_processing_manifest + + records: list[dict[str, object]] = [] + total_points = 0 + started = time.perf_counter() + with zipfile.ZipFile(archive) as zf: + case_members = _archive_case_members(zf.infolist()) + case_names = sorted(case_members) + if len(case_names) < min_cases: + raise RuntimeError(f"AirfRANS archive has {len(case_names)} cases; expected at least {min_cases}") + print(f"stream_airfrans_archive_cases={len(case_names)}", flush=True) + for index, case_name in enumerate(case_names, start=1): + case_dir = scratch_root / case_name + target_path = out_root / f"{case_name}.npz" + if target_path.exists() and not force: + record, points = process_raw_case_to_npz(case_dir, out_root, force=False) + else: + try: + _extract_case_members(zf, case_members[case_name], scratch_root) + record, points = process_raw_case_to_npz(case_dir, out_root, force=force) + finally: + if case_dir.exists(): + shutil.rmtree(case_dir, ignore_errors=True) + records.append(record) + total_points += points + if progress_every is not None and progress_every > 0 and (index % progress_every == 0 or index == len(case_names)): + print(f"streamed_airfrans_cases={index}/{len(case_names)} total_points={total_points}", flush=True) + + try: + scratch_root.rmdir() + except OSError: + pass + + return write_processing_manifest( + out_root, + f"{archive}!OF_dataset", + records=records, + total_points=total_points, + started=started, + ) + + def extract_of_dataset(archive_path: str | Path, extract_root: str | Path, *, min_cases: int = 1000) -> Path: archive = Path(archive_path).expanduser() root = Path(extract_root).expanduser() @@ -155,6 +621,7 @@ def extract_of_dataset(archive_path: str | Path, extract_root: str | Path, *, mi print(f"extract_airfrans_zip archive={archive} root={root}", flush=True) with zipfile.ZipFile(archive) as zf: members = zf.infolist() + _require_extract_space(root, members) for index, member in enumerate(members, start=1): _safe_extract_member(zf, member, root) if index % 1000 == 0 or index == len(members): @@ -206,11 +673,7 @@ def _remote_content_length(url: str) -> int | None: def _safe_extract_member(zf: zipfile.ZipFile, member: zipfile.ZipInfo, root: Path) -> None: - target = root / member.filename - resolved_root = root.resolve() - resolved_target = target.resolve() - if resolved_root != resolved_target and resolved_root not in resolved_target.parents: - raise RuntimeError(f"Unsafe path in AirfRANS archive: {member.filename}") + target = _safe_member_target(member, root) if member.is_dir(): target.mkdir(parents=True, exist_ok=True) return @@ -219,6 +682,94 @@ def _safe_extract_member(zf: zipfile.ZipFile, member: zipfile.ZipInfo, root: Pat shutil.copyfileobj(source, destination, length=16 * 1024 * 1024) +def _require_extract_space(root: Path, members: list[zipfile.ZipInfo]) -> None: + total_uncompressed_bytes = 0 + remaining_uncompressed_bytes = 0 + resolved_root = root.resolve() + for member in members: + if member.is_dir(): + continue + total_uncompressed_bytes += member.file_size + target = _safe_member_target(member, root, resolved_root=resolved_root) + try: + existing_size = target.stat().st_size + except OSError: + existing_size = None + if existing_size == member.file_size: + continue + remaining_uncompressed_bytes += member.file_size + + margin_bytes = max(1024**3, remaining_uncompressed_bytes // 20) if remaining_uncompressed_bytes else 0 + required_free_bytes = remaining_uncompressed_bytes + margin_bytes + usage = shutil.disk_usage(root) + print( + "airfrans_extract_total_uncompressed_bytes=" + f"{total_uncompressed_bytes} airfrans_extract_remaining_uncompressed_bytes={remaining_uncompressed_bytes} " + f"airfrans_extract_free_disk_bytes={usage.free} airfrans_extract_required_free_bytes={required_free_bytes}", + flush=True, + ) + if usage.free < required_free_bytes: + raise RuntimeError( + "Insufficient free disk for AirfRANS extraction: " + f"free={usage.free} required={required_free_bytes} remaining_uncompressed={remaining_uncompressed_bytes}; " + "provision more disk or use a streaming/incremental extraction pipeline" + ) + + +def _safe_member_target(member: zipfile.ZipInfo, root: Path, *, resolved_root: Path | None = None) -> Path: + return _safe_relative_target(root, PurePosixPath(member.filename), resolved_root=resolved_root) + + +def _safe_relative_target(root: Path, relative: PurePosixPath, *, resolved_root: Path | None = None) -> Path: + target = root.joinpath(*relative.parts) + actual_root = resolved_root or root.resolve() + resolved_target = target.resolve() + if actual_root != resolved_target and actual_root not in resolved_target.parents: + raise RuntimeError(f"Unsafe path in AirfRANS archive: {relative}") + return target + + +def _archive_case_members(members: list[zipfile.ZipInfo]) -> dict[str, list[tuple[zipfile.ZipInfo, PurePosixPath]]]: + cases: dict[str, list[tuple[zipfile.ZipInfo, PurePosixPath]]] = {} + for member in members: + parsed = _case_member_parts(member) + if parsed is None: + continue + case_name, relative = parsed + cases.setdefault(case_name, []).append((member, relative)) + return cases + + +def _case_member_parts(member: zipfile.ZipInfo) -> tuple[str, PurePosixPath] | None: + return _case_member_parts_from_name(member.filename) + + +def _case_member_parts_from_name(filename: str) -> tuple[str, PurePosixPath] | None: + parts = PurePosixPath(filename).parts + if any(part == ".." for part in parts): + raise RuntimeError(f"Unsafe path in AirfRANS archive: {filename}") + for index, part in enumerate(parts): + if part.startswith("airFoil2D_"): + return part, PurePosixPath(*parts[index:]) + return None + + +def _extract_case_members( + zf: zipfile.ZipFile, + members: list[tuple[zipfile.ZipInfo, PurePosixPath]], + root: Path, +) -> None: + resolved_root = root.resolve() + for member, relative in members: + target = _safe_relative_target(root, relative, resolved_root=resolved_root) + if member.is_dir(): + target.mkdir(parents=True, exist_ok=True) + continue + target.parent.mkdir(parents=True, exist_ok=True) + with zf.open(member) as source, target.open("wb") as destination: + shutil.copyfileobj(source, destination, length=16 * 1024 * 1024) + + def _find_of_dataset_root(root: Path) -> Path | None: direct = root / "OF_dataset" if direct.is_dir(): @@ -247,6 +798,17 @@ def _optional_secret(name: str) -> str | None: return None +def _remove_file_best_effort(path: Path) -> bool: + try: + path.unlink() + return True + except FileNotFoundError: + return False + except OSError as exc: + print(f"warning: could not remove {path}: {exc}", flush=True) + return False + + def write_json_report(path: str | Path, payload: dict[str, Any]) -> None: report_path = Path(path).expanduser() report_path.parent.mkdir(parents=True, exist_ok=True) diff --git a/src/airfrans_frontier/remote/artifacts.py b/src/airfrans_frontier/remote/artifacts.py index 96dbc8a..48f4ed3 100644 --- a/src/airfrans_frontier/remote/artifacts.py +++ b/src/airfrans_frontier/remote/artifacts.py @@ -50,6 +50,7 @@ def verify_artifacts( if missing_failure: raise ValueError(f"Failed artifact directory missing files: {', '.join(missing_failure)}") + prior_checks = _prior_verification_checks(root / "verification_report.json") checks: dict[str, Any] = { "required_files": {name: True for name in required_names}, "terminal_artifact": "final_metrics.json" if has_final else "failure_report.json" if has_failure else None, @@ -67,19 +68,24 @@ def verify_artifacts( "evaluation_protocol.json", "artifact_manifest.json", "hf_upload_manifest.json", + "artifact_collection_report.json", + "disk_telemetry.json", "verification_report.json", ): path = root / json_name if path.is_file(): _validate_json(path) checks[f"json:{json_name}"] = True - _validate_jsonl(root / "metrics.jsonl") - checks["jsonl:metrics.jsonl"] = True + if (root / "metrics.jsonl").is_file(): + _validate_jsonl(root / "metrics.jsonl") + checks["jsonl:metrics.jsonl"] = True for checkpoint_name in ("checkpoint_latest.pt", "checkpoint_best.pt", "checkpoint_final.pt"): path = root / checkpoint_name if path.is_file(): - _validate_checkpoint_metadata(path) - checks[f"checkpoint:{checkpoint_name}"] = True + checkpoint_check = f"checkpoint:{checkpoint_name}" + if prior_checks.get(checkpoint_check) is not True: + _validate_checkpoint_metadata(path) + checks[checkpoint_check] = True if (root / "hf_upload_manifest.json").is_file(): checks["hf_upload_manifest.json"] = _validate_hf_upload_manifest(root / "hf_upload_manifest.json") @@ -128,6 +134,16 @@ def sha256_file(path: Path) -> str: digest.update(chunk) return digest.hexdigest() +def _prior_verification_checks(path: Path) -> dict[str, Any]: + if not path.is_file(): + return {} + try: + report = json.loads(path.read_text()) + except json.JSONDecodeError: + return {} + checks = report.get("checks") if isinstance(report, dict) else None + return dict(checks) if isinstance(checks, dict) else {} + def _validate_json(path: Path) -> None: try: diff --git a/src/airfrans_frontier/remote/cleanup.py b/src/airfrans_frontier/remote/cleanup.py new file mode 100644 index 0000000..64a3b12 --- /dev/null +++ b/src/airfrans_frontier/remote/cleanup.py @@ -0,0 +1,166 @@ +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +import time +from typing import Any + +_TERMINAL_INSTANCE_STATUSES = { + "deleted", + "destroyed", + "exited", + "offline", + "stopped", + "stopping", + "terminated", +} + + +def reconcile_cleanup( + *, + sky_state: Any, + vast_instances: Sequence[Mapping[str, Any]], + known_run_ids: Sequence[str] = (), + destroy_orphans: bool = False, + destroy_instance: Callable[[int], Any] | None = None, + now: float | None = None, +) -> dict[str, Any]: + """Reconcile Sky's view with Vast API ground truth and report cleanup actions. + + Vast instances are treated as the paid-resource ground truth. Destruction is + opt-in so this can be used as a non-launch-blocking inspection command. + """ + + checked_at = time.time() if now is None else float(now) + sky_refs = _extract_sky_refs(sky_state) + known_runs = tuple(known_run_ids) + records: list[dict[str, Any]] = [] + for instance in vast_instances: + instance_id = _instance_id(instance) + status = _status(instance) + associated_run_id = _associated_run_id(instance, known_runs) + sky_knows = _sky_knows_instance(sky_refs, instance_id=instance_id, run_id=associated_run_id) + live = _is_live_status(status) + unexpected_live = bool(live and not sky_knows) + action = "none" + result = "not_needed" + error = None + if unexpected_live: + action = "destroy_orphan" if destroy_orphans else "report_orphan" + result = "not_attempted" + if destroy_orphans: + if destroy_instance is None: + result = "skipped_no_destroy_function" + elif instance_id is None: + result = "skipped_missing_instance_id" + else: + try: + destroy_instance(int(instance_id)) + except Exception as exc: # pragma: no cover - exercised by callers with fakes. + result = "failed" + error = str(exc) + else: + result = "destroy_requested" + records.append( + { + "vast_instance_id": instance_id, + "host_id": _first_present(instance, "host_id", "machine_id"), + "gpu_type": _first_present(instance, "gpu_name", "gpu", "gpu_type"), + "gpu_count": _first_present(instance, "num_gpus", "gpu_count", "gpus"), + "status": status, + "associated_run_id": associated_run_id, + "hourly_cost": _first_present(instance, "dph_total", "hourly_cost", "cost_per_hour"), + "sky_known": sky_knows, + "live": live, + "unexpected_live": unexpected_live, + "cleanup_action_attempted": action, + "cleanup_result": result, + "error": error, + } + ) + return { + "schema_version": 1, + "checked_at": checked_at, + "sky_instance_ids": sorted(sky_refs["instance_ids"]), + "sky_run_ids": sorted(sky_refs["run_ids"]), + "destroy_orphans": destroy_orphans, + "unexpected_live_count": sum(1 for record in records if record["unexpected_live"]), + "destroy_requested_count": sum(1 for record in records if record["cleanup_result"] == "destroy_requested"), + "instances": records, + } + + +def _extract_sky_refs(value: Any) -> dict[str, set[str]]: + refs = {"instance_ids": set(), "run_ids": set()} + _walk_sky(value, refs) + return refs + + +def _walk_sky(value: Any, refs: dict[str, set[str]]) -> None: + if isinstance(value, Mapping): + for key, item in value.items(): + key_text = str(key).lower() + if key_text in {"id", "instance_id", "vast_instance_id"}: + _add_ref(refs["instance_ids"], item) + elif key_text in {"name", "cluster", "cluster_name", "run_id", "label"}: + _add_ref(refs["run_ids"], item) + _walk_sky(item, refs) + elif isinstance(value, (list, tuple)): + for item in value: + _walk_sky(item, refs) + + +def _add_ref(target: set[str], value: Any) -> None: + if isinstance(value, bool) or value is None: + return + if isinstance(value, (int, float, str)): + text = str(int(value)) if isinstance(value, float) and value.is_integer() else str(value) + if text: + target.add(text) + + +def _sky_knows_instance(refs: Mapping[str, set[str]], *, instance_id: int | None, run_id: str | None) -> bool: + if instance_id is not None and str(instance_id) in refs["instance_ids"]: + return True + if run_id is not None and run_id in refs["run_ids"]: + return True + return False + + +def _instance_id(instance: Mapping[str, Any]) -> int | None: + value = _first_present(instance, "id", "instance_id", "vast_instance_id") + if isinstance(value, bool) or value is None: + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + +def _status(instance: Mapping[str, Any]) -> str | None: + value = _first_present(instance, "actual_status", "status", "state") + return str(value) if value is not None else None + + +def _is_live_status(status: str | None) -> bool: + if status is None: + return True + return status.lower() not in _TERMINAL_INSTANCE_STATUSES + + +def _associated_run_id(instance: Mapping[str, Any], known_run_ids: Sequence[str]) -> str | None: + for key in ("run_id", "label", "name", "cluster_name"): + value = instance.get(key) + if isinstance(value, str) and value: + if value in known_run_ids: + return value + for run_id in known_run_ids: + if run_id and run_id in value: + return run_id + return None + + +def _first_present(instance: Mapping[str, Any], *keys: str) -> Any: + for key in keys: + if key in instance and instance[key] is not None: + return instance[key] + return None diff --git a/src/airfrans_frontier/remote/cli.py b/src/airfrans_frontier/remote/cli.py index d3cf728..3370363 100644 --- a/src/airfrans_frontier/remote/cli.py +++ b/src/airfrans_frontier/remote/cli.py @@ -12,11 +12,15 @@ from pathlib import Path from typing import Any from airfrans_frontier.remote.artifacts import verify_artifacts +from airfrans_frontier.remote.cleanup import reconcile_cleanup +from airfrans_frontier.remote.collection import ARTIFACT_COLLECTION_REPORT, collect_artifact_paths, required_collection_failures from airfrans_frontier.remote.config import RemoteRunConfig, load_remote_run_config +from airfrans_frontier.remote.launch_group import LaunchGroupScheduler +from airfrans_frontier.remote.selection import DEFAULT_SELECTION_MAX_AGE_SECONDS, load_selection_manifest from airfrans_frontier.remote.skypilot import render_skypilot_yaml, write_skyignore from airfrans_frontier.remote.skypilot_patch import apply_patch, patch_status, require_patch from airfrans_frontier.remote.smoke import run_hf_upload_smoke, run_smoke_training, run_wandb_smoke -from airfrans_frontier.remote.vast import SelectionResult, select_offer +from airfrans_frontier.remote.vast import SelectionResult, destroy_instance, list_instances, select_offer, summarize_instances def build_parser() -> argparse.ArgumentParser: @@ -27,14 +31,38 @@ def build_parser() -> argparse.ArgumentParser: doctor.add_argument("--apply-skypilot-patch", action="store_true") doctor.set_defaults(command="doctor") + vast_instances = subparsers.add_parser("vast-instances", help="list Vast.ai instances using the Vast API") + vast_instances.add_argument("--base-url", default="https://cloud.vast.ai") + vast_instances.add_argument("--out") + vast_instances.set_defaults(command="vast-instances") + + cleanup = subparsers.add_parser("cleanup-reconcile", help="reconcile Sky status against Vast API ground truth") + cleanup.add_argument("--sky-status-json", help="local Sky status JSON; omit to call sky status") + cleanup.add_argument("--vast-instances-json", help="local Vast instances JSON; omit to call Vast API") + cleanup.add_argument("--base-url", default="https://cloud.vast.ai") + cleanup.add_argument("--destroy-orphans", action="store_true", help="request Vast destruction for live instances missing from Sky") + cleanup.add_argument("--out") + cleanup.set_defaults(command="cleanup-reconcile") + select = subparsers.add_parser("select", help="select a Vast.ai offer from a remote config") select.add_argument("config") select.add_argument("--out") select.set_defaults(command="select") + launch_group = subparsers.add_parser("launch-group-plan", help="write local launch-group state without provisioning") + launch_group.add_argument("configs", nargs="+") + launch_group.add_argument("--max-active", type=int, default=4) + launch_group.add_argument("--max-fragile", type=int, default=1) + launch_group.add_argument("--state") + launch_group.add_argument("--group-id") + launch_group.add_argument("--allow-duplicate-hosts", action="store_true") + launch_group.set_defaults(command="launch-group-plan") + render = subparsers.add_parser("render", help="render patched SkyPilot YAML") render.add_argument("config") render.add_argument("--selection", required=True) + render.add_argument("--selection-max-age-seconds", type=float, default=DEFAULT_SELECTION_MAX_AGE_SECONDS) + render.add_argument("--allow-stale-selection", action="store_true") render.add_argument("--run-id", required=True) render.add_argument("--out") render.set_defaults(command="render") @@ -85,14 +113,51 @@ def main(argv: list[str] | None = None) -> int: remove_pythonpath_entries() if args.command == "doctor": return _doctor(apply=args.apply_skypilot_patch) + if args.command == "vast-instances": + api_key = os.environ.get("VAST_API_KEY") + if not api_key: + raise RuntimeError("VAST_API_KEY is required to list Vast instances") + instances = list_instances(base_url=args.base_url, api_key=api_key) + _emit_json({"instance_count": len(instances), "instances": summarize_instances(instances)}, args.out) + return 0 + if args.command == "cleanup-reconcile": + sky_state = _load_json_file(Path(args.sky_status_json)) if args.sky_status_json else _load_sky_status() + if args.vast_instances_json: + vast_payload = _load_json_file(Path(args.vast_instances_json)) + instances = _instances_from_json_payload(vast_payload) + else: + api_key = os.environ.get("VAST_API_KEY") + if not api_key: + raise RuntimeError("VAST_API_KEY is required to reconcile live Vast instances") + instances = list_instances(base_url=args.base_url, api_key=api_key) + destroy = None + if args.destroy_orphans: + api_key = os.environ.get("VAST_API_KEY") + if not api_key: + raise RuntimeError("VAST_API_KEY is required to destroy Vast orphan instances") + destroy = lambda instance_id: destroy_instance(base_url=args.base_url, api_key=api_key, instance_id=instance_id) + report = reconcile_cleanup(sky_state=sky_state, vast_instances=instances, destroy_orphans=args.destroy_orphans, destroy_instance=destroy) + _emit_json(report, args.out) + return 0 if args.command == "select": config = load_remote_run_config(args.config) result = select_offer(config) _emit_json(result.to_manifest(), args.out) return 0 + if args.command == "launch-group-plan": + scheduler = LaunchGroupScheduler( + args.configs, + max_active=args.max_active, + max_fragile=args.max_fragile, + state_path=args.state, + group_id=args.group_id, + allow_duplicate_hosts=args.allow_duplicate_hosts, + ) + _emit_json(scheduler.to_payload(), None) + return 0 if args.command == "render": config = load_remote_run_config(args.config) - selection = _selection_from_manifest(Path(args.selection)) + selection = _selection_from_manifest(Path(args.selection), max_age_seconds=args.selection_max_age_seconds, allow_stale=args.allow_stale_selection) text = render_skypilot_yaml(config, selection, run_id=args.run_id) if args.out: Path(args.out).write_text(text) @@ -146,7 +211,7 @@ def _doctor(*, apply: bool) -> int: problems: list[str] = [] if not os.environ.get("VAST_API_KEY"): problems.append("VAST_API_KEY is not set") - sky = shutil.which("sky") + sky = shutil.which("sky", path=_subprocess_env().get("PATH")) if not sky: problems.append("sky executable not found on PATH") if apply: @@ -180,6 +245,21 @@ def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int: local_run_dir = config.run.local_artifact_dir / run_id local_run_dir.mkdir(parents=True, exist_ok=False) state_path = local_run_dir / "orchestrator_state.json" + timeline_path = local_run_dir / "startup_timeline.jsonl" + submitted_at = time.time() + + def timeline(phase: str, event: str, **extra: Any) -> None: + record = { + "run_id": run_id, + "ts": time.time(), + "elapsed_since_submit_seconds": time.time() - submitted_at, + "phase": phase, + "event": event, + **extra, + } + with timeline_path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + def state(phase: str, **extra: Any) -> None: payload = { @@ -190,7 +270,7 @@ def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int: **extra, } state_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - + timeline("orchestrator", phase.lower(), orchestrator_phase=phase, **extra) state("SELECTING_OFFER") selection = select_offer(config) selection_path = local_run_dir / "selection_manifest.json" @@ -231,6 +311,7 @@ def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int: attempt=attempt, resume_checkpoint=str(resume_checkpoint) if resume_checkpoint is not None else None, ) + timeline("sky_launch", "started", attempt=attempt, selected_offer_id=selection.selected_offer_id) return_code = _run_sky_with_periodic_collection( cluster=run_id, sky_yaml_path=sky_yaml_path, @@ -238,6 +319,7 @@ def _run(config_path: str | Path, *, dry_run: bool, skip_down: bool) -> int: local_run_dir=local_run_dir, env=env, ) + timeline("sky_launch", "completed", attempt=attempt, return_code=return_code, selected_offer_id=selection.selected_offer_id) state("REMOTE_FINISHED", selected_offer_id=selection.selected_offer_id, attempt=attempt, return_code=return_code) _collect_terminal_best_effort(cluster=run_id, remote_dir=config.job.artifact_dir, local_dir=local_run_dir, required=config.artifacts.required, env=env) status = _classify_artifacts(local_run_dir) @@ -365,9 +447,13 @@ def _collect_paths_with_rsync( paths: tuple[str, ...], env: dict[str, str], timeout: int, + required: tuple[str, ...] = (), + collection_kind: str = "artifact", + raise_on_required: bool = True, ) -> None: local_dir.mkdir(parents=True, exist_ok=True) - for relative_path in paths: + + def copy_one(relative_path: str) -> int | None: source = f"{cluster}:~/sky_workdir/{remote_dir}/./{relative_path}" _run_checked( [ @@ -383,28 +469,49 @@ def _collect_paths_with_rsync( env=env, timeout=timeout, ) + return 0 + + report = collect_artifact_paths( + local_dir=local_dir, + remote_dir=f"{cluster}:~/sky_workdir/{remote_dir}", + paths=paths, + required=required, + collection_kind=collection_kind, + copy_one=copy_one, + ) + failures = required_collection_failures(report, paths=paths) + if failures and raise_on_required: + names = ", ".join(str(item["expected_path"]) for item in failures) + raise RuntimeError(f"Required artifact collection failed: {names}") def _collect_required_artifacts(*, cluster: str, remote_dir: Path, local_dir: Path, required: tuple[str, ...], env: dict[str, str]) -> None: + large = _large_artifact_names(required) _collect_paths_with_rsync( cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, - paths=_large_artifact_names(required), + paths=large, env=env, timeout=3600, + required=large, + collection_kind="large", ) def _collect_terminal_best_effort(*, cluster: str, remote_dir: Path, local_dir: Path, required: tuple[str, ...], env: dict[str, str]) -> None: try: + terminal = _terminal_artifact_names(required) _collect_paths_with_rsync( cluster=cluster, remote_dir=remote_dir, local_dir=local_dir, - paths=_terminal_artifact_names(required), + paths=terminal, env=env, timeout=120, + required=tuple(name for name in terminal if name in required), + collection_kind="terminal", + raise_on_required=False, ) except Exception: _cleanup_partial_artifacts(local_dir) @@ -419,6 +526,8 @@ def _collect_restart_best_effort(*, cluster: str, remote_dir: Path, local_dir: P paths=("checkpoint_latest.pt",), env=env, timeout=3600, + collection_kind="restart", + raise_on_required=False, ) except Exception: _cleanup_partial_artifacts(local_dir) @@ -427,10 +536,12 @@ def _collect_restart_best_effort(*, cluster: str, remote_dir: Path, local_dir: P _LARGE_ARTIFACT_SUFFIXES = (".pt", ".pth", ".ckpt", ".safetensors") _TERMINAL_ARTIFACT_NAMES = ( "artifact_manifest.json", + ARTIFACT_COLLECTION_REPORT, "checksums.txt", "config.toml", "calibration_manifest.json", "data_manifest.json", + "disk_telemetry.json", "environment_manifest.json", "evaluation_protocol.json", "failure_report.json", @@ -442,6 +553,7 @@ _TERMINAL_ARTIFACT_NAMES = ( "normalization.json", "run_manifest.json", "split_manifest.json", + "startup_timeline.jsonl", "wandb_smoke_manifest.json", "verification_report.json", ) @@ -522,6 +634,9 @@ def _run_best_effort(argv: list[str], *, env: dict[str, str]) -> None: def _subprocess_env() -> dict[str, str]: env = dict(os.environ) env.pop("PYTHONPATH", None) + executable_dir = str(Path(sys.executable).parent) + path = env.get("PATH") + env["PATH"] = executable_dir if not path else f"{executable_dir}{os.pathsep}{path}" return env def _ensure_hf_secret_env(env: dict[str, str]) -> None: @@ -555,6 +670,33 @@ def _load_secret_env(env: dict[str, str], name: str, *, required: bool, purpose: +def _load_json_file(path: Path) -> Any: + return json.loads(path.read_text()) + + +def _load_sky_status() -> Any: + process = subprocess.run( + ["sky", "status", "--format", "json"], + check=True, + capture_output=True, + text=True, + env=_subprocess_env(), + timeout=120, + ) + return json.loads(process.stdout) + + +def _instances_from_json_payload(payload: Any) -> list[dict[str, Any]]: + if isinstance(payload, list): + return [dict(item) for item in payload if isinstance(item, dict)] + if isinstance(payload, dict): + for key in ("instances", "results", "items"): + value = payload.get(key) + if isinstance(value, list): + return [dict(item) for item in value if isinstance(item, dict)] + raise ValueError("Vast instances JSON must be a list or contain instances/results/items") + + def _emit_json(data: dict[str, Any], out: str | None) -> None: text = json.dumps(data, indent=2, sort_keys=True) + "\n" if out: @@ -563,14 +705,17 @@ def _emit_json(data: dict[str, Any], out: str | None) -> None: print(text, end="") -def _selection_from_manifest(path: Path) -> SelectionResult: - from airfrans_frontier.remote.vast import VastOffer, effective_price +def _selection_from_manifest(path: Path, *, max_age_seconds: float = DEFAULT_SELECTION_MAX_AGE_SECONDS, allow_stale: bool = False) -> SelectionResult: + from airfrans_frontier.remote.vast import VastOffer - data = json.loads(path.read_text()) + data = load_selection_manifest(path, max_age_seconds=max_age_seconds, allow_stale=allow_stale) raw_offer = data.get("selected_offer") if not isinstance(raw_offer, dict): raise ValueError(f"Selection manifest missing selected_offer object: {path}") offer = VastOffer.from_mapping({**raw_offer, "id": data.get("selected_offer_id", raw_offer.get("id"))}) + created_at = data.get("created_at") + if not isinstance(created_at, (int, float)): + created_at = time.time() # Preserve manifest values by building a minimal SelectionResult. Effective price is already stored. return SelectionResult( selected_offer=offer, @@ -579,6 +724,7 @@ def _selection_from_manifest(path: Path) -> SelectionResult: effective_price=float(raw_offer.get("effective_price", data.get("effective_price", 0.0))) if raw_offer else 0.0, query=data.get("query", {}) if isinstance(data.get("query"), dict) else {}, policy=data.get("policy", {}) if isinstance(data.get("policy"), dict) else {}, + created_at=float(created_at), ) diff --git a/src/airfrans_frontier/remote/collection.py b/src/airfrans_frontier/remote/collection.py new file mode 100644 index 0000000..3f84e2f --- /dev/null +++ b/src/airfrans_frontier/remote/collection.py @@ -0,0 +1,180 @@ +from __future__ import annotations + +from collections.abc import Callable, Iterable, Mapping +import json +from pathlib import Path +import time +from typing import Any + +ARTIFACT_COLLECTION_REPORT = "artifact_collection_report.json" +_PARTIAL_SUFFIXES = (".tmp", ".part", ".partial") +_RSYNC_TEMP_DIRS = (".rsync-partial", ".~tmp~") + + +def collect_artifact_paths( + *, + local_dir: str | Path, + remote_dir: str | Path, + paths: Iterable[str], + required: Iterable[str] = (), + collection_kind: str, + copy_one: Callable[[str], int | None], + clock: Callable[[], float] | None = None, +) -> dict[str, Any]: + """Copy artifact paths and update artifact_collection_report.json. + + copy_one receives each relative artifact path. It may raise or return a + non-zero return code; both are recorded per path without losing the rest of + the collection report. + """ + + now = clock or time.time + root = Path(local_dir) + root.mkdir(parents=True, exist_ok=True) + report_path = root / ARTIFACT_COLLECTION_REPORT + report = _load_report(report_path) + required_set = set(required) + attempted_paths = tuple(dict.fromkeys(paths)) + batch_started_at = now() + batch_id = f"{collection_kind}-{int(batch_started_at * 1000)}-{len(report['attempts'])}" + report["batches"].append( + { + "batch_id": batch_id, + "collection_kind": collection_kind, + "started_at": batch_started_at, + "paths": list(attempted_paths), + } + ) + _write_report(report_path, _refresh_summary(report, now=now())) + + for relative_path in attempted_paths: + _validate_relative_path(relative_path) + source = f"{str(remote_dir).rstrip('/')}/{relative_path}" + destination = root / relative_path + started = now() + attempt: dict[str, Any] = { + "batch_id": batch_id, + "collection_kind": collection_kind, + "expected_path": relative_path, + "required": relative_path in required_set, + "source_path": source, + "local_destination": str(destination), + "attempted": True, + "started_at": started, + "bytes_copied": None, + "duration_seconds": None, + "return_code": None, + "exception": None, + "final_status": "failed", + "likely_reason": None, + } + try: + return_code = copy_one(relative_path) + if return_code is not None: + attempt["return_code"] = int(return_code) + except Exception as exc: + attempt["exception"] = {"type": type(exc).__name__, "message": str(exc)} + attempt["final_status"] = "failed" + attempt["likely_reason"] = "collection_command_failed" + else: + if attempt["return_code"] not in (None, 0): + attempt["final_status"] = "failed" + attempt["likely_reason"] = "collection_command_failed" + else: + partial = _partial_related_path(root, relative_path) + if partial is not None: + attempt["final_status"] = "partial" + attempt["likely_reason"] = "partial_or_temp_file_present" + attempt["partial_path"] = str(partial) + elif destination.is_file() and not _is_partial_name(destination.name): + attempt["final_status"] = "success" + attempt["bytes_copied"] = destination.stat().st_size + attempt["likely_reason"] = "artifact_collected" + else: + attempt["final_status"] = "missing" + attempt["likely_reason"] = "remote_missing_or_not_produced" + attempt["duration_seconds"] = max(0.0, now() - started) + report["attempts"].append(attempt) + _write_report(report_path, _refresh_summary(report, now=now())) + report["batches"][-1]["finished_at"] = now() + _write_report(report_path, _refresh_summary(report, now=now())) + return report + + +def required_collection_failures(report: Mapping[str, Any], *, paths: Iterable[str] | None = None) -> list[dict[str, Any]]: + selected = set(paths) if paths is not None else None + failures: list[dict[str, Any]] = [] + for raw_attempt in report.get("attempts", []): + if not isinstance(raw_attempt, dict): + continue + if selected is not None and raw_attempt.get("expected_path") not in selected: + continue + if raw_attempt.get("required") and raw_attempt.get("final_status") != "success": + failures.append(dict(raw_attempt)) + return failures + + +def _load_report(path: Path) -> dict[str, Any]: + if path.is_file(): + try: + data = json.loads(path.read_text()) + except json.JSONDecodeError: + data = None + if isinstance(data, dict): + data.setdefault("schema_version", 1) + data.setdefault("attempts", []) + data.setdefault("batches", []) + data.setdefault("summary", {}) + return data + return {"schema_version": 1, "attempts": [], "batches": [], "summary": {}} + + +def _refresh_summary(report: dict[str, Any], *, now: float) -> dict[str, Any]: + counts: dict[str, int] = {} + required_missing: list[str] = [] + for raw_attempt in report.get("attempts", []): + if not isinstance(raw_attempt, dict): + continue + status = str(raw_attempt.get("final_status", "unknown")) + counts[status] = counts.get(status, 0) + 1 + if raw_attempt.get("required") and status != "success": + expected = raw_attempt.get("expected_path") + if isinstance(expected, str): + required_missing.append(expected) + report["summary"] = { + "updated_at": now, + "attempt_count": sum(counts.values()), + "status_counts": dict(sorted(counts.items())), + "required_uncollected": required_missing, + "ok": not required_missing, + } + return report + + +def _write_report(path: Path, report: Mapping[str, Any]) -> None: + path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n") + + +def _validate_relative_path(relative_path: str) -> None: + path = Path(relative_path) + if path.is_absolute() or ".." in path.parts: + raise ValueError(f"Artifact path must be relative and stay under artifact root: {relative_path}") + + +def _partial_related_path(local_dir: Path, relative_path: str) -> Path | None: + destination = local_dir / relative_path + if destination.exists() and _is_partial_name(destination.name): + return destination + for suffix in _PARTIAL_SUFFIXES: + candidate = destination.with_name(f"{destination.name}{suffix}") + if candidate.exists(): + return candidate + for temp_dir in _RSYNC_TEMP_DIRS: + candidate = local_dir / temp_dir / relative_path + if candidate.exists(): + return candidate + return None + + +def _is_partial_name(name: str) -> bool: + return name.endswith(_PARTIAL_SUFFIXES) or name in _RSYNC_TEMP_DIRS diff --git a/src/airfrans_frontier/remote/launch_group.py b/src/airfrans_frontier/remote/launch_group.py new file mode 100644 index 0000000..5a3bbcc --- /dev/null +++ b/src/airfrans_frontier/remote/launch_group.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +import json +from pathlib import Path +import time +from typing import Any, Callable, Iterable + +QUEUED_PHASE = "queued" +HEALTHY_PHASE = "training_healthy" +COMPLETED_PHASE = "completed" +FAILED_PHASE = "failed" + +FRAGILE_PHASES = frozenset( + { + "offer_selection", + "provisioning", + "cluster_startup", + "ssh_reachability", + "workdir_sync", + "environment_setup", + "data_validation", + } +) +TERMINAL_PHASES = frozenset({COMPLETED_PHASE, FAILED_PHASE}) +OBSERVABILITY_EVENTS = ( + "run_queued", + "capacity_acquired", + "capacity_blocked", + "offer_selected", + "provisioning_started", + "cluster_reachable", + "rsync_started", + "rsync_completed", + "setup_started", + "setup_completed", + "data_validation_started", + "data_validation_completed", + "training_healthy", + "run_completed", + "run_failed", + "cleanup_started", + "cleanup_completed", + "retry_scheduled", + "retry_exhausted", +) + +_PHASE_EVENTS = { + QUEUED_PHASE: "run_queued", + "offer_selection": "capacity_acquired", + "provisioning": "provisioning_started", + "cluster_startup": "provisioning_started", + "ssh_reachability": "cluster_reachable", + "workdir_sync": "rsync_started", + "environment_setup": "setup_started", + "data_validation": "data_validation_started", + HEALTHY_PHASE: "training_healthy", + COMPLETED_PHASE: "run_completed", + FAILED_PHASE: "run_failed", +} + + +@dataclass(frozen=True) +class LaunchRunSpec: + run_id: str + config_path: str + + +@dataclass +class LaunchRunState: + run_id: str + config_path: str + phase: str = QUEUED_PHASE + selected_offer_id: int | None = None + selected_host_id: int | None = None + retry_count: int = 0 + last_error: str | None = None + cleanup_state: str = "not_started" + blocked_reason: str | None = None + timestamps: dict[str, float] = field(default_factory=dict) + + def to_payload(self) -> dict[str, Any]: + return { + "run_id": self.run_id, + "config_path": self.config_path, + "phase": self.phase, + "selected_offer_id": self.selected_offer_id, + "selected_host_id": self.selected_host_id, + "retry_count": self.retry_count, + "last_error": self.last_error, + "cleanup_state": self.cleanup_state, + "blocked_reason": self.blocked_reason, + "timestamps": dict(sorted(self.timestamps.items())), + } + + +class LaunchGroupScheduler: + """Local launch-group state machine for bounded fragile-phase scheduling. + + The scheduler does not provision machines. Callers drive phase transitions from + observed launch/training evidence and get a durable state artifact after each + transition. + """ + + def __init__( + self, + run_configs: Iterable[str | Path | LaunchRunSpec], + *, + max_active: int, + max_fragile: int, + state_path: str | Path | None = None, + group_id: str | None = None, + allow_duplicate_hosts: bool = False, + clock: Callable[[], float] | None = None, + ) -> None: + if max_active < 1: + raise ValueError("max_active must be >= 1") + if max_fragile < 1: + raise ValueError("max_fragile must be >= 1") + if max_fragile > max_active: + raise ValueError("max_fragile must be <= max_active") + self.clock = clock or time.time + self.group_id = group_id or f"launch-{int(self.clock())}" + self.max_active = int(max_active) + self.max_fragile = int(max_fragile) + self.allow_duplicate_hosts = bool(allow_duplicate_hosts) + self.state_path = Path(state_path) if state_path is not None else None + self.runs: dict[str, LaunchRunState] = {} + self.events: list[dict[str, Any]] = [] + for spec in _coerce_run_specs(run_configs): + now = self.clock() + run = LaunchRunState(run_id=spec.run_id, config_path=spec.config_path) + run.timestamps["queued_at"] = now + self.runs[run.run_id] = run + self._record_event(run.run_id, "run_queued", phase=QUEUED_PHASE, ts=now) + if not self.runs: + raise ValueError("launch group requires at least one run config") + self.write_state() + + def try_start(self, run_id: str, *, selected_offer_id: int | None = None, selected_host_id: int | None = None) -> bool: + run = self._run(run_id) + if run.phase != QUEUED_PHASE: + raise ValueError(f"Run {run_id} is not queued: {run.phase}") + blocker = self._capacity_blocker(selected_host_id=selected_host_id) + if blocker is not None: + run.blocked_reason = blocker + self._record_event(run_id, "capacity_blocked", phase=run.phase, reason=blocker) + self.write_state() + return False + run.phase = "offer_selection" + run.blocked_reason = None + run.selected_offer_id = selected_offer_id + run.selected_host_id = selected_host_id + now = self.clock() + run.timestamps["capacity_acquired_at"] = now + run.timestamps["offer_selection_at"] = now + self._record_event(run_id, "capacity_acquired", phase=run.phase, ts=now) + if selected_offer_id is not None or selected_host_id is not None: + self._record_event( + run_id, + "offer_selected", + phase=run.phase, + selected_offer_id=selected_offer_id, + selected_host_id=selected_host_id, + ) + self.write_state() + return True + + def assign_offer(self, run_id: str, *, selected_offer_id: int, selected_host_id: int | None) -> bool: + run = self._run(run_id) + if run.phase == QUEUED_PHASE: + return self.try_start(run_id, selected_offer_id=selected_offer_id, selected_host_id=selected_host_id) + if run.phase in TERMINAL_PHASES: + raise ValueError(f"Cannot assign offer to terminal run {run_id}: {run.phase}") + if self._host_collision(selected_host_id, excluding_run_id=run_id): + run.blocked_reason = "host_collision" + self._record_event(run_id, "capacity_blocked", phase=run.phase, reason="host_collision", selected_host_id=selected_host_id) + self.write_state() + return False + run.selected_offer_id = int(selected_offer_id) + run.selected_host_id = selected_host_id + run.blocked_reason = None + run.timestamps["offer_selected_at"] = self.clock() + self._record_event( + run_id, + "offer_selected", + phase=run.phase, + selected_offer_id=selected_offer_id, + selected_host_id=selected_host_id, + ) + self.write_state() + return True + + def transition(self, run_id: str, phase: str, *, error: str | None = None, cleanup_state: str | None = None) -> None: + run = self._run(run_id) + run.phase = phase + run.blocked_reason = None + if error is not None: + run.last_error = error + if cleanup_state is not None: + run.cleanup_state = cleanup_state + now = self.clock() + run.timestamps[f"{phase}_at"] = now + self._record_event(run_id, _PHASE_EVENTS.get(phase, phase), phase=phase, ts=now, error=error, cleanup_state=cleanup_state) + self.write_state() + + def mark_training_healthy(self, run_id: str) -> None: + self.transition(run_id, HEALTHY_PHASE) + + def complete_run(self, run_id: str) -> None: + self.transition(run_id, COMPLETED_PHASE) + + def fail_run(self, run_id: str, error: str) -> None: + self.transition(run_id, FAILED_PHASE, error=error) + + def schedule_retry(self, run_id: str, error: str) -> None: + run = self._run(run_id) + run.retry_count += 1 + run.last_error = error + run.phase = QUEUED_PHASE + run.blocked_reason = None + run.timestamps["retry_scheduled_at"] = self.clock() + self._record_event(run_id, "retry_scheduled", phase=run.phase, retry_count=run.retry_count, error=error) + self.write_state() + + def capacity_snapshot(self) -> dict[str, int]: + return { + "active": self._active_count(), + "fragile": self._fragile_count(), + "pending": len(self._runs_in_phase(QUEUED_PHASE)), + "healthy": len(self._runs_in_phase(HEALTHY_PHASE)), + "completed": len(self._runs_in_phase(COMPLETED_PHASE)), + "failed": len(self._runs_in_phase(FAILED_PHASE)), + } + + def to_payload(self) -> dict[str, Any]: + pending = self._runs_in_phase(QUEUED_PHASE) + healthy = self._runs_in_phase(HEALTHY_PHASE) + completed = self._runs_in_phase(COMPLETED_PHASE) + failed = self._runs_in_phase(FAILED_PHASE) + running = [ + run_id + for run_id, run in self.runs.items() + if run.phase not in {QUEUED_PHASE, HEALTHY_PHASE, COMPLETED_PHASE, FAILED_PHASE} + ] + return { + "schema_version": 1, + "launch_group_id": self.group_id, + "requested_run_configs": [run.config_path for run in self.runs.values()], + "limits": { + "max_active": self.max_active, + "max_fragile": self.max_fragile, + "allow_duplicate_hosts": self.allow_duplicate_hosts, + "fragile_phases": sorted(FRAGILE_PHASES), + }, + "pending_runs": pending, + "running_runs": running, + "healthy_runs": healthy, + "completed_runs": completed, + "failed_runs": failed, + "counts": self.capacity_snapshot(), + "phase_counts": self._phase_counts(), + "runs": {run_id: run.to_payload() for run_id, run in self.runs.items()}, + "events": list(self.events), + "updated_at": self.clock(), + } + + def write_state(self) -> Path | None: + if self.state_path is None: + return None + self.state_path.parent.mkdir(parents=True, exist_ok=True) + self.state_path.write_text(json.dumps(self.to_payload(), indent=2, sort_keys=True) + "\n") + return self.state_path + + def _capacity_blocker(self, *, selected_host_id: int | None) -> str | None: + if self._active_count() >= self.max_active: + return "max_active" + if self._fragile_count() >= self.max_fragile: + return "max_fragile" + if self._host_collision(selected_host_id): + return "host_collision" + return None + + def _host_collision(self, selected_host_id: int | None, *, excluding_run_id: str | None = None) -> bool: + if selected_host_id is None or self.allow_duplicate_hosts: + return False + for run_id, run in self.runs.items(): + if run_id == excluding_run_id: + continue + if run.phase not in FRAGILE_PHASES: + continue + if run.selected_host_id == selected_host_id: + return True + return False + + def _active_count(self) -> int: + return sum(1 for run in self.runs.values() if run.phase not in {QUEUED_PHASE, *TERMINAL_PHASES}) + + def _fragile_count(self) -> int: + return sum(1 for run in self.runs.values() if run.phase in FRAGILE_PHASES) + + def _runs_in_phase(self, phase: str) -> list[str]: + return [run_id for run_id, run in self.runs.items() if run.phase == phase] + + def _phase_counts(self) -> dict[str, int]: + counts: dict[str, int] = {} + for run in self.runs.values(): + counts[run.phase] = counts.get(run.phase, 0) + 1 + return dict(sorted(counts.items())) + + def _run(self, run_id: str) -> LaunchRunState: + try: + return self.runs[run_id] + except KeyError as exc: + raise KeyError(f"Unknown launch run id: {run_id}") from exc + + def _record_event(self, run_id: str, event: str, *, phase: str, ts: float | None = None, **fields: Any) -> None: + record = { + "launch_group_id": self.group_id, + "run_id": run_id, + "event": event, + "phase": phase, + "ts": self.clock() if ts is None else ts, + **{key: value for key, value in fields.items() if value is not None}, + } + self.events.append(record) + + +def _coerce_run_specs(run_configs: Iterable[str | Path | LaunchRunSpec]) -> list[LaunchRunSpec]: + result: list[LaunchRunSpec] = [] + seen: set[str] = set() + for index, item in enumerate(run_configs, start=1): + if isinstance(item, LaunchRunSpec): + spec = item + else: + config_path = str(Path(item)) + base = Path(config_path).stem.replace("_", "-") or f"run-{index}" + run_id = base if base not in seen else f"{base}-{index}" + spec = LaunchRunSpec(run_id=run_id, config_path=config_path) + if spec.run_id in seen: + raise ValueError(f"Duplicate run id in launch group: {spec.run_id}") + seen.add(spec.run_id) + result.append(spec) + return result diff --git a/src/airfrans_frontier/remote/selection.py b/src/airfrans_frontier/remote/selection.py new file mode 100644 index 0000000..26e084f --- /dev/null +++ b/src/airfrans_frontier/remote/selection.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +from datetime import UTC, datetime +import json +from pathlib import Path +import time +from typing import Any, Mapping + +DEFAULT_SELECTION_MAX_AGE_SECONDS = 15 * 60 + + +def selection_freshness_report( + manifest: Mapping[str, Any], + *, + max_age_seconds: float = DEFAULT_SELECTION_MAX_AGE_SECONDS, + now: float | None = None, +) -> dict[str, Any]: + """Return freshness metadata for a Vast offer selection artifact.""" + + if max_age_seconds < 0: + raise ValueError("max_age_seconds must be non-negative") + checked_at = time.time() if now is None else float(now) + created_at = _created_at_seconds(manifest) + if created_at is None: + return { + "created_at": None, + "created_at_iso": None, + "checked_at": checked_at, + "age_seconds": None, + "max_age_seconds": float(max_age_seconds), + "is_fresh": False, + "reason": "missing_creation_time", + } + age = max(0.0, checked_at - created_at) + is_fresh = age <= max_age_seconds + return { + "created_at": created_at, + "created_at_iso": datetime.fromtimestamp(created_at, UTC).isoformat(), + "checked_at": checked_at, + "age_seconds": age, + "max_age_seconds": float(max_age_seconds), + "is_fresh": is_fresh, + "reason": "fresh" if is_fresh else "stale", + } + + +def require_fresh_selection( + manifest: Mapping[str, Any], + *, + max_age_seconds: float = DEFAULT_SELECTION_MAX_AGE_SECONDS, + now: float | None = None, + path: str | Path | None = None, +) -> dict[str, Any]: + report = selection_freshness_report(manifest, max_age_seconds=max_age_seconds, now=now) + if not report["is_fresh"]: + location = f" {path}" if path is not None else "" + age = report["age_seconds"] + if age is None: + raise ValueError(f"Selection artifact{location} has no creation time and is stale by policy") + raise ValueError( + f"Selection artifact{location} is stale: age_seconds={age:.3f} " + f"max_age_seconds={float(max_age_seconds):.3f}" + ) + return report + + +def load_selection_manifest( + path: str | Path, + *, + max_age_seconds: float = DEFAULT_SELECTION_MAX_AGE_SECONDS, + allow_stale: bool = False, + now: float | None = None, +) -> dict[str, Any]: + manifest_path = Path(path) + data = json.loads(manifest_path.read_text()) + if not isinstance(data, dict): + raise ValueError(f"Selection manifest is not a JSON object: {manifest_path}") + report = selection_freshness_report(data, max_age_seconds=max_age_seconds, now=now) + data["freshness"] = report + if not allow_stale: + require_fresh_selection(data, max_age_seconds=max_age_seconds, now=now, path=manifest_path) + return data + + +def _created_at_seconds(manifest: Mapping[str, Any]) -> float | None: + for key in ("created_at", "selected_at", "creation_time"): + value = manifest.get(key) + parsed = _parse_timestamp_seconds(value) + if parsed is not None: + return parsed + for key in ("created_at_iso", "selected_at_iso", "creation_time_iso"): + value = manifest.get(key) + parsed = _parse_timestamp_seconds(value) + if parsed is not None: + return parsed + return None + + +def _parse_timestamp_seconds(value: object) -> float | None: + if isinstance(value, bool) or value is None: + return None + if isinstance(value, (int, float)): + return float(value) + if isinstance(value, str): + raw = value.strip() + if not raw: + return None + try: + return float(raw) + except ValueError: + pass + try: + normalized = raw[:-1] + "+00:00" if raw.endswith("Z") else raw + parsed = datetime.fromisoformat(normalized) + except ValueError: + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=UTC) + return parsed.timestamp() + return None diff --git a/src/airfrans_frontier/remote/skypilot.py b/src/airfrans_frontier/remote/skypilot.py index 08824bf..27dc90b 100644 --- a/src/airfrans_frontier/remote/skypilot.py +++ b/src/airfrans_frontier/remote/skypilot.py @@ -32,6 +32,7 @@ def render_skypilot_yaml( env_lines = [ "envs:", f" AIRFRANS_REMOTE_RUN_ID: {run_id}", + f" AIRFRANS_STARTUP_TIMELINE: {_yaml_scalar(str(config.job.artifact_dir / 'startup_timeline.jsonl'))}", ] if resume_checkpoint is not None: env_lines.append(f" AIRFRANS_RESUME_CHECKPOINT: {_yaml_scalar(str(resume_checkpoint))}") @@ -73,9 +74,66 @@ def _compose_setup(config: RemoteRunConfig) -> str: [ "set -euo pipefail", "export PATH=\"$HOME/.local/bin:$PATH\"", + f"mkdir -p {_sh_quote(str(config.job.artifact_dir))}", + _timeline_shell_function(), + _timeline_event("setup", "started"), "if ! command -v uv >/dev/null 2>&1; then curl -LsSf https://astral.sh/uv/install.sh | sh; fi", "export PATH=\"$HOME/.local/bin:$PATH\"", + _timeline_event("disk_preflight", "started"), + _remote_disk_preflight(config), + _timeline_event("disk_preflight", "completed"), + _timeline_event("bootstrap", "started"), config.bootstrap.command.strip(), + _timeline_event("bootstrap", "completed"), + _timeline_event("setup", "completed"), + ] + ) + + +def _remote_disk_preflight(config: RemoteRunConfig) -> str: + requested_gb = config.provider.disk_gb + minimum_total_kib = int(requested_gb * 1024 * 1024 * 0.90) + telemetry_path = config.job.artifact_dir / "disk_telemetry.json" + return "\n".join( + [ + "echo 'airfrans_disk_df_start'", + "df -h .", + f"airfrans_disk_requested_gb={requested_gb}", + "airfrans_disk_total_kib=$(df -Pk . | tail -n 1 | tr -s ' ' | cut -d ' ' -f 2)", + "airfrans_disk_available_kib=$(df -Pk . | tail -n 1 | tr -s ' ' | cut -d ' ' -f 4)", + f"airfrans_disk_minimum_requested_total_kib={minimum_total_kib}", + "echo \"airfrans_disk_requested_gb=${airfrans_disk_requested_gb}\"", + "echo \"airfrans_disk_total_kib=${airfrans_disk_total_kib}\"", + "echo \"airfrans_disk_available_kib=${airfrans_disk_available_kib}\"", + "echo \"airfrans_disk_minimum_requested_total_kib=${airfrans_disk_minimum_requested_total_kib}\"", + "airfrans_disk_capacity_policy=backpressure_adaptive", + f"if [ \"$airfrans_disk_total_kib\" -lt {minimum_total_kib} ]; then", + f" echo \"warning: effective filesystem total ${{airfrans_disk_total_kib}} KiB is below 90% of requested {requested_gb}GB disk; continuing because runtime cache backpressure can adapt\" >&2", + " airfrans_disk_capacity_status=below_requested", + "else", + " airfrans_disk_capacity_status=ok", + "fi", + f"python3 - \"$airfrans_disk_requested_gb\" \"$airfrans_disk_total_kib\" \"$airfrans_disk_available_kib\" \"$airfrans_disk_minimum_requested_total_kib\" \"$airfrans_disk_capacity_status\" {_sh_quote(str(telemetry_path))} <<'PY'", + "import json, os, sys, time", + "requested_gb, total_kib, available_kib, minimum_total_kib, status, path = sys.argv[1:7]", + "payload = {", + " 'schema_version': 1,", + " 'recorded_at': time.time(),", + " 'requested_gb': int(requested_gb),", + " 'total_kib': int(total_kib),", + " 'available_kib': int(available_kib),", + " 'minimum_requested_total_kib': int(minimum_total_kib),", + " 'capacity_status': status,", + " 'capacity_policy': 'backpressure_adaptive',", + " 'hard_failed': False,", + "}", + "directory = os.path.dirname(path)", + "if directory:", + " os.makedirs(directory, exist_ok=True)", + "with open(path, 'w', encoding='utf-8') as handle:", + " json.dump(payload, handle, indent=2, sort_keys=True)", + " handle.write('\\n')", + "PY", ] ) @@ -85,16 +143,63 @@ def _compose_run(config: RemoteRunConfig, *, run_id: str) -> str: "set -euo pipefail", "export PATH=\"$HOME/.local/bin:$PATH\"", f"mkdir -p {_sh_quote(str(config.job.artifact_dir))}", + _timeline_shell_function(), + _timeline_event("run", "started"), + _timeline_event("gpu_probe", "started"), f"nvidia-smi | tee {_sh_quote(str(config.job.artifact_dir / 'nvidia_smi.txt'))}", + _timeline_event("gpu_probe", "completed"), ] if config.data.validation_command: - lines.append(config.data.validation_command.strip()) - lines.append(config.job.command.strip()) - lines.append(f"uv run --no-dev remote-run verify-artifacts {_sh_quote(str(config.job.artifact_dir))}") - lines.append(f"echo 'remote run {run_id} complete'") + lines.extend( + [ + _timeline_event("data_validation", "started"), + config.data.validation_command.strip(), + _timeline_event("data_validation", "completed"), + ] + ) + lines.extend( + [ + _timeline_event("training_command", "started"), + config.job.command.strip(), + _timeline_event("training_command", "completed"), + _timeline_event("artifact_verification", "started"), + f"uv run --no-dev remote-run verify-artifacts {_sh_quote(str(config.job.artifact_dir))}", + _timeline_event("artifact_verification", "completed"), + _timeline_event("run", "completed"), + f"echo 'remote run {run_id} complete'", + ] + ) return "\n".join(lines) +def _timeline_shell_function() -> str: + return "\n".join( + [ + "airfrans_timeline() {", + " python3 - \"$1\" \"$2\" <<'PY'", + "import json, os, sys, time", + "path = os.environ.get('AIRFRANS_STARTUP_TIMELINE', 'artifacts/current_run/startup_timeline.jsonl')", + "record = {", + " 'run_id': os.environ.get('AIRFRANS_REMOTE_RUN_ID'),", + " 'ts': time.time(),", + " 'phase': sys.argv[1],", + " 'event': sys.argv[2],", + "}", + "directory = os.path.dirname(path)", + "if directory:", + " os.makedirs(directory, exist_ok=True)", + "with open(path, 'a', encoding='utf-8') as handle:", + " handle.write(json.dumps(record, sort_keys=True) + '\\n')", + "PY", + "}", + ] + ) + + +def _timeline_event(phase: str, event: str) -> str: + return f"airfrans_timeline {_sh_quote(phase)} {_sh_quote(event)}" + + def _accelerator(config: RemoteRunConfig) -> str: name = config.provider.gpu.name or "T4" aliases = {"Tesla T4": "T4", "RTX 3060 Ti": "RTX3060"} diff --git a/src/airfrans_frontier/remote/smoke.py b/src/airfrans_frontier/remote/smoke.py index 1b9f380..ad9e1d8 100644 --- a/src/airfrans_frontier/remote/smoke.py +++ b/src/airfrans_frontier/remote/smoke.py @@ -84,6 +84,18 @@ def run_smoke_training( if training_dir is not None: _copy_training_artifacts(training_dir, output_dir) + if error is not None and not (output_dir / "failure_report.json").is_file(): + _write_json( + output_dir / "failure_report.json", + { + "run_id": run_id, + "phase": "training", + "error_type": type(error).__name__, + "error_message": str(error), + "training_run_dir": str(training_dir) if training_dir is not None else None, + "timestamp": time.time(), + }, + ) latest_metrics = _read_json(output_dir / "latest_metrics.json") run_manifest: dict[str, Any] = { @@ -144,6 +156,10 @@ def _copy_training_artifacts(training_dir: Path, output_dir: Path) -> None: "artifact_manifest.json", "checksums.txt", "verification_report.json", + "streaming_events.jsonl", + "streaming_state.json", + "streaming_summary.json", + "processed_upload_manifest.json", ) for name in names: source = training_dir / name @@ -168,7 +184,14 @@ def _latest_training_run_dir(config_path: str | Path) -> Path | None: def _smoke_required(*, success: bool) -> tuple[str, ...]: - required = [ + if not success: + return ( + "heartbeat.json", + "environment_manifest.json", + "run_manifest.json", + "failure_report.json", + ) + return ( "config.toml", "metrics.jsonl", "latest_metrics.json", @@ -185,12 +208,9 @@ def _smoke_required(*, success: bool) -> tuple[str, ...]: "run_manifest.json", "artifact_manifest.json", "checksums.txt", - ] - if success: - required.extend(("checkpoint_final.pt", "final_metrics.json")) - else: - required.append("failure_report.json") - return tuple(required) + "checkpoint_final.pt", + "final_metrics.json", + ) def run_hf_upload_smoke( *, diff --git a/src/airfrans_frontier/remote/vast.py b/src/airfrans_frontier/remote/vast.py index 389440d..a390811 100644 --- a/src/airfrans_frontier/remote/vast.py +++ b/src/airfrans_frontier/remote/vast.py @@ -1,12 +1,15 @@ from __future__ import annotations +from datetime import UTC, datetime import json import math import os +import time +import urllib.error import urllib.parse import urllib.request -from dataclasses import asdict, dataclass -from typing import Any, Mapping +from dataclasses import asdict, dataclass, field +from typing import Any, Iterable, Mapping from airfrans_frontier.remote.config import RemoteRunConfig, SelectionConfig @@ -17,6 +20,7 @@ class VastOffer: gpu_name: str dph_total: float gpu_ram: float | None + disk_space: float | None geolocation: str | None inet_down_cost_per_tb: float inet_up_cost_per_tb: float @@ -36,6 +40,7 @@ class VastOffer: gpu_name=_string(data, "gpu_name"), dph_total=_float(data, "dph_total"), gpu_ram=_optional_float(data, "gpu_ram"), + disk_space=_optional_float(data, "disk_space"), geolocation=_optional_string(data, "geolocation"), inet_down_cost_per_tb=_optional_float(data, "internet_down_cost_per_tb") or 0.0, inet_up_cost_per_tb=_optional_float(data, "internet_up_cost_per_tb") or 0.0, @@ -58,6 +63,7 @@ class SelectionResult: effective_price: float query: dict[str, Any] policy: dict[str, Any] + created_at: float = field(default_factory=time.time) @property def selected_offer_id(self) -> int: @@ -66,6 +72,7 @@ class SelectionResult: def to_manifest(self) -> dict[str, Any]: offer = asdict(self.selected_offer) offer["effective_price"] = self.effective_price + now = time.time() return { "selected_offer_id": self.selected_offer_id, "selected_offer": offer, @@ -73,6 +80,9 @@ class SelectionResult: "survivor_count": self.survivor_count, "query": self.query, "policy": self.policy, + "created_at": self.created_at, + "created_at_iso": datetime.fromtimestamp(self.created_at, UTC).isoformat(), + "age_seconds": max(0.0, now - self.created_at), } @@ -111,6 +121,7 @@ def build_query(config: RemoteRunConfig) -> dict[str, Any]: query["verified"] = {"eq": True} if provider.gpu.min_vram_gb is not None: query["gpu_ram"] = {"gte": provider.gpu.min_vram_gb * 1024} + query["disk_space"] = {"gte": provider.disk_gb} if provider.gpu.name: query["gpu_name"] = {"eq": provider.gpu.name} return query @@ -134,22 +145,106 @@ def search_offers(*, base_url: str, api_key: str, query: Mapping[str, Any]) -> l raise RuntimeError("Vast offer search response missing offers list") return [VastOffer.from_mapping(item) for item in raw_offers if isinstance(item, Mapping)] +def list_instances(*, base_url: str, api_key: str) -> list[dict[str, Any]]: + payload = _vast_api_json_request( + base_url=base_url, + api_key=api_key, + path="/api/v0/instances/", + method="GET", + ) + return _instances_from_payload(payload) -def choose_offer(offers: list[VastOffer], config: RemoteRunConfig, *, query: Mapping[str, Any]) -> SelectionResult: + +def destroy_instance(*, base_url: str, api_key: str, instance_id: int) -> Any: + return _vast_api_json_request( + base_url=base_url, + api_key=api_key, + path=f"/api/v0/instances/{int(instance_id)}/", + method="DELETE", + ) + + +def summarize_instances(instances: list[Mapping[str, Any]]) -> list[dict[str, Any]]: + fields = ( + "id", + "instance_id", + "machine_id", + "host_id", + "label", + "status", + "actual_status", + "gpu_name", + "num_gpus", + "dph_total", + "ssh_host", + "ssh_port", + "start_date", + ) + summaries: list[dict[str, Any]] = [] + for instance in instances: + summary = {field: instance[field] for field in fields if field in instance} + summaries.append(summary) + return summaries + + +def _vast_api_json_request(*, base_url: str, api_key: str, path: str, method: str) -> Any: + url = f"{base_url.rstrip('/')}/{path.lstrip('/')}" + request = urllib.request.Request(url, headers={"Authorization": f"Bearer {api_key}"}, method=method) + try: + with urllib.request.urlopen(request, timeout=45) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + raise RuntimeError(f"Vast API {method} {path} HTTP {exc.code}: {body}") from exc + except OSError as exc: + raise RuntimeError(f"Vast API {method} {path} failed: {exc}") from exc + + +def _instances_from_payload(payload: Any) -> list[dict[str, Any]]: + if isinstance(payload, list): + raw_instances = payload + elif isinstance(payload, Mapping): + raw_instances = None + for key in ("instances", "results", "items"): + value = payload.get(key) + if isinstance(value, list): + raw_instances = value + break + if raw_instances is None: + raise RuntimeError("Vast instances response missing instances list") + else: + raise RuntimeError("Vast instances response is not JSON object or list") + return [dict(item) for item in raw_instances if isinstance(item, Mapping)] + + +def choose_offer( + offers: list[VastOffer], + config: RemoteRunConfig, + *, + query: Mapping[str, Any], + reserved_host_ids: Iterable[int] = (), + allow_reserved_hosts: bool = False, +) -> SelectionResult: survivors = reachable_offers(offers, config.selection) ranked = rank_survivors(survivors, config.selection) if config.provider.max_price_per_hour is not None: ranked = [offer for offer in ranked if effective_price(offer, config.selection) <= config.provider.max_price_per_hour] + reserved_hosts = set(reserved_host_ids) + if reserved_hosts and not allow_reserved_hosts: + ranked = [offer for offer in ranked if offer.host_id is None or offer.host_id not in reserved_hosts] if not ranked: - raise RuntimeError("No Vast offers survived quality filters and price cap") + raise RuntimeError("No Vast offers survived quality, price, and host anti-collision filters") selected = ranked[0] + policy = selection_policy_manifest(config) + policy["reserved_host_ids"] = sorted(reserved_hosts) + policy["allow_reserved_hosts"] = bool(allow_reserved_hosts) return SelectionResult( selected_offer=selected, candidate_count=len(offers), survivor_count=len(ranked), effective_price=effective_price(selected, config.selection), query=dict(query), - policy=selection_policy_manifest(config), + policy=policy, ) diff --git a/src/airfrans_frontier/training/config.py b/src/airfrans_frontier/training/config.py index 1a29b9f..1e62915 100644 --- a/src/airfrans_frontier/training/config.py +++ b/src/airfrans_frontier/training/config.py @@ -38,6 +38,14 @@ class DataConfig: hf_repo_type: str hf_path_prefix: str cache_dir: Path | None + public_source_url: str | None = None + streaming_scratch_dir: Path | None = None + streaming_cache_max_bytes: int = 32 * 1024 * 1024 * 1024 + streaming_cache_high_water_bytes: int = 28 * 1024 * 1024 * 1024 + streaming_cache_low_water_bytes: int = 20 * 1024 * 1024 * 1024 + streaming_queue_max_cases: int = 2 + streaming_upload_processed: bool = False + streaming_upload_batch_size: int = 8 @dataclass(frozen=True) @@ -190,6 +198,24 @@ def load_training_config(path: str | Path) -> TrainingConfig: seed=_integer(run_raw, "seed", minimum=0), artifact_dir=_path(run_raw, "artifact_dir"), ) + streaming_cache_max_bytes = _integer(data_raw, "streaming_cache_max_bytes", minimum=1, default=32 * 1024 * 1024 * 1024) + streaming_high_water_bytes = _integer( + data_raw, + "streaming_cache_high_water_bytes", + minimum=1, + default=max(1, streaming_cache_max_bytes * 9 // 10), + ) + streaming_low_water_bytes = _integer( + data_raw, + "streaming_cache_low_water_bytes", + minimum=1, + default=max(1, streaming_cache_max_bytes * 7 // 10), + ) + if streaming_high_water_bytes > streaming_cache_max_bytes: + raise ValueError("data.streaming_cache_high_water_bytes must be <= data.streaming_cache_max_bytes") + if streaming_low_water_bytes >= streaming_high_water_bytes: + raise ValueError("data.streaming_cache_low_water_bytes must be < data.streaming_cache_high_water_bytes") + data = DataConfig( root=_path(data_raw, "root"), train_cases=_integer(data_raw, "train_cases", minimum=1), @@ -197,11 +223,19 @@ def load_training_config(path: str | Path) -> TrainingConfig: test_cases=_integer(data_raw, "test_cases", minimum=0), points_per_case=_integer(data_raw, "points_per_case", minimum=1), batch_size=_integer(data_raw, "batch_size", minimum=1), - source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface"}, "data.source"), + source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface", "public_zip_streaming"}, "data.source"), hf_repo_id=_optional_string(data_raw, "hf_repo_id"), hf_repo_type=_choice(_string(data_raw, "hf_repo_type", default="dataset"), {"dataset"}, "data.hf_repo_type"), hf_path_prefix=_string(data_raw, "hf_path_prefix", default=""), cache_dir=_path(data_raw, "cache_dir") if "cache_dir" in data_raw else None, + public_source_url=_optional_string(data_raw, "public_source_url"), + streaming_scratch_dir=_path(data_raw, "streaming_scratch_dir") if "streaming_scratch_dir" in data_raw else None, + streaming_cache_max_bytes=streaming_cache_max_bytes, + streaming_cache_high_water_bytes=streaming_high_water_bytes, + streaming_cache_low_water_bytes=streaming_low_water_bytes, + streaming_queue_max_cases=_integer(data_raw, "streaming_queue_max_cases", minimum=1, default=2), + streaming_upload_processed=_boolean(data_raw, "streaming_upload_processed") if "streaming_upload_processed" in data_raw else False, + streaming_upload_batch_size=_integer(data_raw, "streaming_upload_batch_size", minimum=1, default=8), ) model = ModelConfig( type=_choice(_string(model_raw, "type"), _MODEL_TYPES, "model.type"), diff --git a/src/airfrans_frontier/training/hf_upload.py b/src/airfrans_frontier/training/hf_upload.py index 77422a2..8690d9f 100644 --- a/src/airfrans_frontier/training/hf_upload.py +++ b/src/airfrans_frontier/training/hf_upload.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import re import json import os import time @@ -30,8 +31,14 @@ class UploadManifest: uploaded_paths: list[str] = field(default_factory=list) uploaded_files: list[dict[str, Any]] = field(default_factory=list) commits: list[dict[str, Any]] = field(default_factory=list) + suppressed_uploads: list[dict[str, Any]] = field(default_factory=list) last_error: str | None = None - + rate_limit_until: float | None = None + rate_limit_retry_after_seconds: float | None = None + training_success: bool | None = None + publication_complete: bool = False + publication_status: str = "disabled" + finalized_at: float | None = None class HfArtifactUploader: def __init__( @@ -43,6 +50,7 @@ class HfArtifactUploader: repo_type: str | None = None, path_in_repo: str | None = None, private: bool = False, + max_rate_limit_sleep_seconds: float = 300.0, ) -> None: self.enabled = enabled self.run_dir = run_dir @@ -50,6 +58,7 @@ class HfArtifactUploader: self.repo_type = repo_type self.path_in_repo = path_in_repo.strip("/") if path_in_repo else None self.private = private + self.max_rate_limit_sleep_seconds = max_rate_limit_sleep_seconds self._api: Any | None = None self._manifest = UploadManifest( enabled=enabled, @@ -77,51 +86,157 @@ class HfArtifactUploader: def repo_url(self) -> str | None: return self._manifest.repo_url + + @property + def publication_status(self) -> str: + self._refresh_publication_status() + return self._manifest.publication_status + + @property + def publication_complete(self) -> bool: + self._refresh_publication_status() + return self._manifest.publication_complete + + def finalize(self, *, training_success: bool) -> dict[str, Any]: + self._manifest.training_success = bool(training_success) + self._manifest.finalized_at = time.time() + self._refresh_publication_status() + self.write_manifest() + return self.final_report() + + def final_report(self) -> dict[str, Any]: + self._refresh_publication_status() + return { + "hf_publication_status": self._manifest.publication_status, + "hf_publication_complete": self._manifest.publication_complete, + "hf_training_success": self._manifest.training_success, + "hf_rate_limit_until": self._manifest.rate_limit_until, + "hf_last_error": self._manifest.last_error, + } def upload_files(self, names: Iterable[str], *, commit_message: str) -> dict[str, Any]: names = tuple(dict.fromkeys(names)) if not self.enabled: - return {"enabled": False, "uploaded": [], "missing": []} + return {"enabled": False, "uploaded": [], "missing": [], "rate_limited": False} missing = [name for name in names if not (self.run_dir / name).is_file()] if missing: raise FileNotFoundError(f"Cannot upload missing Hugging Face artifacts: {', '.join(missing)}") + suppressed = self._suppress_if_rate_limited(names, commit_message=commit_message) + if suppressed is not None: + return suppressed api = self._ensure_api() - uploaded: list[str] = [] - try: - for name in names: - local_path = self.run_dir / name - repo_path = f"{self.path_in_repo}/{name}" if self.path_in_repo else name - commit = api.upload_file( + paths: list[tuple[Path, str]] = [] + for name in names: + local_path = self.run_dir / name + repo_path = f"{self.path_in_repo}/{name}" if self.path_in_repo else name + paths.append((local_path, repo_path)) + uploaded = [repo_path for _, repo_path in paths] + attempts = 0 + while True: + try: + from huggingface_hub import CommitOperationAdd + + operations = [ + CommitOperationAdd(path_in_repo=repo_path, path_or_fileobj=str(local_path)) + for local_path, repo_path in paths + ] + commit = api.create_commit( repo_id=self.repo_id, repo_type=self.repo_type, - path_or_fileobj=str(local_path), - path_in_repo=repo_path, + operations=operations, commit_message=commit_message, ) - uploaded.append(repo_path) - record = UploadRecord( - local_path=str(local_path), - repo_path=repo_path, - bytes=local_path.stat().st_size, - sha256=_sha256_file(local_path), - uploaded_at=time.time(), - ) - self._manifest.uploaded_paths.append(repo_path) - self._manifest.uploaded_files.append(record.__dict__) + uploaded_at = time.time() + for local_path, repo_path in paths: + record = UploadRecord( + local_path=str(local_path), + repo_path=repo_path, + bytes=local_path.stat().st_size, + sha256=_sha256_file(local_path), + uploaded_at=uploaded_at, + ) + self._manifest.uploaded_paths.append(repo_path) + self._manifest.uploaded_files.append(record.__dict__) self._manifest.commits.append(_commit_payload(commit)) - self._manifest.uploaded_paths = sorted(set(self._manifest.uploaded_paths)) - self._manifest.last_error = None - self.write_manifest() - return {"enabled": True, "uploaded": uploaded, "missing": []} - except Exception as exc: - self._manifest.last_error = str(exc) - self.write_manifest() - raise + self._manifest.uploaded_paths = sorted(set(self._manifest.uploaded_paths)) + self._manifest.last_error = None + self._manifest.rate_limit_until = None + self._manifest.rate_limit_retry_after_seconds = None + self.write_manifest() + return {"enabled": True, "uploaded": uploaded, "missing": [], "rate_limited": False} + except Exception as exc: + retry_after = _retry_after_seconds(exc) + if retry_after is not None: + self._record_rate_limit(exc, retry_after, names=names, commit_message=commit_message) + if attempts == 0 and retry_after <= self.max_rate_limit_sleep_seconds: + attempts += 1 + time.sleep(max(0.0, retry_after)) + continue + self._manifest.last_error = str(exc) + self.write_manifest() + raise + + def _suppress_if_rate_limited(self, names: tuple[str, ...], *, commit_message: str) -> dict[str, Any] | None: + until = self._manifest.rate_limit_until + now = time.time() + if until is None or now >= until: + return None + record = { + "names": list(names), + "commit_message": commit_message, + "suppressed_at": now, + "rate_limit_until": until, + } + self._manifest.suppressed_uploads.append(record) + self._manifest.last_error = f"HF upload suppressed until {until:.3f} after rate limiting" + self.write_manifest() + return {"enabled": True, "uploaded": [], "missing": [], "rate_limited": True, "suppressed_until": until} + + def _record_rate_limit(self, exc: Exception, retry_after: float, *, names: tuple[str, ...], commit_message: str) -> None: + now = time.time() + until = now + retry_after + self._manifest.rate_limit_until = max(self._manifest.rate_limit_until or 0.0, until) + self._manifest.rate_limit_retry_after_seconds = retry_after + self._manifest.last_error = str(exc) + self._manifest.suppressed_uploads.append( + { + "names": list(names), + "commit_message": commit_message, + "rate_limited_at": now, + "rate_limit_until": self._manifest.rate_limit_until, + "retry_after_seconds": retry_after, + } + ) + self.write_manifest() def write_manifest(self) -> Path: + self._refresh_publication_status() path = self.run_dir / "hf_upload_manifest.json" path.write_text(json.dumps(self._manifest.__dict__, indent=2, sort_keys=True) + "\n") return path + def _refresh_publication_status(self) -> None: + if not self.enabled: + self._manifest.publication_status = "disabled" + self._manifest.publication_complete = False + return + if self._manifest.training_success is False: + self._manifest.publication_status = "training_failed" + self._manifest.publication_complete = False + return + incomplete = self._manifest.last_error is not None or self._manifest.rate_limit_until is not None + if incomplete: + self._manifest.publication_status = ( + "training_succeeded_hf_incomplete" if self._manifest.training_success is True else "hf_publication_incomplete" + ) + self._manifest.publication_complete = False + return + if self._manifest.training_success is True: + self._manifest.publication_status = "hf_publication_succeeded" + self._manifest.publication_complete = True + return + self._manifest.publication_status = "in_progress" + self._manifest.publication_complete = False + def _ensure_api(self) -> Any: if self._api is not None: return self._api @@ -199,6 +314,33 @@ def _resolve_secret(name: str, purpose: str) -> str: raise RuntimeError(f"{name} env var or local secret file is required for {purpose}") +def _retry_after_seconds(exc: Exception) -> float | None: + response = getattr(exc, "response", None) + headers = getattr(response, "headers", None) + if headers is not None: + raw = headers.get("Retry-After") or headers.get("retry-after") + if raw is not None: + parsed = _parse_retry_after(raw) + if parsed is not None: + return parsed + match = re.search(r"Retry after\s+(\d+(?:\.\d+)?)\s+seconds", str(exc), flags=re.IGNORECASE) + if match: + return float(match.group(1)) + if "rate limit" not in str(exc).lower() and "too many requests" not in str(exc).lower(): + return None + return 300.0 + + +def _parse_retry_after(value: object) -> float | None: + try: + seconds = float(str(value).strip()) + except ValueError: + return None + if seconds < 0: + return None + return seconds + + def _sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as file: diff --git a/src/airfrans_frontier/training/loop.py b/src/airfrans_frontier/training/loop.py index 47c517a..5f0c4de 100644 --- a/src/airfrans_frontier/training/loop.py +++ b/src/airfrans_frontier/training/loop.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import json import os import random import time @@ -32,6 +33,14 @@ from airfrans_frontier.training.hf_upload import HfArtifactUploader, resolve_res from airfrans_frontier.training.observability import start_observer from airfrans_frontier.training.data import DatasetBundle, build_dataset_bundle, load_processed_dataset from airfrans_frontier.training.metrics import count_parameters, device_metrics, overall_mse, per_channel_mse +from airfrans_frontier.training.streaming_data import ( + PROCESSED_UPLOAD_MANIFEST, + STREAMING_EVENTS, + STREAMING_STATE, + STREAMING_SUMMARY, + StreamingEventRecorder, + StreamingTrainingData, +) from airfrans_frontier.training.normalize import ( NormalizationStats, compute_normalization_stats, @@ -87,15 +96,58 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T "hf_path_in_repo": uploader.path_in_repo, } ) + timeline_path = os.environ.get("AIRFRANS_STARTUP_TIMELINE") + first_metric_timeline_written = False + first_checkpoint_timeline_written = False + first_checkpoint_upload_timeline_written = False + + def record_timeline(phase: str, event: str, **extra: Any) -> None: + if not timeline_path: + return + path = Path(timeline_path) + path.parent.mkdir(parents=True, exist_ok=True) + record = { + "run_id": run_id, + "ts": time.time(), + "phase": phase, + "event": event, + **extra, + } + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(record, sort_keys=True) + "\n") + def record_metrics(metrics: dict[str, Any]) -> None: + nonlocal first_metric_timeline_written writer.append_metrics(metrics) observer.log(metrics) + if not first_metric_timeline_written: + record_timeline("training", "first_metric", step=metrics.get("step"), metric_event=metrics.get("event")) + first_metric_timeline_written = True def publish_artifacts(names: tuple[str, ...], *, event: str, step: int) -> None: + nonlocal first_checkpoint_upload_timeline_written if not config.huggingface.enabled: return - upload_result = uploader.upload_files(names, commit_message=f"{run_id}: {event} step {step}") + try: + upload_result = uploader.upload_files(names, commit_message=f"{run_id}: {event} step {step}") + except Exception as exc: + observer.log( + { + "event": "artifact_upload_failed", + "phase": "artifacts", + "step": step, + "artifact_event": event, + "error_type": type(exc).__name__, + "error_message": str(exc), + "hf_repo_url": uploader.repo_url, + "hf_path_in_repo": uploader.path_in_repo, + } + ) + return + if not first_checkpoint_upload_timeline_written and any(name.startswith("checkpoint_") for name in names) and upload_result["uploaded"]: + record_timeline("artifacts", "first_checkpoint_upload", step=step, artifact_event=event) + first_checkpoint_upload_timeline_written = True observer.log( { "event": "artifact_upload", @@ -103,11 +155,29 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T "step": step, "artifact_event": event, "hf_uploaded_count": len(upload_result["uploaded"]), + "hf_rate_limited": bool(upload_result.get("rate_limited", False)), + "hf_suppressed_until": upload_result.get("suppressed_until"), "hf_repo_url": uploader.repo_url, "hf_path_in_repo": uploader.path_in_repo, } ) + if config.data.source == "public_zip_streaming": + return _train_public_zip_streaming( + config=config, + resume=resume, + resume_info=resume_info, + writer=writer, + observer=observer, + uploader=uploader, + run_id=run_id, + run_manifest=run_manifest, + record_metrics=record_metrics, + record_timeline=record_timeline, + publish_artifacts=publish_artifacts, + device=device, + ) + data_root = resolve_training_data_root(config.data) samples = load_processed_dataset(data_root) bundle = build_dataset_bundle( @@ -303,6 +373,9 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T best_val_loss=best_val_loss, initial_train_loss=initial_train_loss, ) + if not first_checkpoint_timeline_written: + record_timeline("training", "first_checkpoint_written", step=start_step) + first_checkpoint_timeline_written = True writer.write_artifact_manifest() publish_artifacts( ( @@ -478,7 +551,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T latest_grad_norm=last_grad_norm, latest_checkpoint=LATEST_CHECKPOINT, ) - run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1}) + run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, **uploader.finalize(training_success=False)}) writer.write_json("run_manifest.json", run_manifest) writer.write_artifact_manifest() try: @@ -499,6 +572,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T "metrics.jsonl", "latest_metrics.json", "heartbeat.json", + "hf_upload_manifest.json", "run_manifest.json", "artifact_manifest.json", "checksums.txt", @@ -650,9 +724,611 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T event="verification", step=config.optim.steps, ) + run_manifest.update(uploader.finalize(training_success=True)) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + verify_artifacts(writer.run_dir, required=_verification_required(success=True)) observer.finish(exit_code=0) return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics) +def _train_public_zip_streaming( + *, + config: TrainingConfig, + resume: Path | None, + resume_info: dict[str, Any], + writer: ArtifactWriter, + observer: Any, + uploader: HfArtifactUploader, + run_id: str, + run_manifest: dict[str, Any], + record_metrics: Any, + record_timeline: Any, + publish_artifacts: Any, + device: torch.device, +) -> TrainingResult: + recorder = StreamingEventRecorder(writer.run_dir) + streaming = StreamingTrainingData.from_config(config, run_dir=writer.run_dir, recorder=recorder) + started = time.perf_counter() + start_step = 0 + best_val_loss: float | None = None + initial_train_loss: float | None = None + last_grad_norm: float | None = None + last_points_per_sec: float | None = None + first_streaming_metric_written = False + first_streaming_checkpoint_written = False + + def record_streaming_metrics(metrics: dict[str, Any]) -> None: + nonlocal first_streaming_metric_written + metrics.update(_streaming_metric_fields(streaming.telemetry_summary())) + record_metrics(metrics) + if not first_streaming_metric_written: + recorder.emit("first_metric", phase="training", step=metrics.get("step"), metric_event=metrics.get("event")) + first_streaming_metric_written = True + + def mark_streaming_checkpoint(step: int, name: str) -> None: + nonlocal first_streaming_checkpoint_written + if not first_streaming_checkpoint_written: + recorder.emit("first_checkpoint_written", phase="training", step=step, checkpoint=name) + record_timeline("training", "first_checkpoint_written", step=step) + first_streaming_checkpoint_written = True + + try: + streaming.prepare() + stats = streaming.load_or_compute_normalization() + bundle = streaming.schema_bundle() + writer.write_split_manifest(bundle.split.to_dict()) + writer.write_json("data_manifest.json", streaming.data_manifest()) + writer.write_normalization(stats.to_dict()) + + model = _build_model(config, bundle, output_dim=bundle.train.targets.shape[1]).to(device) + optimizer = torch.optim.AdamW( + model.parameters(), + lr=config.optim.lr, + weight_decay=config.optim.weight_decay, + ) + + calibration_fields = static_calibration_fields(config, model) + protocol_fields = _evaluation_protocol(config.model.type) + writer.write_json("calibration_manifest.json", calibration_fields) + writer.write_json("evaluation_protocol.json", protocol_fields) + observer.update_config({**calibration_fields, **protocol_fields, **_streaming_metric_fields(streaming.telemetry_summary())}) + run_manifest.update( + { + "phase": "initialized", + "data_mode": "public_zip_streaming", + "parameter_count": count_parameters(model), + **calibration_fields, + **protocol_fields, + **_streaming_metric_fields(streaming.telemetry_summary()), + } + ) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + publish_artifacts( + _existing_artifact_names( + writer.run_dir, + ( + "config.toml", + "environment_manifest.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "calibration_manifest.json", + "evaluation_protocol.json", + "run_manifest.json", + STREAMING_EVENTS, + STREAMING_STATE, + STREAMING_SUMMARY, + PROCESSED_UPLOAD_MANIFEST, + "artifact_manifest.json", + "checksums.txt", + ), + ), + event="initialized", + step=0, + ) + + rng = np.random.default_rng(config.run.seed + 404) + if resume is not None: + try: + checkpoint = _load_checkpoint(resume, device) + _validate_resume_checkpoint(checkpoint, config, bundle, stats) + model.load_state_dict(checkpoint["model_state_dict"]) + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + start_step = int(checkpoint["step"]) + best_val_loss = _optional_float(checkpoint.get("best_val_loss")) + initial_train_loss = _optional_float(checkpoint.get("initial_train_loss")) + _restore_rng_state(checkpoint, rng) + except Exception as exc: + _write_failure( + writer, + phase="resume", + step=0, + error_type=type(exc).__name__, + error_message=str(exc), + latest_checkpoint=str(resume), + ) + raise + record_streaming_metrics( + _log_metrics( + event="resume", + step=start_step, + train_loss=None, + val_loss=best_val_loss, + elapsed_seconds=0.0, + lr=_learning_rate(optimizer), + grad_norm=None, + points_per_sec=None, + device=device, + latest_checkpoint=LATEST_CHECKPOINT, + ) + ) + + initial_train = _evaluate_streaming_split( + model, + streaming, + "train", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + initial_val = ( + _evaluate_streaming_split( + model, + streaming, + "val", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + if bundle.split.val_ids + else None + ) + if initial_train_loss is None: + initial_train_loss = initial_train["loss"] + if best_val_loss is None and initial_val is not None: + best_val_loss = initial_val["loss"] + + record_streaming_metrics( + _log_metrics( + event="initial_eval" if start_step == 0 else "resume_eval", + step=start_step, + train_loss=initial_train["loss"], + val_loss=initial_val["loss"] if initial_val is not None else None, + elapsed_seconds=0.0, + lr=_learning_rate(optimizer), + grad_norm=None, + points_per_sec=None, + device=device, + latest_checkpoint=LATEST_CHECKPOINT, + ) + ) + _save_training_checkpoint( + writer, + LATEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=start_step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + _save_training_checkpoint( + writer, + BEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=start_step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + mark_streaming_checkpoint(start_step, LATEST_CHECKPOINT) + writer.write_artifact_manifest() + publish_artifacts( + _existing_artifact_names( + writer.run_dir, + ( + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + LATEST_CHECKPOINT, + BEST_CHECKPOINT, + STREAMING_EVENTS, + STREAMING_STATE, + STREAMING_SUMMARY, + PROCESSED_UPLOAD_MANIFEST, + "artifact_manifest.json", + "checksums.txt", + ), + ), + event="initial_checkpoint", + step=start_step, + ) + + log_interval = config.optim.log_interval or max(1, config.optim.steps // 10) + last_checkpoint_at = time.monotonic() + last_log_at = time.perf_counter() + last_log_step = start_step + model.train() + for step in range(start_step + 1, config.optim.steps + 1): + batch_features, batch_targets = streaming.sample_train_batch( + rng, + batch_size=config.data.batch_size, + step=step, + ) + features_tensor = _to_device(batch_features, device) + targets_tensor = _to_device(batch_targets, device) + recorder.mark_first_gpu_batch(step=step) + + optimizer.zero_grad(set_to_none=True) + with _autocast_context(config, device): + predictions = model(features_tensor) + loss = F.mse_loss(predictions, targets_tensor) + if not torch.isfinite(loss): + _write_failure( + writer, + phase="training", + step=step, + error_type="NonFiniteLoss", + error_message="loss is NaN or Inf", + latest_loss=float(loss.detach().cpu().item()), + latest_grad_norm=last_grad_norm, + latest_checkpoint=LATEST_CHECKPOINT, + ) + raise RuntimeError("nonfinite loss") + loss.backward() + try: + grad_norm_tensor = torch.nn.utils.clip_grad_norm_( + model.parameters(), + config.stability.max_grad_norm if config.stability.max_grad_norm is not None else float("inf"), + error_if_nonfinite=True, + ) + except RuntimeError as exc: + _write_failure( + writer, + phase="training", + step=step, + error_type="NonFiniteGradient", + error_message=str(exc), + latest_loss=float(loss.detach().cpu().item()), + latest_grad_norm=last_grad_norm, + latest_checkpoint=LATEST_CHECKPOINT, + ) + raise RuntimeError("nonfinite gradients") from exc + last_grad_norm = float(grad_norm_tensor.detach().cpu().item()) + optimizer.step() + + now = time.monotonic() + should_checkpoint = ( + config.checkpoint.interval_seconds == 0 + or now - last_checkpoint_at >= config.checkpoint.interval_seconds + or step == config.optim.steps + ) + if should_checkpoint: + _save_training_checkpoint( + writer, + LATEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + mark_streaming_checkpoint(step, LATEST_CHECKPOINT) + last_checkpoint_at = now + writer.write_artifact_manifest() + publish_artifacts( + _existing_artifact_names( + writer.run_dir, + (LATEST_CHECKPOINT, "metrics.jsonl", "latest_metrics.json", "heartbeat.json", STREAMING_EVENTS, STREAMING_STATE, STREAMING_SUMMARY, PROCESSED_UPLOAD_MANIFEST, "artifact_manifest.json", "checksums.txt"), + ), + event="latest_checkpoint", + step=step, + ) + + if step % log_interval == 0 or step == config.optim.steps: + train_eval = _evaluate_streaming_split( + model, + streaming, + "train", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + val_eval = ( + _evaluate_streaming_split( + model, + streaming, + "val", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + if bundle.split.val_ids + else None + ) + current_metric = val_eval["loss"] if val_eval is not None else train_eval["loss"] + if best_val_loss is None or current_metric < best_val_loss: + best_val_loss = current_metric + _save_training_checkpoint( + writer, + BEST_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=step, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + ) + writer.write_artifact_manifest() + publish_artifacts( + _existing_artifact_names( + writer.run_dir, + (BEST_CHECKPOINT, "metrics.jsonl", "latest_metrics.json", "heartbeat.json", STREAMING_EVENTS, STREAMING_STATE, STREAMING_SUMMARY, PROCESSED_UPLOAD_MANIFEST, "artifact_manifest.json", "checksums.txt"), + ), + event="best_checkpoint", + step=step, + ) + elapsed = time.perf_counter() - started + interval_elapsed = max(time.perf_counter() - last_log_at, 1e-9) + points_per_sec = (step - last_log_step) * config.data.batch_size / interval_elapsed + last_points_per_sec = points_per_sec + record_streaming_metrics( + _log_metrics( + event="train_eval", + step=step, + train_loss=train_eval["loss"], + val_loss=val_eval["loss"] if val_eval is not None else None, + elapsed_seconds=elapsed, + lr=_learning_rate(optimizer), + grad_norm=last_grad_norm, + points_per_sec=points_per_sec, + device=device, + latest_checkpoint=LATEST_CHECKPOINT if should_checkpoint else None, + ) + ) + last_log_at = time.perf_counter() + last_log_step = step + model.train() + + final_train = _evaluate_streaming_split( + model, + streaming, + "train", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + validation_started = time.perf_counter() + final_val = ( + _evaluate_streaming_split( + model, + streaming, + "val", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + if bundle.split.val_ids + else None + ) + validation_runtime_seconds = time.perf_counter() - validation_started if final_val is not None else None + final_test = ( + _evaluate_streaming_split( + model, + streaming, + "test", + batch_size=config.data.batch_size, + device=device, + target_names=bundle.target_names, + ) + if bundle.split.test_ids + else None + ) + streaming.finish(success=True) + streaming_summary = streaming.telemetry_summary() + elapsed = time.perf_counter() - started + final_metrics: dict[str, Any] = { + "initial_train_loss": initial_train_loss, + "train_loss": final_train["loss"], + "train_mse_per_channel": final_train["per_channel_mse"], + "val_loss": final_val["loss"] if final_val is not None else None, + "val_mse_per_channel": final_val["per_channel_mse"] if final_val is not None else None, + "test_loss": final_test["loss"] if final_test is not None else None, + "test_mse_per_channel": final_test["per_channel_mse"] if final_test is not None else None, + "best_val_loss": best_val_loss, + "parameter_count": count_parameters(model), + "model_type": config.model.type, + "model_family": config.model.type, + "precision": config.precision.dtype, + "train_cases": len(bundle.split.train_ids), + "val_cases": len(bundle.split.val_ids), + "test_cases": len(bundle.split.test_ids), + "points_per_case": config.data.points_per_case, + "steps": config.optim.steps, + "elapsed_seconds": elapsed, + "points_per_sec": last_points_per_sec, + "step_time_seconds": (config.data.batch_size / last_points_per_sec) if last_points_per_sec else None, + "validation_runtime_seconds": validation_runtime_seconds, + "checkpoint_interval_seconds": config.checkpoint.interval_seconds, + "data_source": config.data.source, + "data_mode": "public_zip_streaming", + "data_public_source_url": config.data.public_source_url, + "data_hf_repo_id": config.data.hf_repo_id, + "data_hf_path_prefix": config.data.hf_path_prefix, + "data_cache_dir": str(config.data.cache_dir) if config.data.cache_dir is not None else None, + "resumed_from": str(resume) if resume is not None else None, + **resume_info, + **calibration_fields, + **protocol_fields, + **device_metrics(device), + **_streaming_metric_fields(streaming_summary), + } + if observer.url is not None: + final_metrics["wandb_run_url"] = observer.url + if uploader.repo_url is not None: + final_metrics["hf_repo_url"] = uploader.repo_url + final_metrics["hf_path_in_repo"] = uploader.path_in_repo + writer.write_final_metrics(final_metrics) + _save_training_checkpoint( + writer, + FINAL_CHECKPOINT, + config=config, + bundle=bundle, + stats=stats, + model=model, + optimizer=optimizer, + rng=rng, + step=config.optim.steps, + best_val_loss=best_val_loss, + initial_train_loss=initial_train_loss, + final_metrics=final_metrics, + ) + final_metrics.update( + { + "checkpoint_latest_bytes": checkpoint_size_bytes(writer.run_dir, LATEST_CHECKPOINT), + "checkpoint_best_bytes": checkpoint_size_bytes(writer.run_dir, BEST_CHECKPOINT), + "checkpoint_final_bytes": checkpoint_size_bytes(writer.run_dir, FINAL_CHECKPOINT), + } + ) + writer.write_final_metrics(final_metrics) + observer.update_summary(final_metrics) + record_streaming_metrics( + _log_metrics( + event="completed", + phase="completed", + step=config.optim.steps, + train_loss=final_train["loss"], + val_loss=final_val["loss"] if final_val is not None else None, + elapsed_seconds=elapsed, + lr=_learning_rate(optimizer), + grad_norm=last_grad_norm, + points_per_sec=None, + device=device, + latest_checkpoint=FINAL_CHECKPOINT, + ) + ) + run_manifest.update( + { + "phase": "completed", + "finished_at": time.time(), + "exit_code": 0, + "data_mode": "public_zip_streaming", + "final_metrics_path": str(writer.run_dir / "final_metrics.json"), + "checkpoint_latest_path": str(writer.run_dir / LATEST_CHECKPOINT), + "checkpoint_best_path": str(writer.run_dir / BEST_CHECKPOINT), + "checkpoint_final_path": str(writer.run_dir / FINAL_CHECKPOINT), + "wandb_run_url": observer.url, + "hf_repo_url": uploader.repo_url, + "hf_path_in_repo": uploader.path_in_repo, + **calibration_fields, + **protocol_fields, + **_streaming_metric_fields(streaming.telemetry_summary()), + } + ) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + verify_artifacts(writer.run_dir, required=_streaming_verification_required(success=True)) + publish_artifacts(_streaming_final_upload_names(), event="completed", step=config.optim.steps) + writer.write_artifact_manifest() + verify_artifacts(writer.run_dir, required=_streaming_verification_required(success=True)) + publish_artifacts( + ("hf_upload_manifest.json", STREAMING_EVENTS, STREAMING_STATE, STREAMING_SUMMARY, PROCESSED_UPLOAD_MANIFEST, "artifact_manifest.json", "checksums.txt", "verification_report.json"), + event="verification", + step=config.optim.steps, + ) + run_manifest.update(uploader.finalize(training_success=True)) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + verify_artifacts(writer.run_dir, required=_streaming_verification_required(success=True)) + observer.finish(exit_code=0) + return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics) + except Exception as exc: + failure_step = int(locals().get("step", start_step)) + try: + streaming.finish(success=False) + except Exception: + pass + if not (writer.run_dir / "metrics.jsonl").is_file(): + record_streaming_metrics( + _log_metrics( + event="failed", + phase="failed", + step=failure_step, + train_loss=None, + val_loss=None, + elapsed_seconds=time.perf_counter() - started, + lr=0.0, + grad_norm=last_grad_norm, + points_per_sec=None, + device=device, + latest_checkpoint=LATEST_CHECKPOINT if (writer.run_dir / LATEST_CHECKPOINT).is_file() else None, + ) + ) + if not (writer.run_dir / "failure_report.json").is_file(): + _write_failure( + writer, + phase="streaming_training", + step=failure_step, + error_type=type(exc).__name__, + error_message=str(exc), + latest_grad_norm=last_grad_norm, + latest_checkpoint=LATEST_CHECKPOINT if (writer.run_dir / LATEST_CHECKPOINT).is_file() else None, + streaming_summary=streaming.telemetry_summary(), + ) + run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, "data_mode": "public_zip_streaming", **uploader.finalize(training_success=False)}) + writer.write_json("run_manifest.json", run_manifest) + writer.write_artifact_manifest() + try: + verify_artifacts(writer.run_dir, required=_streaming_verification_required(success=False)) + except Exception as verification_exc: + writer.write_json( + "verification_report.json", + { + "ok": False, + "error_type": type(verification_exc).__name__, + "error_message": str(verification_exc), + "checked_at": time.time(), + }, + ) + failure_names = _existing_artifact_names( + writer.run_dir, + ( + "failure_report.json", + "metrics.jsonl", + "latest_metrics.json", + "heartbeat.json", + "run_manifest.json", + STREAMING_EVENTS, + STREAMING_STATE, + STREAMING_SUMMARY, + "hf_upload_manifest.json", + PROCESSED_UPLOAD_MANIFEST, + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", + ), + ) + if failure_names: + publish_artifacts(failure_names, event="failure", step=failure_step) + observer.finish(exit_code=1) + raise + + def _autocast_context(config: TrainingConfig, device: torch.device): if config.precision.dtype == "float32" or device.type != "cuda": return torch.autocast(device_type=device.type, enabled=False) @@ -799,6 +1475,32 @@ def evaluate_arrays( "per_channel_mse": per_channel_mse(squared_error_sum, count, target_names), } +def _evaluate_streaming_split( + model: torch.nn.Module, + streaming: StreamingTrainingData, + split_name: str, + *, + batch_size: int, + device: torch.device, + target_names: tuple[str, ...], +) -> dict[str, Any]: + model.eval() + target_dim = len(target_names) + squared_error_sum = torch.zeros(target_dim, dtype=torch.float64) + count = 0 + with torch.no_grad(): + for batch_features_np, batch_targets_np in streaming.iter_split_batches(split_name, batch_size=batch_size): + batch_features = _to_device(batch_features_np, device) + batch_targets = _to_device(batch_targets_np, device) + predictions = model(batch_features) + errors = predictions - batch_targets + squared_error_sum += (errors.double().pow(2).sum(dim=0)).detach().cpu() + count += int(batch_targets_np.shape[0]) + return { + "loss": overall_mse(squared_error_sum, count, target_dim), + "per_channel_mse": per_channel_mse(squared_error_sum, count, target_names), + } + def _sample_batch( features: np.ndarray, @@ -975,6 +1677,48 @@ def _final_upload_names() -> tuple[str, ...]: "verification_report.json", ) +def _streaming_artifact_names() -> tuple[str, ...]: + return (STREAMING_EVENTS, STREAMING_STATE, STREAMING_SUMMARY, PROCESSED_UPLOAD_MANIFEST) + + +def _streaming_verification_required(*, success: bool) -> tuple[str, ...]: + return _verification_required(success=success) + _streaming_artifact_names() + + +def _streaming_final_upload_names() -> tuple[str, ...]: + return _final_upload_names() + _streaming_artifact_names() + + +def _existing_artifact_names(run_dir: Path, names: tuple[str, ...]) -> tuple[str, ...]: + return tuple(name for name in names if (run_dir / name).is_file()) + + +def _streaming_metric_fields(summary: dict[str, Any]) -> dict[str, Any]: + return { + "streaming_source_bytes": summary.get("source_bytes"), + "streaming_total_downloaded_bytes": summary.get("total_downloaded_bytes"), + "streaming_total_processed_bytes": summary.get("total_processed_bytes"), + "streaming_total_processed_cases": summary.get("total_processed_cases"), + "streaming_processed_cache_high_water_bytes": summary.get("processed_cache_high_water_bytes"), + "streaming_peak_local_disk_usage_bytes": summary.get("peak_local_disk_usage_bytes"), + "streaming_minimum_free_disk_bytes": summary.get("minimum_free_disk_bytes"), + "streaming_producer_idle_backpressure_seconds": summary.get("producer_idle_backpressure_seconds"), + "streaming_trainer_idle_data_starvation_seconds": summary.get("trainer_idle_data_starvation_seconds"), + "streaming_time_to_first_batch_ready_seconds": summary.get("time_to_first_batch_ready_seconds"), + "streaming_time_to_first_gpu_batch_seconds": summary.get("time_to_first_gpu_batch_seconds"), + "streaming_normalization_runtime_seconds": summary.get("normalization_runtime_seconds"), + "streaming_max_inflight_reserved_bytes": summary.get("max_inflight_reserved_bytes"), + "streaming_max_processed_unit_bytes": summary.get("max_processed_unit_bytes"), + "streaming_cache_high_water_events": summary.get("cache_high_water_events"), + "streaming_cache_low_water_events": summary.get("cache_low_water_events"), + "streaming_producer_pause_events": summary.get("producer_pause_events"), + "streaming_producer_resume_events": summary.get("producer_resume_events"), + "streaming_evicted_units": summary.get("evicted_units"), + "streaming_upload_queue_depth": summary.get("upload_queue_depth"), + "streaming_upload_lag_seconds": summary.get("upload_lag_seconds"), + "streaming_upload_suppressed_until": summary.get("upload_suppressed_until"), + } + def _memory_metrics(device: torch.device) -> dict[str, int | None]: if device.type != "cuda": @@ -1140,6 +1884,7 @@ def _write_failure( latest_loss: float | None = None, latest_grad_norm: float | None = None, latest_checkpoint: str | None = None, + **extra: Any, ) -> None: writer.write_failure_report( { @@ -1153,6 +1898,7 @@ def _write_failure( "latest_grad_norm": latest_grad_norm, "latest_checkpoint": latest_checkpoint, "timestamp": time.time(), + **extra, } ) writer.write_artifact_manifest() diff --git a/src/airfrans_frontier/training/streaming_data.py b/src/airfrans_frontier/training/streaming_data.py new file mode 100644 index 0000000..7fa8705 --- /dev/null +++ b/src/airfrans_frontier/training/streaming_data.py @@ -0,0 +1,1132 @@ +from __future__ import annotations + +import json +import os +import shutil +import time +from dataclasses import dataclass, field +from pathlib import Path, PurePosixPath +from typing import Any, Iterable, Iterator, Mapping + +import numpy as np +from numpy.typing import NDArray + +from airfrans_frontier.raw.public import ( + PUBLIC_OF_DATASET_URL, + RangeReader, + RemoteZipMember, + _extract_remote_case_members, + _range_reader_for, + _read_zip_central_directory, + _remote_archive_case_members, +) +from airfrans_frontier.raw.process import process_raw_case_to_npz +from airfrans_frontier.training.config import DataConfig, TrainingConfig +from airfrans_frontier.training.data import CaseSplit, DatasetBundle, SimulationSample, SplitArrays, create_case_split, load_simulation_npz +from airfrans_frontier.training.hf_upload import _retry_after_seconds +from airfrans_frontier.training.normalize import NormalizationStats, load_normalization_stats + +FloatArray = NDArray[np.float32] +IntArray = NDArray[np.int64] + +STREAMING_EVENTS = "streaming_events.jsonl" +STREAMING_STATE = "streaming_state.json" +STREAMING_SUMMARY = "streaming_summary.json" +PROCESSED_UPLOAD_MANIFEST = "processed_upload_manifest.json" +_SAMPLE_SEEDS = {"train": 101, "val": 202, "test": 303} + + +@dataclass(frozen=True) +class StreamingCaseMember: + member: RemoteZipMember + relative: PurePosixPath + + +@dataclass(frozen=True) +class SamplingSpec: + case_id: str + split_name: str + count: int + mode: str + indices_path: Path | None = None + + +@dataclass +class StreamingSummary: + source_bytes: int = 0 + total_downloaded_bytes: int = 0 + total_processed_bytes: int = 0 + total_processed_cases: int = 0 + processed_cache_high_water_bytes: int = 0 + peak_local_disk_usage_bytes: int = 0 + minimum_free_disk_bytes: int | None = None + producer_idle_backpressure_seconds: float = 0.0 + trainer_idle_data_starvation_seconds: float = 0.0 + time_to_first_batch_ready_seconds: float | None = None + time_to_first_gpu_batch_seconds: float | None = None + normalization_runtime_seconds: float | None = None + max_inflight_reserved_bytes: int = 0 + max_processed_unit_bytes: int = 0 + cache_high_water_events: int = 0 + cache_low_water_events: int = 0 + producer_pause_events: int = 0 + producer_resume_events: int = 0 + evicted_units: int = 0 + upload_queue_depth: int = 0 + upload_lag_seconds: float | None = None + upload_suppressed_until: float | None = None + + +class StreamingEventRecorder: + def __init__(self, run_dir: Path) -> None: + self.run_dir = run_dir + self.events_path = run_dir / STREAMING_EVENTS + self.summary_path = run_dir / STREAMING_SUMMARY + self.started_at = time.time() + self.started_perf = time.perf_counter() + self.summary = StreamingSummary() + run_dir.mkdir(parents=True, exist_ok=True) + if not self.events_path.exists(): + self.events_path.write_text("") + self.write_summary() + + def emit(self, event: str, *, phase: str, **fields: Any) -> None: + now = time.time() + record = { + "event": event, + "phase": phase, + "timestamp": now, + "elapsed_seconds": time.perf_counter() - self.started_perf, + **fields, + } + _append_jsonl(self.events_path, record) + self._apply_event_to_summary(event, fields) + self.write_summary() + + def observe_cache(self, cache_dir: Path, *, cache_bytes: int) -> None: + usage = shutil.disk_usage(cache_dir) + used = int(usage.used) + free = int(usage.free) + self.summary.peak_local_disk_usage_bytes = max(self.summary.peak_local_disk_usage_bytes, used) + if self.summary.minimum_free_disk_bytes is None: + self.summary.minimum_free_disk_bytes = free + else: + self.summary.minimum_free_disk_bytes = min(self.summary.minimum_free_disk_bytes, free) + self.summary.processed_cache_high_water_bytes = max(self.summary.processed_cache_high_water_bytes, cache_bytes) + self.write_summary() + + def add_downloaded_bytes(self, delta: int) -> None: + if delta > 0: + self.summary.total_downloaded_bytes += int(delta) + self.write_summary() + + def mark_first_batch_ready(self) -> None: + if self.summary.time_to_first_batch_ready_seconds is None: + elapsed = time.perf_counter() - self.started_perf + self.summary.time_to_first_batch_ready_seconds = elapsed + self.emit("first_batch_ready", phase="training", elapsed_to_first_batch_seconds=elapsed) + + def mark_first_gpu_batch(self, *, step: int) -> None: + if self.summary.time_to_first_gpu_batch_seconds is None: + elapsed = time.perf_counter() - self.started_perf + self.summary.time_to_first_gpu_batch_seconds = elapsed + self.emit("first_gpu_batch_consumed", phase="training", step=step, elapsed_to_first_gpu_batch_seconds=elapsed) + + def add_trainer_wait(self, seconds: float, *, step: int) -> None: + if seconds <= 0: + return + self.summary.trainer_idle_data_starvation_seconds += float(seconds) + self.emit("data_loader_starvation", phase="training", step=step, wait_seconds=float(seconds)) + + def set_upload_state(self, *, queue_depth: int, lag_seconds: float | None, suppressed_until: float | None) -> None: + self.summary.upload_queue_depth = int(queue_depth) + self.summary.upload_lag_seconds = lag_seconds + self.summary.upload_suppressed_until = suppressed_until + self.write_summary() + + def write_summary(self) -> None: + _atomic_write_json(self.summary_path, self.to_dict()) + + def to_dict(self) -> dict[str, Any]: + return dict(self.summary.__dict__) + + def _apply_event_to_summary(self, event: str, fields: Mapping[str, Any]) -> None: + if event == "cache_high_water": + self.summary.cache_high_water_events += 1 + elif event == "cache_low_water": + self.summary.cache_low_water_events += 1 + elif event == "producer_paused": + self.summary.producer_pause_events += 1 + elif event == "producer_resumed": + self.summary.producer_resume_events += 1 + idle = fields.get("idle_seconds") + if idle is not None: + self.summary.producer_idle_backpressure_seconds += float(idle) + elif event == "cleanup_eviction": + self.summary.evicted_units += 1 + elif event == "processing_end": + processed_bytes = int(fields.get("processed_bytes", 0) or 0) + self.summary.total_processed_cases += 1 + self.summary.total_processed_bytes += processed_bytes + self.summary.max_processed_unit_bytes = max(self.summary.max_processed_unit_bytes, processed_bytes) + elif event == "inflight_reserved": + reserved = int(fields.get("inflight_reserved_bytes", 0) or 0) + self.summary.max_inflight_reserved_bytes = max(self.summary.max_inflight_reserved_bytes, reserved) + elif event == "normalization_end": + runtime = fields.get("normalization_runtime_seconds") + if runtime is not None: + self.summary.normalization_runtime_seconds = float(runtime) + + +class PublicZipStreamingCache: + def __init__(self, data_config: DataConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> None: + self.data_config = data_config + self.run_dir = run_dir + self.recorder = recorder + self.source_url = data_config.public_source_url or PUBLIC_OF_DATASET_URL + self.cache_dir = data_config.cache_dir or data_config.root + self.scratch_dir = data_config.streaming_scratch_dir or (self.cache_dir / "_streaming_raw") + self.state_path = run_dir / STREAMING_STATE + self.indices_dir = run_dir / "streaming_sample_indices" + self.cache_dir.mkdir(parents=True, exist_ok=True) + self.scratch_dir.mkdir(parents=True, exist_ok=True) + self.indices_dir.mkdir(parents=True, exist_ok=True) + self.reader: RangeReader | None = None + self.members: list[RemoteZipMember] = [] + self.case_members: dict[str, list[StreamingCaseMember]] = {} + self.state: dict[str, Any] = _read_json(self.state_path) if self.state_path.is_file() else self._empty_state() + self._active_cases: set[str] = set() + self._last_access: dict[str, float] = {} + self._loaded: dict[str, SimulationSample] = {} + self._cache_pressure_active = False + + def enumerate_cases(self) -> tuple[str, ...]: + self.recorder.emit("dataset_enumeration_start", phase="data", source_url=self.source_url) + before = 0 + self.reader = _range_reader_for(self.source_url) + before = self.reader.bytes_read + self.members = _read_zip_central_directory(self.reader) + self.recorder.add_downloaded_bytes(self.reader.bytes_read - before) + raw_cases = _remote_archive_case_members(self.members) + self.case_members = { + case_id: [StreamingCaseMember(member=member, relative=relative) for member, relative in entries] + for case_id, entries in raw_cases.items() + } + case_ids = tuple(sorted(self.case_members)) + self.recorder.summary.source_bytes = int(self.reader.size) + self.state["source_url"] = self.source_url + self.state["source_bytes"] = int(self.reader.size) + self.state["enumerated_case_count"] = len(case_ids) + self._write_state() + self.recorder.emit( + "dataset_enumeration_end", + phase="data", + source_url=self.source_url, + source_bytes=int(self.reader.size), + zip_members=len(self.members), + case_count=len(case_ids), + ranged_bytes_read=int(self.reader.bytes_read), + ) + return case_ids + + def ensure_case(self, case_id: str) -> SimulationSample: + self._discard_partial(case_id) + cached = self._load_valid_cached_case(case_id) + if cached is not None: + self._active_cases.add(case_id) + self._last_access[case_id] = time.time() + return cached + members = self.case_members.get(case_id) + if not members: + raise KeyError(f"Unknown streaming AirfRANS case: {case_id}") + assert self.reader is not None + reserved = self._case_reserved_bytes(members) + self._reserve_inflight(case_id, reserved) + self._mark_case(case_id, status="processing", inflight_reserved_bytes=reserved, processing_started_at=time.time()) + case_dir = self.scratch_dir / case_id + fetch_started = time.perf_counter() + before = self.reader.bytes_read + self.recorder.emit("range_fetch_start", phase="data", case_id=case_id, reserved_bytes=reserved, member_count=len(members)) + try: + _extract_remote_case_members(self.reader, [(entry.member, entry.relative) for entry in members], self.scratch_dir) + except Exception as exc: + self._mark_case(case_id, status="failed", error_type=type(exc).__name__, error_message=str(exc), failed_at=time.time()) + self.recorder.emit("range_fetch_failure", phase="data", case_id=case_id, error_type=type(exc).__name__, error_message=str(exc)) + raise + finally: + self.recorder.add_downloaded_bytes(self.reader.bytes_read - before) + self.recorder.emit( + "range_fetch_end", + phase="data", + case_id=case_id, + downloaded_bytes=self.reader.bytes_read - before, + ranged_bytes_read=int(self.reader.bytes_read), + fetch_seconds=time.perf_counter() - fetch_started, + ) + process_started = time.perf_counter() + self.recorder.emit("processing_start", phase="data", case_id=case_id) + try: + record, points = process_raw_case_to_npz(case_dir, self.cache_dir, force=True) + sample = load_simulation_npz(self.cache_dir / f"{case_id}.npz") + except Exception as exc: + self._mark_case(case_id, status="failed", error_type=type(exc).__name__, error_message=str(exc), failed_at=time.time()) + self.recorder.emit("processing_failure", phase="data", case_id=case_id, error_type=type(exc).__name__, error_message=str(exc)) + raise + finally: + if case_dir.exists(): + shutil.rmtree(case_dir, ignore_errors=True) + processed_bytes = sample.source_path.stat().st_size + self._loaded[case_id] = sample + self._active_cases.add(case_id) + self._last_access[case_id] = time.time() + self._mark_case( + case_id, + status="processed", + path=str(sample.source_path), + processed_bytes=processed_bytes, + points=int(points), + processed_at=time.time(), + processing_seconds=time.perf_counter() - process_started, + metadata=record.get("metadata"), + ) + self.recorder.emit( + "processing_end", + phase="data", + case_id=case_id, + processed_bytes=processed_bytes, + points=int(points), + processing_seconds=time.perf_counter() - process_started, + ) + self._observe_cache() + return sample + + def release_case(self, case_id: str, *, consumed: bool = True) -> None: + self._active_cases.discard(case_id) + self._loaded.pop(case_id, None) + self._last_access[case_id] = time.time() + entry = self.state.setdefault("cases", {}).setdefault(case_id, {}) + if consumed: + entry["consumed_count"] = int(entry.get("consumed_count", 0) or 0) + 1 + entry["last_consumed_at"] = time.time() + self._write_state() + self._evict_if_needed() + + def cache_size_bytes(self) -> int: + return _tree_size_bytes(self.cache_dir, suffix=".npz") + + def data_manifest(self, *, split: CaseSplit, feature_names: tuple[str, ...], target_names: tuple[str, ...]) -> dict[str, Any]: + selected = set(split.train_ids) | set(split.val_ids) | set(split.test_ids) + cases_state = self.state.get("cases", {}) if isinstance(self.state.get("cases"), dict) else {} + return { + "source": "public_zip_streaming", + "source_url": self.source_url, + "source_bytes": self.state.get("source_bytes"), + "configured_root": str(self.data_config.root), + "cache_dir": str(self.cache_dir), + "scratch_dir": str(self.scratch_dir), + "case_count": self.state.get("enumerated_case_count"), + "selected_case_count": len(selected), + "train_cases": len(split.train_ids), + "val_cases": len(split.val_ids), + "test_cases": len(split.test_ids), + "feature_names": list(feature_names), + "target_names": list(target_names), + "cache_max_bytes": self.data_config.streaming_cache_max_bytes, + "cache_high_water_bytes": self.data_config.streaming_cache_high_water_bytes, + "cache_low_water_bytes": self.data_config.streaming_cache_low_water_bytes, + "queue_max_cases": self.data_config.streaming_queue_max_cases, + "ranged_bytes_read": self.reader.bytes_read if self.reader is not None else None, + "cases": [ + { + "case_id": case_id, + "split": _split_name_for_case(split, case_id), + "status": cases_state.get(case_id, {}).get("status", "pending") if isinstance(cases_state.get(case_id), dict) else "pending", + "points": cases_state.get(case_id, {}).get("points") if isinstance(cases_state.get(case_id), dict) else None, + "processed_bytes": cases_state.get(case_id, {}).get("processed_bytes") if isinstance(cases_state.get(case_id), dict) else None, + "member_count": len(self.case_members.get(case_id, ())), + "compressed_bytes": sum(entry.member.compress_size for entry in self.case_members.get(case_id, ())), + "uncompressed_bytes": sum(entry.member.file_size for entry in self.case_members.get(case_id, ())), + } + for case_id in sorted(selected) + ], + } + + def _load_valid_cached_case(self, case_id: str) -> SimulationSample | None: + if case_id in self._loaded: + return self._loaded[case_id] + path = self.cache_dir / f"{case_id}.npz" + if not path.is_file(): + return None + try: + sample = load_simulation_npz(path) + except Exception: + path.unlink(missing_ok=True) + self.recorder.emit("partial_unit_discarded", phase="data", case_id=case_id, path=str(path)) + self._mark_case(case_id, status="pending", discarded_invalid_at=time.time()) + return None + self._loaded[case_id] = sample + self._last_access[case_id] = time.time() + entry = self.state.setdefault("cases", {}).setdefault(case_id, {}) + if entry.get("status") == "processed": + self.recorder.emit("resume_validated_unit_reused", phase="data", case_id=case_id, path=str(path), processed_bytes=path.stat().st_size) + self._mark_case(case_id, status="processed", path=str(path), processed_bytes=path.stat().st_size, points=sample.num_points) + self._observe_cache() + return sample + + def _reserve_inflight(self, case_id: str, reserved: int) -> None: + cache_bytes = self.cache_size_bytes() + self.recorder.emit("inflight_reserved", phase="data", case_id=case_id, inflight_reserved_bytes=reserved, cache_bytes=cache_bytes) + self._observe_cache(cache_bytes=cache_bytes) + if cache_bytes + reserved <= self.data_config.streaming_cache_high_water_bytes: + return + pause_started = time.perf_counter() + self._cache_pressure_active = True + self.recorder.emit( + "cache_high_water", + phase="data", + case_id=case_id, + cache_bytes=cache_bytes, + inflight_reserved_bytes=reserved, + high_water_bytes=self.data_config.streaming_cache_high_water_bytes, + allowed_inflight_slack_bytes=reserved, + ) + self.recorder.emit( + "producer_paused", + phase="data", + case_id=case_id, + reason="cache_high_water", + cache_bytes=cache_bytes, + inflight_reserved_bytes=reserved, + ) + self._evict_to_target(self.data_config.streaming_cache_low_water_bytes, protected={case_id}) + resumed_cache = self.cache_size_bytes() + idle_seconds = time.perf_counter() - pause_started + self.recorder.emit( + "cache_low_water", + phase="data", + case_id=case_id, + cache_bytes=resumed_cache, + low_water_bytes=self.data_config.streaming_cache_low_water_bytes, + ) + self.recorder.emit( + "producer_resumed", + phase="data", + case_id=case_id, + reason="cache_low_water", + cache_bytes=resumed_cache, + idle_seconds=idle_seconds, + ) + self._cache_pressure_active = False + + def _evict_if_needed(self) -> None: + cache_bytes = self.cache_size_bytes() + self._observe_cache(cache_bytes=cache_bytes) + if cache_bytes <= self.data_config.streaming_cache_high_water_bytes: + return + self.recorder.emit( + "cache_high_water", + phase="data", + cache_bytes=cache_bytes, + high_water_bytes=self.data_config.streaming_cache_high_water_bytes, + ) + pause_started = time.perf_counter() + self.recorder.emit("producer_paused", phase="data", reason="cache_high_water", cache_bytes=cache_bytes) + self._evict_to_target(self.data_config.streaming_cache_low_water_bytes, protected=set()) + resumed_cache = self.cache_size_bytes() + self.recorder.emit( + "cache_low_water", + phase="data", + cache_bytes=resumed_cache, + low_water_bytes=self.data_config.streaming_cache_low_water_bytes, + ) + self.recorder.emit( + "producer_resumed", + phase="data", + reason="cache_low_water", + cache_bytes=resumed_cache, + idle_seconds=time.perf_counter() - pause_started, + ) + + def _evict_to_target(self, target_bytes: int, *, protected: set[str]) -> None: + candidates: list[tuple[float, str, Path]] = [] + for path in self.cache_dir.glob("*.npz"): + case_id = path.stem + if case_id in protected or case_id in self._active_cases: + continue + candidates.append((self._last_access.get(case_id, 0.0), case_id, path)) + for _last_access, case_id, path in sorted(candidates): + if self.cache_size_bytes() <= target_bytes: + break + bytes_before = path.stat().st_size if path.exists() else 0 + path.unlink(missing_ok=True) + self._loaded.pop(case_id, None) + self._mark_case(case_id, status="consumed", evicted_at=time.time(), evicted_bytes=bytes_before) + self.recorder.emit("cleanup_eviction", phase="data", case_id=case_id, bytes=bytes_before, cache_bytes=self.cache_size_bytes()) + self._observe_cache() + + def _discard_partial(self, case_id: str) -> None: + partials = list(self.cache_dir.glob(f"{case_id}.npz.tmp*")) + list(self.cache_dir.glob(f"{case_id}.tmp*")) + for path in partials: + path.unlink(missing_ok=True) + self.recorder.emit("partial_unit_discarded", phase="data", case_id=case_id, path=str(path)) + + def _case_reserved_bytes(self, members: Iterable[StreamingCaseMember]) -> int: + total = 0 + for entry in members: + if entry.member.is_dir: + continue + total += 30 + len(entry.member.filename.encode()) + entry.member.compress_size + return int(total) + + def _observe_cache(self, *, cache_bytes: int | None = None) -> None: + self.recorder.observe_cache(self.cache_dir, cache_bytes=self.cache_size_bytes() if cache_bytes is None else cache_bytes) + + def _mark_case(self, case_id: str, **fields: Any) -> None: + cases = self.state.setdefault("cases", {}) + entry = cases.setdefault(case_id, {}) + entry.update(fields) + entry["updated_at"] = time.time() + self._write_state() + + def _write_state(self) -> None: + self.state["schema_version"] = 1 + self.state["cache_dir"] = str(self.cache_dir) + self.state["scratch_dir"] = str(self.scratch_dir) + self.state["updated_at"] = time.time() + _atomic_write_json(self.state_path, self.state) + + @staticmethod + def _empty_state() -> dict[str, Any]: + return {"schema_version": 1, "cases": {}, "sampling": {"train": {}, "val": {}, "test": {}}} + + +class StreamingTrainingData: + def __init__(self, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> None: + self.config = config + self.recorder = recorder + self.cache = PublicZipStreamingCache(config.data, run_dir=run_dir, recorder=recorder) + self.split: CaseSplit | None = None + self.feature_names: tuple[str, ...] | None = None + self.target_names: tuple[str, ...] | None = None + self.stats: NormalizationStats | None = None + self._split_case_ids: dict[str, tuple[str, ...]] = {} + self._sampling_specs: dict[str, dict[str, SamplingSpec]] = {"train": {}, "val": {}, "test": {}} + self._train_counts: list[int] = [] + self._train_offsets: NDArray[np.int64] | None = None + self._upload_queue = ProcessedDataUploadQueue.from_config(config, run_dir=run_dir, recorder=recorder) + + @classmethod + def from_config(cls, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> StreamingTrainingData: + if config.data.source != "public_zip_streaming": + raise ValueError(f"StreamingTrainingData requires data.source = 'public_zip_streaming', got {config.data.source!r}") + return cls(config, run_dir=run_dir, recorder=recorder) + + def prepare(self) -> None: + case_ids = self.cache.enumerate_cases() + split = create_case_split( + case_ids, + train_cases=self.config.data.train_cases, + val_cases=self.config.data.val_cases, + test_cases=self.config.data.test_cases, + seed=self.config.run.seed, + ) + self.split = split + self._split_case_ids = {"train": split.train_ids, "val": split.val_ids, "test": split.test_ids} + self.cache.state["split"] = split.to_dict() + self.cache._write_state() + self.recorder.emit( + "split_selection", + phase="data", + train_cases=len(split.train_ids), + val_cases=len(split.val_ids), + test_cases=len(split.test_ids), + seed=self.config.run.seed, + ) + first_case = split.train_ids[0] + sample = self.cache.ensure_case(first_case) + self.feature_names = sample.feature_names + self.target_names = sample.target_names + self._upload_queue.enqueue(sample.source_path) + self.cache.release_case(first_case, consumed=False) + + def load_or_compute_normalization(self) -> NormalizationStats: + if self.split is None: + raise RuntimeError("Streaming data must be prepared before normalization") + existing = self.cache.run_dir / "normalization.json" + if existing.is_file(): + stats = load_normalization_stats(existing) + self.stats = stats + self._restore_sampling_specs("train") + self._build_train_offsets() + self.recorder.emit("normalization_end", phase="normalization", reused=True, normalization_runtime_seconds=0.0) + return stats + assert self.feature_names is not None + assert self.target_names is not None + self.recorder.emit("normalization_start", phase="normalization", train_cases=len(self.split.train_ids)) + started = time.perf_counter() + accumulator = _StatsAccumulator(feature_names=self.feature_names, target_names=self.target_names) + rng = np.random.default_rng(self.config.run.seed + _SAMPLE_SEEDS["train"]) + for case_id in self.split.train_ids: + sample = self.cache.ensure_case(case_id) + self._validate_schema(sample) + indices = self._sampling_spec_for_case("train", sample, rng) + features = _selected_rows(sample.features, indices) + targets = _selected_rows(sample.targets, indices) + accumulator.update(features, targets) + self._upload_queue.enqueue(sample.source_path) + self.cache.release_case(case_id) + self._upload_queue.drain() + stats = accumulator.finish() + self.stats = stats + self._build_train_offsets() + runtime = time.perf_counter() - started + self.recorder.emit( + "normalization_end", + phase="normalization", + train_cases=len(self.split.train_ids), + sample_count=accumulator.count, + normalization_runtime_seconds=runtime, + ) + return stats + + def schema_bundle(self) -> DatasetBundle: + if self.split is None or self.feature_names is None or self.target_names is None: + raise RuntimeError("Streaming schema is not initialized") + train = SplitArrays( + features=np.zeros((1, len(self.feature_names)), dtype=np.float32), + targets=np.zeros((1, len(self.target_names)), dtype=np.float32), + case_ids=self.split.train_ids, + ) + val = SplitArrays(train.features, train.targets, self.split.val_ids) if self.split.val_ids else None + test = SplitArrays(train.features, train.targets, self.split.test_ids) if self.split.test_ids else None + return DatasetBundle(train=train, val=val, test=test, split=self.split, feature_names=self.feature_names, target_names=self.target_names) + + def data_manifest(self) -> dict[str, Any]: + if self.split is None or self.feature_names is None or self.target_names is None: + raise RuntimeError("Streaming data is not ready") + payload = self.cache.data_manifest(split=self.split, feature_names=self.feature_names, target_names=self.target_names) + payload["processed_upload"] = self._upload_queue.manifest_payload() + return payload + + def telemetry_summary(self) -> dict[str, Any]: + self.recorder.set_upload_state(**self._upload_queue.telemetry_state()) + return self.recorder.to_dict() + + def sample_train_batch(self, rng: np.random.Generator, *, batch_size: int, step: int) -> tuple[FloatArray, FloatArray]: + if self.stats is None or self.split is None or self._train_offsets is None: + raise RuntimeError("Streaming normalization must be computed before sampling") + started = time.perf_counter() + total = int(self._train_offsets[-1]) if self._train_offsets.size else 0 + if total <= 0: + raise ValueError("Streaming train split has no sampled points") + global_indices = rng.integers(0, total, size=batch_size) + case_positions = np.searchsorted(self._train_offsets[1:], global_indices, side="right") + features = np.empty((batch_size, len(self.stats.feature_names)), dtype=np.float32) + targets = np.empty((batch_size, len(self.stats.target_names)), dtype=np.float32) + for case_position in np.unique(case_positions): + mask = case_positions == case_position + case_id = self.split.train_ids[int(case_position)] + sample = self.cache.ensure_case(case_id) + self._validate_schema(sample) + spec = self._sampling_specs["train"][case_id] + local_indices = global_indices[mask] - self._train_offsets[int(case_position)] + source_indices = _source_indices_for_local(spec, local_indices) + selected_features = sample.features[source_indices] + selected_targets = sample.targets[source_indices] + features[mask] = ((selected_features - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False) + targets[mask] = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False) + self._upload_queue.enqueue(sample.source_path) + self.cache.release_case(case_id) + wait_seconds = time.perf_counter() - started + self.recorder.add_trainer_wait(wait_seconds, step=step) + self.recorder.mark_first_batch_ready() + self._upload_queue.drain() + return np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32) + + def iter_split_batches(self, split_name: str, *, batch_size: int) -> Iterator[tuple[FloatArray, FloatArray]]: + if self.stats is None: + raise RuntimeError("Streaming normalization must be computed before evaluation") + case_ids = self._split_case_ids.get(split_name) + if case_ids is None: + raise ValueError(f"Unknown split: {split_name}") + rng = np.random.default_rng(self.config.run.seed + _SAMPLE_SEEDS[split_name]) + for case_id in case_ids: + sample = self.cache.ensure_case(case_id) + self._validate_schema(sample) + spec = self._sampling_spec_for_case(split_name, sample, rng) + for start in range(0, spec.count, batch_size): + stop = min(start + batch_size, spec.count) + source_indices = _source_indices_for_local(spec, np.arange(start, stop, dtype=np.int64)) + features = ((sample.features[source_indices] - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False) + targets = ((sample.targets[source_indices] - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False) + yield np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32) + self._upload_queue.enqueue(sample.source_path) + self.cache.release_case(case_id) + self._upload_queue.drain() + + def finish(self, *, success: bool) -> None: + self._upload_queue.drain(force=True) + self._upload_queue.finalize(training_success=success) + self.cache.state["finished_at"] = time.time() + self.cache.state["training_success"] = bool(success) + self.cache.state["processed_upload"] = self._upload_queue.manifest_payload() + self.cache._write_state() + self.recorder.set_upload_state(**self._upload_queue.telemetry_state()) + self.recorder.write_summary() + + def _validate_schema(self, sample: SimulationSample) -> None: + if self.feature_names is None or self.target_names is None: + self.feature_names = sample.feature_names + self.target_names = sample.target_names + return + if sample.feature_names != self.feature_names: + raise ValueError(f"Feature schema mismatch in streaming case {sample.case_id}") + if sample.target_names != self.target_names: + raise ValueError(f"Target schema mismatch in streaming case {sample.case_id}") + + def _sampling_spec_for_case(self, split_name: str, sample: SimulationSample, rng: np.random.Generator) -> SamplingSpec: + existing = self._sampling_specs[split_name].get(sample.case_id) + if existing is not None: + return existing + if self.config.data.points_per_case >= sample.num_points: + spec = SamplingSpec(case_id=sample.case_id, split_name=split_name, count=sample.num_points, mode="all") + else: + indices = np.sort(rng.choice(sample.num_points, size=self.config.data.points_per_case, replace=False)).astype(np.int64) + path = self.cache.indices_dir / split_name / f"{sample.case_id}.npy" + path.parent.mkdir(parents=True, exist_ok=True) + _atomic_save_npy(path, indices) + spec = SamplingSpec(case_id=sample.case_id, split_name=split_name, count=int(indices.shape[0]), mode="indexed", indices_path=path) + self._sampling_specs[split_name][sample.case_id] = spec + sampling = self.cache.state.setdefault("sampling", {}).setdefault(split_name, {}) + sampling[sample.case_id] = _sampling_spec_json(spec) + self.cache._write_state() + return spec + + def _restore_sampling_specs(self, split_name: str) -> None: + raw = self.cache.state.get("sampling", {}).get(split_name, {}) + if not isinstance(raw, dict): + return + for case_id, data in raw.items(): + if not isinstance(data, dict): + continue + path = Path(data["indices_path"]) if data.get("indices_path") else None + self._sampling_specs[split_name][case_id] = SamplingSpec( + case_id=case_id, + split_name=split_name, + count=int(data["count"]), + mode=str(data["mode"]), + indices_path=path, + ) + + def _build_train_offsets(self) -> None: + if self.split is None: + raise RuntimeError("split is missing") + counts: list[int] = [] + for case_id in self.split.train_ids: + spec = self._sampling_specs["train"].get(case_id) + if spec is None: + raise RuntimeError(f"Missing train sampling spec for {case_id}") + counts.append(spec.count) + self._train_counts = counts + self._train_offsets = np.concatenate(([0], np.cumsum(np.asarray(counts, dtype=np.int64)))).astype(np.int64) + + +@dataclass +class ProcessedUploadItem: + local_path: str + repo_path: str + bytes: int + queued_at: float + uploaded_at: float | None = None + sha256: str | None = None + + +@dataclass +class ProcessedUploadManifest: + enabled: bool + repo_id: str | None + repo_type: str | None + path_in_repo: str | None + queue_depth: int = 0 + uploaded_paths: list[str] = field(default_factory=list) + uploaded_files: list[dict[str, Any]] = field(default_factory=list) + suppressed_uploads: list[dict[str, Any]] = field(default_factory=list) + commits: list[dict[str, Any]] = field(default_factory=list) + last_error: str | None = None + rate_limit_until: float | None = None + rate_limit_retry_after_seconds: float | None = None + upload_lag_seconds: float | None = None + training_success: bool | None = None + publication_complete: bool = False + publication_status: str = "disabled" + finalized_at: float | None = None + + +class ProcessedDataUploadQueue: + def __init__( + self, + *, + enabled: bool, + run_dir: Path, + repo_id: str | None, + repo_type: str | None, + path_in_repo: str | None, + private: bool, + batch_size: int, + recorder: StreamingEventRecorder, + ) -> None: + self.enabled = enabled + self.run_dir = run_dir + self.repo_id = repo_id + self.repo_type = repo_type + self.path_in_repo = path_in_repo.strip("/") if path_in_repo else None + self.private = private + self.batch_size = max(1, batch_size) + self.recorder = recorder + self._api: Any | None = None + self._pending: dict[str, ProcessedUploadItem] = {} + self._manifest = ProcessedUploadManifest(enabled=enabled, repo_id=repo_id, repo_type=repo_type, path_in_repo=self.path_in_repo) + self.write_manifest() + + @classmethod + def from_config(cls, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> ProcessedDataUploadQueue: + enabled = bool(config.data.streaming_upload_processed and config.data.hf_repo_id) + return cls( + enabled=enabled, + run_dir=run_dir, + repo_id=config.data.hf_repo_id, + repo_type=config.data.hf_repo_type, + path_in_repo=config.data.hf_path_prefix, + private=False, + batch_size=config.data.streaming_upload_batch_size, + recorder=recorder, + ) + + def enqueue(self, local_path: Path) -> None: + if not self.enabled or not local_path.is_file(): + self._sync_manifest() + return + repo_path = f"{self.path_in_repo}/{local_path.name}" if self.path_in_repo else local_path.name + if repo_path in self._manifest.uploaded_paths or repo_path in self._pending: + self._sync_manifest() + return + self._pending[repo_path] = ProcessedUploadItem( + local_path=str(local_path), + repo_path=repo_path, + bytes=local_path.stat().st_size, + queued_at=time.time(), + ) + self.recorder.emit("processed_data_upload_queued", phase="artifacts", repo_path=repo_path, queue_depth=len(self._pending)) + self._sync_manifest() + + def drain(self, *, force: bool = False) -> None: + if not self.enabled: + self._sync_manifest() + return + if not self._pending: + self._manifest.publication_complete = True + self._sync_manifest() + return + now = time.time() + if self._manifest.rate_limit_until is not None and now < self._manifest.rate_limit_until: + self._manifest.suppressed_uploads.append( + {"suppressed_at": now, "rate_limit_until": self._manifest.rate_limit_until, "queue_depth": len(self._pending)} + ) + self.recorder.emit( + "processed_data_upload_suppressed", + phase="artifacts", + queue_depth=len(self._pending), + suppressed_until=self._manifest.rate_limit_until, + ) + self._sync_manifest() + return + if not force and len(self._pending) < self.batch_size: + self._sync_manifest() + return + batch = list(self._pending.values())[: self.batch_size] + self.recorder.emit("processed_data_upload_start", phase="artifacts", upload_count=len(batch), queue_depth=len(self._pending)) + try: + api = self._ensure_api() + from huggingface_hub import CommitOperationAdd + + operations = [CommitOperationAdd(path_in_repo=item.repo_path, path_or_fileobj=item.local_path) for item in batch] + commit = api.create_commit( + repo_id=self.repo_id, + repo_type=self.repo_type, + operations=operations, + commit_message=f"Upload {len(batch)} streamed AirfRANS processed cases", + ) + except Exception as exc: + retry_after = _retry_after_seconds(exc) + if retry_after is not None: + until = time.time() + retry_after + self._manifest.rate_limit_until = max(self._manifest.rate_limit_until or 0.0, until) + self._manifest.rate_limit_retry_after_seconds = retry_after + self._manifest.suppressed_uploads.append( + { + "rate_limited_at": time.time(), + "rate_limit_until": self._manifest.rate_limit_until, + "retry_after_seconds": retry_after, + "queue_depth": len(self._pending), + } + ) + self.recorder.emit( + "processed_data_upload_rate_limited", + phase="artifacts", + queue_depth=len(self._pending), + retry_after_seconds=retry_after, + suppressed_until=self._manifest.rate_limit_until, + ) + else: + self.recorder.emit( + "processed_data_upload_end", + phase="artifacts", + ok=False, + queue_depth=len(self._pending), + error_type=type(exc).__name__, + error_message=str(exc), + ) + self._manifest.last_error = str(exc) + self._sync_manifest() + return + uploaded_at = time.time() + for item in batch: + path = Path(item.local_path) + item.uploaded_at = uploaded_at + item.sha256 = _sha256_file(path) + self._manifest.uploaded_paths.append(item.repo_path) + self._manifest.uploaded_files.append(item.__dict__.copy()) + self._pending.pop(item.repo_path, None) + self._manifest.uploaded_paths = sorted(set(self._manifest.uploaded_paths)) + self._manifest.commits.append(_commit_payload(commit)) + self._manifest.last_error = None + self._manifest.rate_limit_until = None + self._manifest.rate_limit_retry_after_seconds = None + self.recorder.emit("processed_data_upload_end", phase="artifacts", ok=True, uploaded_count=len(batch), queue_depth=len(self._pending)) + self._sync_manifest() + + def telemetry_state(self) -> dict[str, Any]: + return { + "queue_depth": len(self._pending), + "lag_seconds": self._oldest_lag_seconds(), + "suppressed_until": self._manifest.rate_limit_until, + } + + def manifest_payload(self) -> dict[str, Any]: + self._sync_manifest() + return dict(self._manifest.__dict__) + + def finalize(self, *, training_success: bool) -> dict[str, Any]: + self._manifest.training_success = bool(training_success) + self._manifest.finalized_at = time.time() + self._sync_manifest() + return self.manifest_payload() + + def write_manifest(self) -> Path: + self._refresh_publication_status() + path = self.run_dir / PROCESSED_UPLOAD_MANIFEST + _atomic_write_json(path, dict(self._manifest.__dict__)) + return path + + def _ensure_api(self) -> Any: + if self._api is not None: + return self._api + try: + from huggingface_hub import HfApi + except ModuleNotFoundError as exc: + raise RuntimeError("huggingface_hub is required for streaming processed uploads") from exc + token = os.environ.get("HF_TOKEN") + if not token: + raise RuntimeError("HF_TOKEN is required for streaming processed uploads") + api = HfApi(token=token) + assert self.repo_id is not None + assert self.repo_type is not None + api.create_repo(repo_id=self.repo_id, repo_type=self.repo_type, private=self.private, exist_ok=True) + self._api = api + return api + + def _refresh_publication_status(self) -> None: + if not self.enabled: + self._manifest.publication_status = "disabled" + self._manifest.publication_complete = False + return + if self._manifest.training_success is False: + self._manifest.publication_status = "training_failed" + self._manifest.publication_complete = False + return + incomplete = bool(self._pending) or self._manifest.last_error is not None or self._manifest.rate_limit_until is not None + if incomplete: + self._manifest.publication_status = ( + "training_succeeded_hf_incomplete" if self._manifest.training_success is True else "hf_publication_incomplete" + ) + self._manifest.publication_complete = False + return + if self._manifest.training_success is True: + self._manifest.publication_status = "hf_publication_succeeded" + self._manifest.publication_complete = True + return + self._manifest.publication_status = "in_progress" + self._manifest.publication_complete = False + + def _sync_manifest(self) -> None: + self._manifest.queue_depth = len(self._pending) + self._manifest.upload_lag_seconds = self._oldest_lag_seconds() + self._refresh_publication_status() + self.write_manifest() + self.recorder.set_upload_state(**self.telemetry_state()) + + def _oldest_lag_seconds(self) -> float | None: + if not self._pending: + return None + oldest = min(item.queued_at for item in self._pending.values()) + return time.time() - oldest + + +class _StatsAccumulator: + def __init__(self, *, feature_names: tuple[str, ...], target_names: tuple[str, ...], min_std: float = 1e-6) -> None: + self.feature_names = feature_names + self.target_names = target_names + self.min_std = min_std + self.count = 0 + self.feature_sum = np.zeros(len(feature_names), dtype=np.float64) + self.feature_sum_sq = np.zeros(len(feature_names), dtype=np.float64) + self.target_sum = np.zeros(len(target_names), dtype=np.float64) + self.target_sum_sq = np.zeros(len(target_names), dtype=np.float64) + + def update(self, features: FloatArray, targets: FloatArray) -> None: + _validate_matrix(features, "features", len(self.feature_names)) + _validate_matrix(targets, "targets", len(self.target_names)) + if features.shape[0] != targets.shape[0]: + raise ValueError("Feature/target row mismatch while computing streaming normalization") + self.count += int(features.shape[0]) + self.feature_sum += features.sum(axis=0, dtype=np.float64) + self.feature_sum_sq += np.square(features, dtype=np.float64).sum(axis=0, dtype=np.float64) + self.target_sum += targets.sum(axis=0, dtype=np.float64) + self.target_sum_sq += np.square(targets, dtype=np.float64).sum(axis=0, dtype=np.float64) + + def finish(self) -> NormalizationStats: + if self.count <= 0: + raise ValueError("Cannot compute streaming normalization from zero rows") + feature_mean64 = self.feature_sum / self.count + target_mean64 = self.target_sum / self.count + feature_var = np.maximum(self.feature_sum_sq / self.count - np.square(feature_mean64), 0.0) + target_var = np.maximum(self.target_sum_sq / self.count - np.square(target_mean64), 0.0) + feature_std = np.sqrt(feature_var).astype(np.float32) + target_std = np.sqrt(target_var).astype(np.float32) + feature_std[feature_std < self.min_std] = 1.0 + target_std[target_std < self.min_std] = 1.0 + return NormalizationStats( + feature_mean=feature_mean64.astype(np.float32), + feature_std=feature_std, + target_mean=target_mean64.astype(np.float32), + target_std=target_std, + feature_names=self.feature_names, + target_names=self.target_names, + ) + + +def _selected_rows(values: FloatArray, spec: SamplingSpec) -> FloatArray: + if spec.mode == "all": + return np.ascontiguousarray(values, dtype=np.float32) + if spec.indices_path is None: + raise ValueError(f"Indexed sampling spec missing indices path for {spec.case_id}") + indices = np.load(spec.indices_path, allow_pickle=False) + return np.ascontiguousarray(values[indices], dtype=np.float32) + + +def _source_indices_for_local(spec: SamplingSpec, local_indices: NDArray[np.integer[Any]]) -> IntArray: + local = np.asarray(local_indices, dtype=np.int64) + if spec.mode == "all": + return local + if spec.indices_path is None: + raise ValueError(f"Indexed sampling spec missing indices path for {spec.case_id}") + indices = np.load(spec.indices_path, allow_pickle=False) + return np.asarray(indices[local], dtype=np.int64) + + +def _sampling_spec_json(spec: SamplingSpec) -> dict[str, Any]: + return { + "case_id": spec.case_id, + "split_name": spec.split_name, + "count": spec.count, + "mode": spec.mode, + "indices_path": str(spec.indices_path) if spec.indices_path is not None else None, + } + + +def _split_name_for_case(split: CaseSplit, case_id: str) -> str | None: + if case_id in split.train_ids: + return "train" + if case_id in split.val_ids: + return "val" + if case_id in split.test_ids: + return "test" + return None + + +def _validate_matrix(values: FloatArray, name: str, width: int) -> None: + if values.ndim != 2: + raise ValueError(f"Expected {name} to be a 2D array") + if values.shape[1] != width: + raise ValueError(f"Expected {name} width {width}, got {values.shape[1]}") + if not np.all(np.isfinite(values)): + raise ValueError(f"Expected finite values in {name}") + + +def _tree_size_bytes(root: Path, *, suffix: str | None = None) -> int: + if not root.exists(): + return 0 + total = 0 + for path in root.rglob("*"): + if path.is_file() and (suffix is None or path.suffix == suffix): + total += path.stat().st_size + return total + + +def _append_jsonl(path: Path, record: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with path.open("a", encoding="utf-8") as handle: + handle.write(json.dumps(dict(record), sort_keys=True) + "\n") + handle.flush() + os.fsync(handle.fileno()) + + +def _read_json(path: Path) -> dict[str, Any]: + try: + data = json.loads(path.read_text()) + except FileNotFoundError: + return {} + if not isinstance(data, dict): + raise ValueError(f"Expected JSON object in {path}") + return data + + +def _atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + tmp = path.with_name(f"{path.name}.tmp") + tmp.write_text(json.dumps(dict(payload), indent=2, sort_keys=True) + "\n") + tmp.replace(path) + + +def _atomic_save_npy(path: Path, values: NDArray[np.int64]) -> None: + tmp = path.with_name(f"{path.name}.tmp.npy") + np.save(tmp, values) + tmp.replace(path) + + +def _sha256_file(path: Path) -> str: + import hashlib + + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _commit_payload(commit: Any) -> dict[str, Any]: + return { + "commit_url": getattr(commit, "commit_url", None), + "oid": getattr(commit, "oid", None), + "pr_url": getattr(commit, "pr_url", None), + } diff --git a/tests/test_hf_upload.py b/tests/test_hf_upload.py index f3a902f..fe6b3ef 100644 --- a/tests/test_hf_upload.py +++ b/tests/test_hf_upload.py @@ -7,7 +7,7 @@ import tempfile import types import unittest from pathlib import Path -from unittest.mock import patch +from unittest.mock import Mock, patch from airfrans_frontier.training.hf_upload import HfArtifactUploader, resolve_resume_checkpoint @@ -15,7 +15,12 @@ from airfrans_frontier.training.hf_upload import HfArtifactUploader, resolve_res class HuggingFaceUploadTests(unittest.TestCase): def test_uploader_creates_repo_uploads_file_and_writes_manifest(self) -> None: created: list[tuple[str, str, bool]] = [] - uploaded: list[tuple[str, str, str]] = [] + committed: list[tuple[str, str, tuple[str, ...], str]] = [] + + class FakeCommitOperationAdd: + def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None: + self.path_in_repo = path_in_repo + self.path_or_fileobj = path_or_fileobj class FakeApi: def __init__(self, token: str) -> None: @@ -24,11 +29,11 @@ class HuggingFaceUploadTests(unittest.TestCase): def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None: created.append((repo_id, repo_type, private)) - def upload_file(self, *, repo_id: str, repo_type: str, path_or_fileobj: str, path_in_repo: str, commit_message: str): - uploaded.append((repo_id, repo_type, path_in_repo)) + def create_commit(self, *, repo_id: str, repo_type: str, operations: list[FakeCommitOperationAdd], commit_message: str): + committed.append((repo_id, repo_type, tuple(operation.path_in_repo for operation in operations), commit_message)) return types.SimpleNamespace(commit_url="https://huggingface.co/repo/commit/abc", oid="abc") - fake_module = types.SimpleNamespace(HfApi=FakeApi) + fake_module = types.SimpleNamespace(HfApi=FakeApi, CommitOperationAdd=FakeCommitOperationAdd) with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {"HF_TOKEN": "token"}): root = Path(tmp) (root / "checkpoint_latest.pt").write_bytes(b"checkpoint") @@ -45,12 +50,50 @@ class HuggingFaceUploadTests(unittest.TestCase): self.assertEqual(result["uploaded"], ["runs/model/run-1/checkpoint_latest.pt"]) self.assertEqual(created, [("owner/repo", "model", False)]) - self.assertEqual(uploaded, [("owner/repo", "model", "runs/model/run-1/checkpoint_latest.pt")]) + self.assertEqual(committed, [("owner/repo", "model", ("runs/model/run-1/checkpoint_latest.pt",), "upload checkpoint")]) manifest = json.loads((root / "hf_upload_manifest.json").read_text()) self.assertTrue(manifest["enabled"]) self.assertEqual(manifest["repo_id"], "owner/repo") self.assertIn("runs/model/run-1/checkpoint_latest.pt", manifest["uploaded_paths"]) + def test_uploader_suppresses_uploads_after_hf_retry_after_limit(self) -> None: + class FakeRateLimitError(RuntimeError): + def __init__(self) -> None: + super().__init__("429 Too Many Requests: Retry after 600 seconds") + self.response = types.SimpleNamespace(headers={"Retry-After": "600"}) + + class FakeCommitOperationAdd: + def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None: + self.path_in_repo = path_in_repo + self.path_or_fileobj = path_or_fileobj + + fake_module = types.SimpleNamespace(CommitOperationAdd=FakeCommitOperationAdd) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}): + root = Path(tmp) + (root / "metrics.jsonl").write_text("{}\n") + uploader = HfArtifactUploader( + enabled=True, + run_dir=root, + repo_id="owner/repo", + repo_type="model", + path_in_repo="runs/model/run-1", + private=False, + max_rate_limit_sleep_seconds=0, + ) + fake_api = types.SimpleNamespace(create_commit=Mock(side_effect=FakeRateLimitError())) + uploader._api = fake_api + + with self.assertRaises(FakeRateLimitError): + uploader.upload_files(("metrics.jsonl",), commit_message="first") + suppressed = uploader.upload_files(("metrics.jsonl",), commit_message="second") + + self.assertTrue(suppressed["rate_limited"]) + self.assertEqual(fake_api.create_commit.call_count, 1) + manifest = json.loads((root / "hf_upload_manifest.json").read_text()) + self.assertGreater(manifest["rate_limit_until"], 0) + self.assertEqual(manifest["rate_limit_retry_after_seconds"], 600.0) + self.assertEqual(len(manifest["suppressed_uploads"]), 2) + def test_resolve_resume_checkpoint_downloads_hf_uri(self) -> None: calls: list[tuple[str, str]] = [] diff --git a/tests/test_preflight_polish.py b/tests/test_preflight_polish.py new file mode 100644 index 0000000..c11b155 --- /dev/null +++ b/tests/test_preflight_polish.py @@ -0,0 +1,279 @@ +from __future__ import annotations + +import json +import shutil +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import Mock, patch + +from airfrans_frontier.runtime import remove_pythonpath_entries + +remove_pythonpath_entries() + +from airfrans_frontier.remote.cleanup import reconcile_cleanup +from airfrans_frontier.remote.collection import ARTIFACT_COLLECTION_REPORT, collect_artifact_paths, required_collection_failures +from airfrans_frontier.remote.config import load_remote_run_config +from airfrans_frontier.remote.launch_group import LaunchGroupScheduler, LaunchRunSpec +from airfrans_frontier.remote.selection import require_fresh_selection, selection_freshness_report +from airfrans_frontier.remote.skypilot import render_skypilot_yaml +from airfrans_frontier.remote.vast import VastOffer, choose_offer +from airfrans_frontier.training.hf_upload import HfArtifactUploader +from airfrans_frontier.training.streaming_data import StreamingEventRecorder + + +class LaunchGroupSchedulingTests(unittest.TestCase): + def test_healthy_runs_release_fragile_launch_capacity_without_serializing_training(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + state_path = Path(tmp) / "launch_state.json" + scheduler = LaunchGroupScheduler( + [ + LaunchRunSpec("run-a", "configs/a.toml"), + LaunchRunSpec("run-b", "configs/b.toml"), + LaunchRunSpec("run-c", "configs/c.toml"), + ], + max_active=3, + max_fragile=1, + state_path=state_path, + group_id="group-local", + ) + + self.assertTrue(scheduler.try_start("run-a", selected_offer_id=101, selected_host_id=11)) + self.assertFalse(scheduler.try_start("run-b", selected_offer_id=102, selected_host_id=12)) + self.assertEqual(scheduler.capacity_snapshot()["fragile"], 1) + + scheduler.mark_training_healthy("run-a") + self.assertTrue(scheduler.try_start("run-b", selected_offer_id=102, selected_host_id=12)) + self.assertEqual(scheduler.capacity_snapshot()["active"], 2) + self.assertEqual(scheduler.capacity_snapshot()["fragile"], 1) + + payload = json.loads(state_path.read_text()) + self.assertEqual(payload["launch_group_id"], "group-local") + self.assertEqual(payload["healthy_runs"], ["run-a"]) + self.assertEqual(payload["running_runs"], ["run-b"]) + self.assertEqual(payload["runs"]["run-b"]["selected_host_id"], 12) + self.assertIn("capacity_blocked", [event["event"] for event in payload["events"]]) + + +class HostAntiCollisionTests(unittest.TestCase): + def test_active_launches_avoid_duplicate_hosts_unless_allowed(self) -> None: + scheduler = LaunchGroupScheduler( + [LaunchRunSpec("run-a", "a.toml"), LaunchRunSpec("run-b", "b.toml")], + max_active=2, + max_fragile=2, + ) + + self.assertTrue(scheduler.try_start("run-a", selected_offer_id=1, selected_host_id=9)) + self.assertFalse(scheduler.try_start("run-b", selected_offer_id=2, selected_host_id=9)) + self.assertEqual(scheduler.to_payload()["runs"]["run-b"]["blocked_reason"], "host_collision") + + allowed = LaunchGroupScheduler( + [LaunchRunSpec("run-a", "a.toml"), LaunchRunSpec("run-b", "b.toml")], + max_active=2, + max_fragile=2, + allow_duplicate_hosts=True, + ) + self.assertTrue(allowed.try_start("run-a", selected_offer_id=1, selected_host_id=9)) + self.assertTrue(allowed.try_start("run-b", selected_offer_id=2, selected_host_id=9)) + + def test_offer_selection_skips_reserved_active_hosts(self) -> None: + config = load_remote_run_config("configs/remote_smoke.toml") + result = choose_offer( + [offer(10, price=0.20, host=1), offer(11, price=0.22, host=2)], + config, + query={"test": True}, + reserved_host_ids=(1,), + ) + + self.assertEqual(result.selected_offer.host_id, 2) + self.assertEqual(result.policy["reserved_host_ids"], [1]) + + +class SelectionFreshnessTests(unittest.TestCase): + def test_selection_artifacts_record_and_enforce_freshness(self) -> None: + fresh = {"selected_offer_id": 1, "created_at": 1000.0} + report = selection_freshness_report(fresh, max_age_seconds=60, now=1020.0) + self.assertTrue(report["is_fresh"]) + self.assertEqual(report["age_seconds"], 20.0) + + stale = {"selected_offer_id": 1, "created_at": 1000.0} + with self.assertRaisesRegex(ValueError, "stale"): + require_fresh_selection(stale, max_age_seconds=60, now=1100.0, path="selection.json") + + config = load_remote_run_config("configs/remote_smoke.toml") + manifest = choose_offer([offer(20, price=0.20, host=3)], config, query={}).to_manifest() + self.assertIn("created_at", manifest) + self.assertIn("created_at_iso", manifest) + self.assertIn("age_seconds", manifest) + + +class CleanupReconciliationTests(unittest.TestCase): + def test_reconciliation_uses_vast_ground_truth_for_orphans_and_records_actions(self) -> None: + destroyed: list[int] = [] + report = reconcile_cleanup( + sky_state={"clusters": [{"name": "known-run", "instance_id": 77}]}, + vast_instances=[ + {"id": 77, "actual_status": "running", "gpu_name": "RTX 4090", "num_gpus": 1, "dph_total": 0.40}, + {"id": 88, "host_id": 123, "actual_status": "running", "gpu_name": "RTX 4090", "num_gpus": 1, "dph_total": 0.45, "label": "orphan-run"}, + ], + known_run_ids=("known-run", "orphan-run"), + destroy_orphans=True, + destroy_instance=lambda instance_id: destroyed.append(instance_id), + now=1234.0, + ) + + orphan = next(item for item in report["instances"] if item["vast_instance_id"] == 88) + self.assertEqual(report["unexpected_live_count"], 1) + self.assertEqual(destroyed, [88]) + self.assertEqual(orphan["cleanup_action_attempted"], "destroy_orphan") + self.assertEqual(orphan["cleanup_result"], "destroy_requested") + self.assertEqual(orphan["hourly_cost"], 0.45) + + +class HfSafetyTests(unittest.TestCase): + def test_rate_limit_suppression_preserves_training_success_as_hf_incomplete(self) -> None: + class FakeRateLimitError(RuntimeError): + def __init__(self) -> None: + super().__init__("429 Too Many Requests") + self.response = types.SimpleNamespace(headers={"Retry-After": "600"}) + + class FakeCommitOperationAdd: + def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None: + self.path_in_repo = path_in_repo + self.path_or_fileobj = path_or_fileobj + + fake_module = types.SimpleNamespace(CommitOperationAdd=FakeCommitOperationAdd) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}): + run_dir = Path(tmp) + (run_dir / "metrics.jsonl").write_text("{}\n") + uploader = HfArtifactUploader( + enabled=True, + run_dir=run_dir, + repo_id="owner/repo", + repo_type="model", + path_in_repo="runs/run-1", + max_rate_limit_sleep_seconds=0, + ) + uploader._api = types.SimpleNamespace(create_commit=Mock(side_effect=FakeRateLimitError())) + + with self.assertRaises(FakeRateLimitError): + uploader.upload_files(("metrics.jsonl",), commit_message="upload metrics") + suppressed = uploader.upload_files(("metrics.jsonl",), commit_message="retry metrics") + final = uploader.finalize(training_success=True) + + self.assertTrue(suppressed["rate_limited"]) + self.assertEqual(final["hf_publication_status"], "training_succeeded_hf_incomplete") + manifest = json.loads((run_dir / "hf_upload_manifest.json").read_text()) + self.assertTrue(manifest["training_success"]) + self.assertFalse(manifest["publication_complete"]) + self.assertGreater(manifest["rate_limit_until"], 0) + + def test_final_reporting_distinguishes_training_failure_from_hf_success(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + disabled = HfArtifactUploader(enabled=False, run_dir=Path(tmp)) + self.assertEqual(disabled.finalize(training_success=False)["hf_publication_status"], "disabled") + + with tempfile.TemporaryDirectory() as tmp: + uploader = HfArtifactUploader(enabled=True, run_dir=Path(tmp), repo_id="owner/repo", repo_type="model", path_in_repo="run") + self.assertEqual(uploader.finalize(training_success=False)["hf_publication_status"], "training_failed") + + with tempfile.TemporaryDirectory() as tmp: + uploader = HfArtifactUploader(enabled=True, run_dir=Path(tmp), repo_id="owner/repo", repo_type="model", path_in_repo="run") + self.assertEqual(uploader.finalize(training_success=True)["hf_publication_status"], "hf_publication_succeeded") + + +class ArtifactCollectionReportTests(unittest.TestCase): + def test_collection_report_classifies_produced_missing_partial_and_failed_copy(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + + def copy_one(relative_path: str) -> int | None: + if relative_path == "produced.json": + (root / relative_path).write_text("{}\n") + return 0 + if relative_path == "missing.json": + return 0 + if relative_path == "partial.pt": + partial = root / ".rsync-partial" / relative_path + partial.parent.mkdir(parents=True) + partial.write_bytes(b"partial") + return 0 + raise RuntimeError("rsync failed") + + report = collect_artifact_paths( + local_dir=root, + remote_dir="remote:~/artifacts", + paths=("produced.json", "missing.json", "partial.pt", "failed.json"), + required=("produced.json", "partial.pt", "failed.json"), + collection_kind="terminal", + copy_one=copy_one, + ) + + by_path = {attempt["expected_path"]: attempt for attempt in report["attempts"]} + self.assertEqual(by_path["produced.json"]["final_status"], "success") + self.assertEqual(by_path["missing.json"]["likely_reason"], "remote_missing_or_not_produced") + self.assertEqual(by_path["partial.pt"]["final_status"], "partial") + self.assertEqual(by_path["failed.json"]["likely_reason"], "collection_command_failed") + self.assertEqual( + {item["expected_path"] for item in required_collection_failures(report)}, + {"partial.pt", "failed.json"}, + ) + saved = json.loads((root / ARTIFACT_COLLECTION_REPORT).read_text()) + self.assertFalse(saved["summary"]["ok"]) + + +class DiskPhilosophyTests(unittest.TestCase): + def test_disk_paths_record_telemetry_and_backpressure_state_instead_of_capacity_mismatch_hard_fail(self) -> None: + config = load_remote_run_config("configs/remote_smoke.toml") + yaml = render_skypilot_yaml(config, choose_offer([offer(30, price=0.20, host=4)], config, query={}), run_id="disk-check") + + self.assertIn("disk_telemetry.json", yaml) + self.assertIn("backpressure_adaptive", yaml) + self.assertIn("airfrans_disk_capacity_status=below_requested", yaml) + self.assertNotIn("exit 74", yaml) + + with tempfile.TemporaryDirectory() as tmp: + run_dir = Path(tmp) + recorder = StreamingEventRecorder(run_dir) + usage = shutil._ntuple_diskusage(total=1000, used=900, free=100) + with patch("airfrans_frontier.training.streaming_data.shutil.disk_usage", return_value=usage): + recorder.observe_cache(run_dir, cache_bytes=950) + recorder.emit("cache_high_water", phase="data", cache_bytes=950, high_water_bytes=900) + recorder.emit("producer_paused", phase="data", reason="cache_high_water") + recorder.emit("cache_low_water", phase="data", cache_bytes=500, low_water_bytes=600) + recorder.emit("producer_resumed", phase="data", reason="cache_low_water", idle_seconds=1.25) + + summary = recorder.to_dict() + self.assertEqual(summary["minimum_free_disk_bytes"], 100) + self.assertEqual(summary["cache_high_water_events"], 1) + self.assertEqual(summary["cache_low_water_events"], 1) + self.assertEqual(summary["producer_pause_events"], 1) + self.assertEqual(summary["producer_resume_events"], 1) + self.assertGreater(summary["producer_idle_backpressure_seconds"], 0) + + +def offer(offer_id: int, *, price: float, host: int) -> VastOffer: + return VastOffer( + id=offer_id, + gpu_name="RTX 4090", + dph_total=price, + gpu_ram=24_000, + disk_space=256.0, + geolocation="US", + inet_down_cost_per_tb=0.0, + inet_up_cost_per_tb=0.0, + host_id=host, + verification="verified", + reliability2=0.99, + cuda_max_good=12.8, + direct_port_count=1, + inet_down=500.0, + inet_up=100.0, + verified=True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_public_data.py b/tests/test_public_data.py index c6c5250..f307671 100644 --- a/tests/test_public_data.py +++ b/tests/test_public_data.py @@ -1,14 +1,38 @@ from __future__ import annotations +import gzip import sys import tempfile import types +import shutil import unittest import zipfile from pathlib import Path from unittest.mock import patch -from airfrans_frontier.raw.public import ensure_public_airfrans_processed_hf, extract_of_dataset +from airfrans_frontier.runtime import remove_pythonpath_entries + +remove_pythonpath_entries() + +from airfrans_frontier.raw.public import ensure_public_airfrans_processed_hf, extract_of_dataset, process_of_dataset_url_streaming + + +def write_minimal_airfrans_archive(archive: Path, case_names: list[str]) -> None: + with zipfile.ZipFile(archive, "w") as zf: + for case_name in case_names: + base = f"OF_dataset/{case_name}" + zf.writestr(f"{base}/constant/transportProperties", "nu 1e-5;\n") + zf.writestr( + f"{base}/constant/polyMesh/boundary", + "\naerofoil\n{\n type wall;\n nFaces 1;\n startFace 0;\n}\nfarfield\n{\n type patch;\n nFaces 3;\n startFace 1;\n}\n", + ) + zf.writestr(f"{base}/constant/polyMesh/points.gz", gzip.compress(b"4\n(\n(0 0 0)\n(1 0 0)\n(1 1 0)\n(0 1 0)\n)\n")) + zf.writestr(f"{base}/constant/polyMesh/faces.gz", gzip.compress(b"4\n(\n2(0 1)\n2(1 2)\n2(2 3)\n2(3 0)\n)\n")) + zf.writestr(f"{base}/constant/polyMesh/owner.gz", gzip.compress(b"4\n(\n0\n0\n0\n0\n)\n")) + zf.writestr(f"{base}/constant/polyMesh/neighbour.gz", gzip.compress(b"0\n(\n)\n")) + zf.writestr(f"{base}/1/U.gz", gzip.compress(b"1\n(\n(1 0 0)\n)\n")) + zf.writestr(f"{base}/1/p.gz", gzip.compress(b"1\n(\n0.5\n)\n")) + zf.writestr(f"{base}/1/nut.gz", gzip.compress(b"1\n(\n0.01\n)\n")) class PublicAirfransDataTests(unittest.TestCase): @@ -62,6 +86,76 @@ class PublicAirfransDataTests(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "Unsafe path"): extract_of_dataset(archive, tmp_path / "raw", min_cases=1) + def test_extract_of_dataset_fails_before_partial_extract_when_disk_is_too_small(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr("OF_dataset/airFoil2D_SST_demo/system/controlDict", "ok") + tiny_disk = shutil._ntuple_diskusage(total=10, used=10, free=0) + with patch("airfrans_frontier.raw.public.shutil.disk_usage", return_value=tiny_disk): + with self.assertRaisesRegex(RuntimeError, "Insufficient free disk"): + extract_of_dataset(archive, tmp_path / "raw", min_cases=1) + self.assertFalse((tmp_path / "raw" / "OF_dataset").exists()) + + def test_range_streaming_processing_writes_npz_and_discards_raw_case(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + case_name = "airFoil2D_SST_10.0_5.0_0012" + write_minimal_airfrans_archive(archive, [case_name]) + + streamed = process_of_dataset_url_streaming( + str(archive), + tmp_path / "processed", + scratch_dir=tmp_path / "streaming_raw", + min_cases=1, + progress_every=1, + ) + result = streamed.processing + + self.assertEqual(result.case_count, 1) + self.assertTrue((tmp_path / "processed" / f"{case_name}.npz").is_file()) + self.assertTrue(result.manifest_path.is_file()) + self.assertFalse((tmp_path / "streaming_raw" / case_name).exists()) + self.assertGreater(streamed.ranged_bytes_read, 0) + + def test_prepare_public_hf_streams_archive_before_publish(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + source_archive = tmp_path / "source_OF_dataset.zip" + case_name = "airFoil2D_SST_10.0_5.0_0012" + write_minimal_airfrans_archive(source_archive, [case_name]) + statuses = [ + {"file_count": 0, "npz_file_count": 0, "has_manifest": False}, + {"file_count": 2, "npz_file_count": 1, "has_manifest": True}, + ] + + def fake_publish(**kwargs): + data_root = Path(kwargs["data_root"]) + self.assertTrue((data_root / f"{case_name}.npz").is_file()) + self.assertFalse((tmp_path / "work" / "streaming_raw" / case_name).exists()) + return {"repo_url": "https://huggingface.co/datasets/owner/repo", "npz_file_count": 1} + + with patch("airfrans_frontier.raw.public._hf_dataset_status", side_effect=statuses), patch( + "airfrans_frontier.raw.public.publish_processed_dataset", side_effect=fake_publish + ): + report = ensure_public_airfrans_processed_hf( + repo_id="owner/repo", + path_in_repo="processed/full", + work_dir=tmp_path / "work", + output_dir=tmp_path / "processed", + source_url=str(source_archive), + min_cases=1, + ) + self.assertFalse((tmp_path / "work" / "OF_dataset.zip").exists()) + + self.assertTrue(report["ok"]) + self.assertTrue(report["streaming"]) + self.assertEqual(report["streaming_mode"], "zip_range") + self.assertEqual(report["download"]["mode"], "zip_range") + self.assertEqual(report["processed_case_count"], 1) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_remote_run.py b/tests/test_remote_run.py index b68b0e4..1d97058 100644 --- a/tests/test_remote_run.py +++ b/tests/test_remote_run.py @@ -1,9 +1,12 @@ from __future__ import annotations +from contextlib import redirect_stdout +from io import StringIO import json import tempfile import shutil import unittest +from unittest.mock import patch from pathlib import Path from airfrans_frontier.runtime import remove_pythonpath_entries @@ -13,7 +16,7 @@ remove_pythonpath_entries() import torch from airfrans_frontier.remote.artifacts import verify_artifacts -from airfrans_frontier.remote.cli import _classify_artifacts, _stage_resume_checkpoint +from airfrans_frontier.remote.cli import _classify_artifacts, _stage_resume_checkpoint, _terminal_artifact_names, main as remote_main from airfrans_frontier.remote.config import load_remote_run_config from airfrans_frontier.remote.skypilot import render_skypilot_yaml from airfrans_frontier.remote.vast import VastOffer, choose_offer @@ -79,6 +82,16 @@ class VastSelectionTests(unittest.TestCase): self.assertNotIn("sky launch", yaml) self.assertIn("remote-run smoke-train", yaml) self.assertIn("configs/aggressive_smoke.toml", yaml) + self.assertIn("df -h .", yaml) + self.assertIn("airfrans_disk_requested_gb=128", yaml) + self.assertIn("AIRFRANS_STARTUP_TIMELINE: artifacts/current_run/startup_timeline.jsonl", yaml) + self.assertIn("airfrans_timeline 'setup' 'started'", yaml) + self.assertIn("airfrans_timeline 'data_validation' 'started'", yaml) + self.assertIn("airfrans_timeline 'training_command' 'started'", yaml) + + def test_terminal_collection_includes_startup_timeline(self) -> None: + self.assertIn("startup_timeline.jsonl", _terminal_artifact_names(())) + def test_rendered_yaml_can_pass_resume_checkpoint(self) -> None: config = load_remote_run_config("configs/remote_smoke.toml") @@ -94,6 +107,22 @@ class VastSelectionTests(unittest.TestCase): self.assertIn("AIRFRANS_RESUME_CHECKPOINT: .airfrans_resume/airfrans-test/checkpoint_latest.pt", yaml) + +class VastInstanceCliTests(unittest.TestCase): + def test_vast_instances_reports_api_ground_truth(self) -> None: + stdout = StringIO() + with patch.dict("os.environ", {"VAST_API_KEY": "token"}), patch( + "airfrans_frontier.remote.cli.list_instances", + return_value=[{"id": 123, "actual_status": "running", "gpu_name": "RTX 4090"}], + ), redirect_stdout(stdout): + code = remote_main(["vast-instances"]) + + self.assertEqual(code, 0) + payload = json.loads(stdout.getvalue()) + self.assertEqual(payload["instance_count"], 1) + self.assertEqual(payload["instances"][0]["id"], 123) + self.assertEqual(payload["instances"][0]["actual_status"], "running") + class ArtifactVerificationTests(unittest.TestCase): def test_verify_artifacts_requires_contract_files_and_writes_manifest(self) -> None: with tempfile.TemporaryDirectory() as tmp: @@ -174,6 +203,7 @@ def offer( gpu_name="RTX 4090", dph_total=price, gpu_ram=24_000, + disk_space=256.0, geolocation=geo, inet_down_cost_per_tb=0.0, inet_up_cost_per_tb=0.0, diff --git a/tests/test_remote_smoke.py b/tests/test_remote_smoke.py new file mode 100644 index 0000000..5f476f5 --- /dev/null +++ b/tests/test_remote_smoke.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +import json +import sys +import tempfile +import types +import unittest +from pathlib import Path +from unittest.mock import patch + +from airfrans_frontier.remote.smoke import run_smoke_training + + +class SmokeTrainingFailureTests(unittest.TestCase): + def test_pre_checkpoint_training_error_writes_terminal_failure_report(self) -> None: + fake_loop = types.ModuleType("airfrans_frontier.training.loop") + + def fail_train(*_: object, **__: object) -> object: + raise RuntimeError("hub commit rate limited") + + fake_loop.train_from_config_path = fail_train + fake_torch = types.ModuleType("torch") + fake_torch.__version__ = "fake" + fake_torch.version = types.SimpleNamespace(cuda=None) + fake_torch.cuda = types.SimpleNamespace( + is_available=lambda: False, + get_device_name=lambda _index: None, + ) + + with tempfile.TemporaryDirectory() as tmp, patch.dict( + sys.modules, + { + "airfrans_frontier.training.loop": fake_loop, + "torch": fake_torch, + }, + ): + artifact_dir = Path(tmp) + with self.assertRaisesRegex(RuntimeError, "hub commit rate limited"): + run_smoke_training("missing-config.toml", artifact_dir=artifact_dir, run_id="smoke-fail") + + report = json.loads((artifact_dir / "failure_report.json").read_text()) + self.assertEqual(report["run_id"], "smoke-fail") + self.assertEqual(report["error_type"], "RuntimeError") + self.assertEqual(report["error_message"], "hub commit rate limited") + verification = json.loads((artifact_dir / "verification_report.json").read_text()) + self.assertTrue(verification["ok"]) + self.assertEqual(verification["checks"]["terminal_artifact"], "failure_report.json") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_streaming_data.py b/tests/test_streaming_data.py new file mode 100644 index 0000000..45302ac --- /dev/null +++ b/tests/test_streaming_data.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import gzip +import json +import os +import sys +import tempfile +import types +import unittest +import zipfile +from pathlib import Path +from unittest.mock import patch + +from airfrans_frontier.runtime import remove_pythonpath_entries + +remove_pythonpath_entries() + +import numpy as np + +from airfrans_frontier.raw.public import process_of_dataset_url_streaming +from airfrans_frontier.training.config import load_training_config +from airfrans_frontier.training.data import build_dataset_bundle, load_processed_dataset +from airfrans_frontier.training.loop import train +from airfrans_frontier.training.normalize import compute_normalization_stats +from airfrans_frontier.training.streaming_data import StreamingEventRecorder, StreamingTrainingData + + +def write_minimal_airfrans_archive(archive: Path, case_names: list[str]) -> None: + with zipfile.ZipFile(archive, "w") as zf: + for index, case_name in enumerate(case_names): + base = f"OF_dataset/{case_name}" + u_value = 1.0 + 0.1 * index + p_value = 0.5 + 0.2 * index + nut_value = 0.01 + 0.001 * index + zf.writestr(f"{base}/constant/transportProperties", "nu 1e-5;\n") + zf.writestr( + f"{base}/constant/polyMesh/boundary", + "\naerofoil\n{\n type wall;\n nFaces 1;\n startFace 0;\n}\nfarfield\n{\n type patch;\n nFaces 3;\n startFace 1;\n}\n", + ) + zf.writestr(f"{base}/constant/polyMesh/points.gz", gzip.compress(b"4\n(\n(0 0 0)\n(1 0 0)\n(1 1 0)\n(0 1 0)\n)\n")) + zf.writestr(f"{base}/constant/polyMesh/faces.gz", gzip.compress(b"4\n(\n2(0 1)\n2(1 2)\n2(2 3)\n2(3 0)\n)\n")) + zf.writestr(f"{base}/constant/polyMesh/owner.gz", gzip.compress(b"4\n(\n0\n0\n0\n0\n)\n")) + zf.writestr(f"{base}/constant/polyMesh/neighbour.gz", gzip.compress(b"0\n(\n)\n")) + zf.writestr(f"{base}/1/U.gz", gzip.compress(f"1\n(\n({u_value} 0 0)\n)\n".encode())) + zf.writestr(f"{base}/1/p.gz", gzip.compress(f"1\n(\n{p_value}\n)\n".encode())) + zf.writestr(f"{base}/1/nut.gz", gzip.compress(f"1\n(\n{nut_value}\n)\n".encode())) + + +def write_malformed_airfrans_archive(archive: Path, case_name: str) -> None: + with zipfile.ZipFile(archive, "w") as zf: + zf.writestr(f"OF_dataset/{case_name}/constant/transportProperties", "nu 1e-5;\n") + + +def write_streaming_config( + path: Path, + *, + archive: Path, + cache_dir: Path, + artifact_dir: Path, + train_cases: int = 2, + val_cases: int = 1, + test_cases: int = 1, + steps: int = 2, + log_interval: int = 1, + batch_size: int = 2, + high_water_bytes: int = 32 * 1024 * 1024, + low_water_bytes: int = 16 * 1024 * 1024, + upload_processed: bool = False, + upload_batch_size: int = 1, +) -> None: + path.write_text( + f""" +[run] +name = "streaming_test" +seed = 7 +artifact_dir = "{artifact_dir}" + +[data] +root = "{cache_dir}" +source = "public_zip_streaming" +public_source_url = "{archive}" +cache_dir = "{cache_dir}" +streaming_scratch_dir = "{cache_dir / '_raw'}" +train_cases = {train_cases} +val_cases = {val_cases} +test_cases = {test_cases} +points_per_case = 999999999 +batch_size = {batch_size} +streaming_cache_max_bytes = {max(high_water_bytes, high_water_bytes + 1)} +streaming_cache_high_water_bytes = {high_water_bytes} +streaming_cache_low_water_bytes = {low_water_bytes} +streaming_queue_max_cases = 1 +streaming_upload_processed = {str(upload_processed).lower()} +streaming_upload_batch_size = {upload_batch_size} +hf_repo_id = "owner/airfrans-processed" +hf_repo_type = "dataset" +hf_path_prefix = "processed/full" + +[model] +type = "mlp" +hidden_width = 16 +depth = 2 +activation = "gelu" + +[optim] +lr = 0.01 +weight_decay = 0.0 +steps = {steps} +log_interval = {log_interval} + +[device] +type = "cpu" +allow_cpu_fallback = false +benchmark_kernels = false + +[loss] +type = "normalized_mse" + +[checkpoint] +interval_seconds = 0 +""".strip() + + "\n" + ) + + +def read_events(run_dir: Path) -> list[dict[str, object]]: + return [json.loads(line) for line in (run_dir / "streaming_events.jsonl").read_text().splitlines() if line.strip()] + + +class FullDataBackpressureStreamingTests(unittest.TestCase): + def test_streaming_training_smoke_writes_artifacts_without_eager_concatenation(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + case_names = [f"airFoil2D_SST_{10 + index}.0_5.0_0012" for index in range(5)] + write_minimal_airfrans_archive(archive, case_names) + config_path = tmp_path / "streaming.toml" + artifact_dir = tmp_path / "artifacts" + write_streaming_config(config_path, archive=archive, cache_dir=tmp_path / "cache", artifact_dir=artifact_dir) + config = load_training_config(config_path) + + with patch("airfrans_frontier.training.loop.load_processed_dataset", side_effect=AssertionError("eager load called")), patch( + "airfrans_frontier.training.loop.build_dataset_bundle", side_effect=AssertionError("eager concat called") + ): + result = train(config) + + self.assertTrue(np.isfinite(result.final_metrics["train_loss"])) + self.assertEqual(result.final_metrics["data_mode"], "public_zip_streaming") + for name in ( + "metrics.jsonl", + "checkpoint_latest.pt", + "checkpoint_best.pt", + "checkpoint_final.pt", + "final_metrics.json", + "split_manifest.json", + "data_manifest.json", + "normalization.json", + "streaming_events.jsonl", + "streaming_state.json", + "streaming_summary.json", + "processed_upload_manifest.json", + "artifact_manifest.json", + "checksums.txt", + "verification_report.json", + ): + self.assertTrue((result.run_dir / name).is_file(), name) + events = read_events(result.run_dir) + event_names = {event["event"] for event in events} + self.assertIn("dataset_enumeration_start", event_names) + self.assertIn("dataset_enumeration_end", event_names) + self.assertIn("split_selection", event_names) + self.assertIn("normalization_start", event_names) + self.assertIn("normalization_end", event_names) + self.assertIn("first_batch_ready", event_names) + self.assertIn("first_gpu_batch_consumed", event_names) + self.assertIn("first_metric", event_names) + self.assertIn("first_checkpoint_written", event_names) + selected_cases = set(json.loads((result.run_dir / "data_manifest.json").read_text())["cases"][index]["case_id"] for index in range(4)) + processed_cases = {str(event["case_id"]) for event in events if event["event"] == "processing_end"} + self.assertLessEqual(processed_cases, selected_cases) + self.assertFalse(any((tmp_path / "cache" / "_raw").glob("airFoil2D_*"))) + + def test_backpressure_pauses_resumes_and_bounds_cache_with_inflight_slack(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + case_names = [f"airFoil2D_SST_{10 + index}.0_5.0_0012" for index in range(4)] + write_minimal_airfrans_archive(archive, case_names) + config_path = tmp_path / "streaming.toml" + high_water = 256 + write_streaming_config( + config_path, + archive=archive, + cache_dir=tmp_path / "cache", + artifact_dir=tmp_path / "artifacts", + high_water_bytes=high_water, + low_water_bytes=128, + steps=2, + ) + + result = train(load_training_config(config_path)) + + summary = json.loads((result.run_dir / "streaming_summary.json").read_text()) + self.assertGreater(summary["cache_high_water_events"], 0) + self.assertGreater(summary["cache_low_water_events"], 0) + self.assertGreater(summary["producer_pause_events"], 0) + self.assertGreater(summary["producer_resume_events"], 0) + self.assertGreater(summary["evicted_units"], 0) + self.assertLessEqual(summary["processed_cache_high_water_bytes"], high_water + summary["max_processed_unit_bytes"]) + event_names = {event["event"] for event in read_events(result.run_dir)} + self.assertIn("producer_paused", event_names) + self.assertIn("producer_resumed", event_names) + self.assertIn("cleanup_eviction", event_names) + + def test_streaming_normalization_matches_eager_train_split_statistics(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + case_names = [f"airFoil2D_SST_{10 + index}.0_5.0_0012" for index in range(4)] + write_minimal_airfrans_archive(archive, case_names) + config_path = tmp_path / "streaming.toml" + write_streaming_config(config_path, archive=archive, cache_dir=tmp_path / "cache", artifact_dir=tmp_path / "artifacts", steps=1) + config = load_training_config(config_path) + run_dir = tmp_path / "run" + recorder = StreamingEventRecorder(run_dir) + streaming = StreamingTrainingData.from_config(config, run_dir=run_dir, recorder=recorder) + + streaming.prepare() + streaming_stats = streaming.load_or_compute_normalization() + eager_root = tmp_path / "eager_processed" + process_of_dataset_url_streaming(str(archive), eager_root, scratch_dir=tmp_path / "eager_raw", min_cases=4) + eager_bundle = build_dataset_bundle( + load_processed_dataset(eager_root), + train_cases=config.data.train_cases, + val_cases=config.data.val_cases, + test_cases=config.data.test_cases, + points_per_case=config.data.points_per_case, + seed=config.run.seed, + ) + eager_stats = compute_normalization_stats( + eager_bundle.train.features, + eager_bundle.train.targets, + feature_names=eager_bundle.feature_names, + target_names=eager_bundle.target_names, + ) + + np.testing.assert_allclose(streaming_stats.feature_mean, eager_stats.feature_mean, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(streaming_stats.feature_std, eager_stats.feature_std, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(streaming_stats.target_mean, eager_stats.target_mean, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(streaming_stats.target_std, eager_stats.target_std, rtol=1e-6, atol=1e-6) + + def test_resume_reuses_validated_units_and_discards_partial_units(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + case_names = [f"airFoil2D_SST_{10 + index}.0_5.0_0012" for index in range(3)] + write_minimal_airfrans_archive(archive, case_names) + config_path = tmp_path / "streaming.toml" + write_streaming_config( + config_path, + archive=archive, + cache_dir=tmp_path / "cache", + artifact_dir=tmp_path / "artifacts", + train_cases=1, + val_cases=1, + test_cases=1, + steps=1, + ) + config = load_training_config(config_path) + run_dir = tmp_path / "run" + first = StreamingTrainingData.from_config(config, run_dir=run_dir, recorder=StreamingEventRecorder(run_dir)) + first.prepare() + assert first.split is not None + first_case = first.split.train_ids[0] + (tmp_path / "cache" / f"{first_case}.npz.tmp.npz").write_bytes(b"partial") + + second = StreamingTrainingData.from_config(config, run_dir=run_dir, recorder=StreamingEventRecorder(run_dir)) + second.prepare() + + events = read_events(run_dir) + self.assertTrue(any(event["event"] == "partial_unit_discarded" and event.get("case_id") == first_case for event in events)) + self.assertTrue(any(event["event"] == "resume_validated_unit_reused" and event.get("case_id") == first_case for event in events)) + processing_events = [event for event in events if event["event"] == "processing_end" and event.get("case_id") == first_case] + self.assertEqual(len(processing_events), 1) + + def test_processed_upload_rate_limit_does_not_fail_training(self) -> None: + class FakeRateLimitError(RuntimeError): + def __init__(self) -> None: + super().__init__("429 Too Many Requests") + self.response = types.SimpleNamespace(headers={"Retry-After": "600"}) + + class FakeCommitOperationAdd: + def __init__(self, *, path_in_repo: str, path_or_fileobj: str) -> None: + self.path_in_repo = path_in_repo + self.path_or_fileobj = path_or_fileobj + + class FakeApi: + def __init__(self, token: str) -> None: + self.token = token + + def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None: + return None + + def create_commit(self, **kwargs): + raise FakeRateLimitError() + + fake_module = types.SimpleNamespace(HfApi=FakeApi, CommitOperationAdd=FakeCommitOperationAdd) + with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(os.environ, {"HF_TOKEN": "token"}): + tmp_path = Path(tmp) + archive = tmp_path / "OF_dataset.zip" + case_names = [f"airFoil2D_SST_{10 + index}.0_5.0_0012" for index in range(4)] + write_minimal_airfrans_archive(archive, case_names) + config_path = tmp_path / "streaming.toml" + write_streaming_config( + config_path, + archive=archive, + cache_dir=tmp_path / "cache", + artifact_dir=tmp_path / "artifacts", + upload_processed=True, + upload_batch_size=1, + steps=1, + ) + + result = train(load_training_config(config_path)) + + self.assertTrue(np.isfinite(result.final_metrics["train_loss"])) + manifest = json.loads((result.run_dir / "processed_upload_manifest.json").read_text()) + self.assertTrue(manifest["enabled"]) + self.assertGreater(manifest["queue_depth"], 0) + self.assertGreater(manifest["rate_limit_until"], 0) + self.assertEqual(manifest["rate_limit_retry_after_seconds"], 600.0) + events = {event["event"] for event in read_events(result.run_dir)} + self.assertIn("processed_data_upload_rate_limited", events) + self.assertIn("processed_data_upload_suppressed", events) + run_manifest = json.loads((result.run_dir / "run_manifest.json").read_text()) + self.assertEqual(run_manifest["phase"], "completed") + + def test_streaming_failure_writes_diagnostic_artifacts(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + tmp_path = Path(tmp) + archive = tmp_path / "bad.zip" + case_name = "airFoil2D_SST_10.0_5.0_0012" + write_malformed_airfrans_archive(archive, case_name) + config_path = tmp_path / "streaming.toml" + artifact_dir = tmp_path / "artifacts" + write_streaming_config( + config_path, + archive=archive, + cache_dir=tmp_path / "cache", + artifact_dir=artifact_dir, + train_cases=1, + val_cases=0, + test_cases=0, + steps=1, + ) + + with self.assertRaises(Exception): + train(load_training_config(config_path)) + + run_dir = next(path for path in artifact_dir.iterdir() if path.is_dir()) + for name in ("failure_report.json", "metrics.jsonl", "streaming_events.jsonl", "streaming_state.json", "streaming_summary.json", "verification_report.json"): + self.assertTrue((run_dir / name).is_file(), name) + report = json.loads((run_dir / "failure_report.json").read_text()) + self.assertEqual(report["phase"], "streaming_training") + events = {event["event"] for event in read_events(run_dir)} + self.assertIn("processing_failure", events) + verification = json.loads((run_dir / "verification_report.json").read_text()) + self.assertFalse(verification["ok"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_training_loop.py b/tests/test_training_loop.py index 8bd8051..75c9c89 100644 --- a/tests/test_training_loop.py +++ b/tests/test_training_loop.py @@ -18,6 +18,7 @@ import torch from unittest.mock import Mock from airfrans_frontier.training.config import load_training_config +from airfrans_frontier.training.hf_upload import HfArtifactUploader from airfrans_frontier.training.loop import train, select_device @@ -107,6 +108,40 @@ interval_seconds = {checkpoint_interval_seconds} ) +class HfArtifactUploaderTests(unittest.TestCase): + def test_upload_files_commits_batch_once_to_reduce_hub_rate_limit(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + run_dir = Path(tmp) + (run_dir / "metrics.jsonl").write_text("{}\n") + (run_dir / "latest_metrics.json").write_text("{}\n") + uploader = HfArtifactUploader( + enabled=True, + run_dir=run_dir, + repo_id="owner/repo", + repo_type="model", + path_in_repo="runs/model", + ) + fake_api = Mock() + fake_api.create_commit.return_value = Mock(oid="abc123", commit_url="https://hf/commit/abc123", pr_url=None) + uploader._api = fake_api + + result = uploader.upload_files(("metrics.jsonl", "latest_metrics.json"), commit_message="batch artifacts") + + self.assertEqual(result["uploaded"], ["runs/model/metrics.jsonl", "runs/model/latest_metrics.json"]) + fake_api.create_commit.assert_called_once() + call_kwargs = fake_api.create_commit.call_args.kwargs + self.assertEqual(call_kwargs["repo_id"], "owner/repo") + self.assertEqual(call_kwargs["repo_type"], "model") + self.assertEqual(call_kwargs["commit_message"], "batch artifacts") + self.assertEqual( + [operation.path_in_repo for operation in call_kwargs["operations"]], + ["runs/model/metrics.jsonl", "runs/model/latest_metrics.json"], + ) + manifest = json.loads((run_dir / "hf_upload_manifest.json").read_text()) + self.assertEqual(len(manifest["commits"]), 1) + self.assertEqual(set(manifest["uploaded_paths"]), set(result["uploaded"])) + + class TrainingLoopTests(unittest.TestCase): def test_cuda_config_fails_clearly_when_cuda_unavailable(self) -> None: with tempfile.TemporaryDirectory() as tmp: diff --git a/uv.lock b/uv.lock index 116cf6b..729577b 100644 --- a/uv.lock +++ b/uv.lock @@ -199,7 +199,7 @@ dev = [ requires-dist = [ { name = "huggingface-hub", specifier = ">=0.36.0" }, { name = "numpy", specifier = ">=2.4.0" }, - { name = "torch", specifier = ">=2.8.0" }, + { name = "torch", specifier = ">=2.7.1,<2.8.0" }, { name = "wandb", specifier = ">=0.23.0" }, ] @@ -761,83 +761,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/aa/50/a9caea39ad19c431c1a3f8a31114df65b260cdfe67786b6c7e7c040c4c44/cryptography-49.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:be9fcb48a55f023493482827d4f459bd263cc20efde64f204b97c123201850c6", size = 3783731, upload-time = "2026-06-12T20:02:43.319Z" }, ] -[[package]] -name = "cuda-bindings" -version = "13.3.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "cuda-pathfinder", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] -wheels = [ - { url = "https://files.pythonhosted.org/packages/51/6b/457ca12dad3ee9bfcc9a545cfd6b64b359ba49de40f776f6e028e678f262/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c5879712accf6e14bb01aa5e67440eb84998b8d104b509cc7a6dc0b8f656a474", size = 6053539, upload-time = "2026-05-29T23:11:43.19Z" }, - { url = "https://files.pythonhosted.org/packages/95/7a/c5e3c34a409b148f5c0f5a4ea374158f95d488862c1dffedf9aa5c639df9/cuda_bindings-13.3.1-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:04436a9364059c84b8f9636f359eccda1cf814341f5b670c71d80d2f79dbc708", size = 6674166, upload-time = "2026-05-29T23:11:45.478Z" }, - { url = "https://files.pythonhosted.org/packages/ce/67/5e7dba1ba576dd73da5dee894ca076ca5e959450dfff66d6d510a255d1f7/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c7855c4868aabc0cfae28abbe83d56734bdfbd08f08fc234ac1912a12858bf49", size = 6025351, upload-time = "2026-05-29T23:11:49.685Z" }, - { url = "https://files.pythonhosted.org/packages/39/2a/6d2e9047d1fb243dbaa364b01e0297534b9ed7fd27dba1c9f361519cf69b/cuda_bindings-13.3.1-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e32d08f71ebcdf00f0f41eab2eb37e8da94c8ed411cc9f7f7a019ce6b34abe3a", size = 6657965, upload-time = "2026-05-29T23:11:52.227Z" }, - { url = "https://files.pythonhosted.org/packages/cc/6e/2394f8163360f8391f8f1b7e72d300a82724edb81a7b7084c799fbd4c91f/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9efb21c1ee64981e184b9e0ba5eb3179e5ba3d4b51665a6cb52b8ef3d01a7cbf", size = 5920504, upload-time = "2026-05-29T23:11:56.883Z" }, - { url = "https://files.pythonhosted.org/packages/34/c2/ef9b6a63f7dc432712a462c816662e662e00d38caa9b861c8c2588195d03/cuda_bindings-13.3.1-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2732904099e0a4d4db774a5fc6d91ee95fae065b4d2ecabb4968c5fe2406c9d7", size = 6476660, upload-time = "2026-05-29T23:11:59.188Z" }, - { url = "https://files.pythonhosted.org/packages/b1/81/bff68ce829999c1e4209c761bbf903b1c06ec570416ddb25020864ad5907/cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:1ab2f74ed65bfef4163ba07a8db16f1085e0729291db12a2423aff84ee8278b8", size = 6013639, upload-time = "2026-05-29T23:12:03.509Z" }, - { url = "https://files.pythonhosted.org/packages/d4/e0/c8a1f0c8f9ffdea4f5fe6dbab89b326cef4d85caf489dad39e209da89416/cuda_bindings-13.3.1-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:efd4c814d311ec08c981f6dded1dbe7d4b371067ee4f6c14cccec4bde9590f80", size = 6534419, upload-time = "2026-05-29T23:12:05.633Z" }, - { url = "https://files.pythonhosted.org/packages/52/b8/83b1f563925b290f2d11a01a77a84013ba56052fe3653a5bef3ccfbb43d6/cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c3c772dfff49681541d59630c90f858e173ac926b9c593a2b7123f2a1043cc76", size = 5809771, upload-time = "2026-05-29T23:12:10.422Z" }, - { url = "https://files.pythonhosted.org/packages/12/20/e79b4bfe98f075195afb6343d41c498f9dbd2d161d7021d4d28bceb83581/cuda_bindings-13.3.1-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:36febb7c1079d68a981dbbd8d5a67235b399802b82075c9388624719607e52b9", size = 6358584, upload-time = "2026-05-29T23:12:12.767Z" }, -] - -[[package]] -name = "cuda-pathfinder" -version = "1.5.6" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d2/53/8fc9b0cdc5b7f62746e6a01b85b6461e5ae27f871010a5fcf8fa6950766d/cuda_pathfinder-1.5.6-py3-none-any.whl", hash = "sha256:7e4c07c117b78ba1fb35dac4c444d21f3677b1b1ff56175c53a8e3025c5b43c0", size = 52972, upload-time = "2026-06-30T00:58:04.34Z" }, -] - -[[package]] -name = "cuda-toolkit" -version = "13.0.3.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/c7/a79086a62c98befcdb8349656c6f114e2db3b8b2422f6e25c97a7f2a9a3c/cuda_toolkit-13.0.3.0-py2.py3-none-any.whl", hash = "sha256:d693caaa261214ddd7dbb60d68e71cbed884e68c2be7509778f3051da0b91c3f", size = 2512, upload-time = "2026-04-14T00:50:08.173Z" }, -] - -[package.optional-dependencies] -cublas = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cudart = [ - { name = "nvidia-cuda-runtime", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cufft = [ - { name = "nvidia-cufft", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cufile = [ - { name = "nvidia-cufile", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cupti = [ - { name = "nvidia-cuda-cupti", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -curand = [ - { name = "nvidia-curand", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cusolver = [ - { name = "nvidia-cublas", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusolver", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -cusparse = [ - { name = "nvidia-cusparse", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] -nvtx = [ - { name = "nvidia-nvtx", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, -] - [[package]] name = "cycler" version = "0.12.1" @@ -2314,155 +2237,136 @@ wheels = [ ] [[package]] -name = "nvidia-cublas" -version = "13.1.1.3" +name = "nvidia-cublas-cu12" +version = "12.6.4.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/af/eb/ff4b8c503fa1f1796679dce648854d58751982426e4e4b37d6fce49d259c/nvidia_cublas_cu12-12.6.4.1-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08ed2686e9875d01b58e3cb379c6896df8e76c75e0d4a7f7dace3d7b6d9ef8eb", size = 393138322, upload-time = "2024-11-20T17:40:25.65Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.6.80" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/60/7b6497946d74bcf1de852a21824d63baad12cd417db4195fc1bfe59db953/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6768bad6cab4f19e8292125e5f1ac8aa7d1718704012a0e3272a6f61c4bce132", size = 8917980, upload-time = "2024-11-20T17:36:04.019Z" }, + { url = "https://files.pythonhosted.org/packages/a5/24/120ee57b218d9952c379d1e026c4479c9ece9997a4fb46303611ee48f038/nvidia_cuda_cupti_cu12-12.6.80-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a3eff6cdfcc6a4c35db968a06fcadb061cbc7d6dde548609a941ff8701b98b73", size = 8917972, upload-time = "2024-10-01T16:58:06.036Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.6.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/75/2e/46030320b5a80661e88039f59060d1790298b4718944a65a7f2aeda3d9e9/nvidia_cuda_nvrtc_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:35b0cc6ee3a9636d5409133e79273ce1f3fd087abb0532d2d2e8fff1fe9efc53", size = 23650380, upload-time = "2024-10-01T17:00:14.643Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.6.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e1/23/e717c5ac26d26cf39a27fbc076240fad2e3b817e5889d671b67f4f9f49c5/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba3b56a4f896141e25e19ab287cd71e52a6a0f4b29d0d31609f60e3b4d5219b7", size = 897690, upload-time = "2024-11-20T17:35:30.697Z" }, + { url = "https://files.pythonhosted.org/packages/f0/62/65c05e161eeddbafeca24dc461f47de550d9fa8a7e04eb213e32b55cfd99/nvidia_cuda_runtime_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:a84d15d5e1da416dd4774cb42edf5e954a3e60cc945698dc1d5be02321c44dc8", size = 897678, upload-time = "2024-10-01T16:57:33.821Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "9.5.1.17" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cuda-nvrtc", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/a7/a1/0bd24ee8c8d03adac032fd2909426a00c88f8c57961b1277ded97f91119f/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:b7a210458267ac818974c53038fbec2e969d5c99f305ab15c72522fa9f001dd5", size = 542848918, upload-time = "2026-04-08T18:46:22.985Z" }, - { url = "https://files.pythonhosted.org/packages/3b/cd/154ca20c38269e05eff77c1464e6c1da89f50a6390b565e9d82e06bc11e1/nvidia_cublas-13.1.1.3-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:37936a16db8fe4ac1f065c2139360608a543a09275cb1a1af612e08cfa065436", size = 423138758, upload-time = "2026-04-08T18:46:58.655Z" }, + { url = "https://files.pythonhosted.org/packages/2a/78/4535c9c7f859a64781e43c969a3a7e84c54634e319a996d43ef32ce46f83/nvidia_cudnn_cu12-9.5.1.17-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:30ac3869f6db17d170e0e556dd6cc5eee02647abc31ca856634d5a40f82c15b2", size = 570988386, upload-time = "2024-10-25T19:54:26.39Z" }, ] [[package]] -name = "nvidia-cuda-cupti" -version = "13.0.85" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/2a/2a/80353b103fc20ce05ef51e928daed4b6015db4aaa9162ed0997090fe2250/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_aarch64.whl", hash = "sha256:796bd679890ee55fb14a94629b698b6db54bcfd833d391d5e94017dd9d7d3151", size = 10310827, upload-time = "2025-09-04T08:26:42.012Z" }, - { url = "https://files.pythonhosted.org/packages/33/6d/737d164b4837a9bbd202f5ae3078975f0525a55730fe871d8ed4e3b952b0/nvidia_cuda_cupti-13.0.85-py3-none-manylinux_2_25_x86_64.whl", hash = "sha256:4eb01c08e859bf924d222250d2e8f8b8ff6d3db4721288cf35d14252a4d933c8", size = 10715597, upload-time = "2025-09-04T08:26:51.312Z" }, -] - -[[package]] -name = "nvidia-cuda-nvrtc" -version = "13.0.88" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/68/483a78f5e8f31b08fb1bb671559968c0ca3a065ac7acabfc7cee55214fd6/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:ad9b6d2ead2435f11cbb6868809d2adeeee302e9bb94bcf0539c7a40d80e8575", size = 90215200, upload-time = "2025-09-04T08:28:44.204Z" }, - { url = "https://files.pythonhosted.org/packages/b7/dc/6bb80850e0b7edd6588d560758f17e0550893a1feaf436807d64d2da040f/nvidia_cuda_nvrtc-13.0.88-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:d27f20a0ca67a4bb34268a5e951033496c5b74870b868bacd046b1b8e0c3267b", size = 43015449, upload-time = "2025-09-04T08:28:20.239Z" }, -] - -[[package]] -name = "nvidia-cuda-runtime" -version = "13.0.96" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/87/4f/17d7b9b8e285199c58ce28e31b5c5bbaa4d8271af06a89b6405258245de2/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ef9bcbe90493a2b9d810e43d249adb3d02e98dd30200d86607d8d02687c43f55", size = 2261060, upload-time = "2025-10-09T08:55:15.78Z" }, - { url = "https://files.pythonhosted.org/packages/2e/24/d1558f3b68b1d26e706813b1d10aa1d785e4698c425af8db8edc3dced472/nvidia_cuda_runtime-13.0.96-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7f82250d7782aa23b6cfe765ecc7db554bd3c2870c43f3d1821f1d18aebf0548", size = 2243632, upload-time = "2025-10-09T08:55:36.117Z" }, -] - -[[package]] -name = "nvidia-cudnn-cu13" -version = "9.20.0.48" +name = "nvidia-cufft-cu12" +version = "11.3.0.4" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/56/c5/83384d846b2fd17c44bd499b36c75a45ed4f095fbbb2252294e89cea5c5c/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:e31454ae00094b0c55319d9d15b6fa2fc50a9e1c0f5c8c80fb75258234e731e1", size = 444574296, upload-time = "2026-03-09T19:28:27.751Z" }, - { url = "https://files.pythonhosted.org/packages/6e/5e/edb9c0ae051602c3ccaffe424256463636d639e27d7f302dde9975ef9e7a/nvidia_cudnn_cu13-9.20.0.48-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0c45dd8eeb50b603f07995b1b300c62ffe6a1980482b82b3bcf94a4ca9d49304", size = 366173588, upload-time = "2026-03-09T19:29:34.474Z" }, + { url = "https://files.pythonhosted.org/packages/8f/16/73727675941ab8e6ffd86ca3a4b7b47065edcca7a997920b831f8147c99d/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ccba62eb9cef5559abd5e0d54ceed2d9934030f51163df018532142a8ec533e5", size = 200221632, upload-time = "2024-11-20T17:41:32.357Z" }, + { url = "https://files.pythonhosted.org/packages/60/de/99ec247a07ea40c969d904fc14f3a356b3e2a704121675b75c366b694ee1/nvidia_cufft_cu12-11.3.0.4-py3-none-manylinux2014_x86_64.whl", hash = "sha256:768160ac89f6f7b459bee747e8d175dbf53619cfe74b2a5636264163138013ca", size = 200221622, upload-time = "2024-10-01T17:03:58.79Z" }, ] [[package]] -name = "nvidia-cufft" -version = "12.0.0.61" +name = "nvidia-cufile-cu12" +version = "1.11.1.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/66/cc9876340ac68ae71b15c743ddb13f8b30d5244af344ec8322b449e35426/nvidia_cufile_cu12-1.11.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:cc23469d1c7e52ce6c1d55253273d32c565dd22068647f3aa59b3c6b005bf159", size = 1142103, upload-time = "2024-11-20T17:42:11.83Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.7.77" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/73/1b/44a01c4e70933637c93e6e1a8063d1e998b50213a6b65ac5a9169c47e98e/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a42cd1344297f70b9e39a1e4f467a4e1c10f1da54ff7a85c12197f6c652c8bdf", size = 56279010, upload-time = "2024-11-20T17:42:50.958Z" }, + { url = "https://files.pythonhosted.org/packages/4a/aa/2c7ff0b5ee02eaef890c0ce7d4f74bc30901871c5e45dee1ae6d0083cd80/nvidia_curand_cu12-10.3.7.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:99f1a32f1ac2bd134897fc7a203f779303261268a65762a623bf30cc9fe79117", size = 56279000, upload-time = "2024-10-01T17:04:45.274Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.7.1.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cublas-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-cusparse-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, - { url = "https://files.pythonhosted.org/packages/a8/2f/7b57e29836ea8714f81e9898409196f47d772d5ddedddf1592eadb8ab743/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:6c44f692dce8fd5ffd3e3df134b6cdb9c2f72d99cf40b62c32dde45eea9ddad3", size = 214085489, upload-time = "2025-09-04T08:31:56.044Z" }, + { url = "https://files.pythonhosted.org/packages/f0/6e/c2cf12c9ff8b872e92b4a5740701e51ff17689c4d726fca91875b07f655d/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:e9e49843a7707e42022babb9bcfa33c29857a93b88020c4e4434656a655b698c", size = 158229790, upload-time = "2024-11-20T17:43:43.211Z" }, + { url = "https://files.pythonhosted.org/packages/9f/81/baba53585da791d043c10084cf9553e074548408e04ae884cfe9193bd484/nvidia_cusolver_cu12-11.7.1.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6cf28f17f64107a0c4d7802be5ff5537b2130bfc112f25d5a30df227058ca0e6", size = 158229780, upload-time = "2024-10-01T17:05:39.875Z" }, ] [[package]] -name = "nvidia-cufile" -version = "1.15.1.6" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/3f/70/4f193de89a48b71714e74602ee14d04e4019ad36a5a9f20c425776e72cd6/nvidia_cufile-1.15.1.6-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:08a3ecefae5a01c7f5117351c64f17c7c62efa5fffdbe24fc7d298da19cd0b44", size = 1223672, upload-time = "2025-09-04T08:32:22.779Z" }, - { url = "https://files.pythonhosted.org/packages/ab/73/cc4a14c9813a8a0d509417cf5f4bdaba76e924d58beb9864f5a7baceefbf/nvidia_cufile-1.15.1.6-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:bdc0deedc61f548bddf7733bdc216456c2fdb101d020e1ab4b88d232d5e2f6d1", size = 1136992, upload-time = "2025-09-04T08:32:14.119Z" }, -] - -[[package]] -name = "nvidia-curand" -version = "10.4.0.35" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/72/7c2ae24fb6b63a32e6ae5d241cc65263ea18d08802aaae087d9f013335a2/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:133df5a7509c3e292aaa2b477afd0194f06ce4ea24d714d616ff36439cee349a", size = 61962106, upload-time = "2025-08-04T10:21:41.128Z" }, - { url = "https://files.pythonhosted.org/packages/a5/9f/be0a41ca4a4917abf5cb9ae0daff1a6060cc5de950aec0396de9f3b52bc5/nvidia_curand-10.4.0.35-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:1aee33a5da6e1db083fe2b90082def8915f30f3248d5896bcec36a579d941bfc", size = 59544258, upload-time = "2025-08-04T10:22:03.992Z" }, -] - -[[package]] -name = "nvidia-cusolver" -version = "12.0.4.66" +name = "nvidia-cusparse-cu12" +version = "12.5.4.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-cusparse", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "nvidia-nvjitlink-cu12", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, - { url = "https://files.pythonhosted.org/packages/5f/67/cba3777620cdacb99102da4042883709c41c709f4b6323c10781a9c3aa34/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_x86_64.whl", hash = "sha256:0a759da5dea5c0ea10fd307de75cdeb59e7ea4fcb8add0924859b944babf1112", size = 200941980, upload-time = "2025-09-04T08:33:22.767Z" }, + { url = "https://files.pythonhosted.org/packages/06/1e/b8b7c2f4099a37b96af5c9bb158632ea9e5d9d27d7391d7eb8fc45236674/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:7556d9eca156e18184b94947ade0fba5bb47d69cec46bf8660fd2c71a4b48b73", size = 216561367, upload-time = "2024-11-20T17:44:54.824Z" }, + { url = "https://files.pythonhosted.org/packages/43/ac/64c4316ba163e8217a99680c7605f779accffc6a4bcd0c778c12948d3707/nvidia_cusparse_cu12-12.5.4.2-py3-none-manylinux2014_x86_64.whl", hash = "sha256:23749a6571191a215cb74d1cdbff4a86e7b19f1200c071b3fcf844a5bea23a2f", size = 216561357, upload-time = "2024-10-01T17:06:29.861Z" }, ] [[package]] -name = "nvidia-cusparse" -version = "12.6.3.3" +name = "nvidia-cusparselt-cu12" +version = "0.6.3" source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "nvidia-nvjitlink", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, -] wheels = [ - { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, - { url = "https://files.pythonhosted.org/packages/fa/18/623c77619c31d62efd55302939756966f3ecc8d724a14dab2b75f1508850/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:2b3c89c88d01ee0e477cb7f82ef60a11a4bcd57b6b87c33f789350b59759360b", size = 145942937, upload-time = "2025-09-04T08:33:58.029Z" }, + { url = "https://files.pythonhosted.org/packages/3b/9a/72ef35b399b0e183bc2e8f6f558036922d453c4d8237dab26c666a04244b/nvidia_cusparselt_cu12-0.6.3-py3-none-manylinux2014_x86_64.whl", hash = "sha256:e5c8a26c36445dd2e6812f1177978a24e2d37cacce7e090f297a688d1ec44f46", size = 156785796, upload-time = "2024-10-15T21:29:17.709Z" }, ] [[package]] -name = "nvidia-cusparselt-cu13" -version = "0.8.1" +name = "nvidia-nccl-cu12" +version = "2.26.2" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/46/e1/cdc1797eadf82d3a9a575a19b33fdc871a97edbec42c00b5b5e914f4aff4/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl", hash = "sha256:4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f", size = 221051344, upload-time = "2025-09-05T18:49:51.289Z" }, - { url = "https://files.pythonhosted.org/packages/34/7d/2661f2fb3ac4302f3a246f5fc030213ac60c1fe0bce84f9783dbd831dbb7/nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_x86_64.whl", hash = "sha256:786ce87568c303fadb5afcc7102d454cd3040d75f6f8626f5db460d1871f4dd0", size = 170148586, upload-time = "2025-09-05T18:50:50.248Z" }, + { url = "https://files.pythonhosted.org/packages/67/ca/f42388aed0fddd64ade7493dbba36e1f534d4e6fdbdd355c6a90030ae028/nvidia_nccl_cu12-2.26.2-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:694cf3879a206553cc9d7dbda76b13efaf610fdb70a50cba303de1b0d1530ac6", size = 201319755, upload-time = "2025-03-13T00:29:55.296Z" }, ] [[package]] -name = "nvidia-nccl-cu13" -version = "2.29.7" +name = "nvidia-nvjitlink-cu12" +version = "12.6.85" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/72/0d/daf50d44177ee0cbc7ff0a0c91eb5ff676c82be42f9a970bc7597f440c3a/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_aarch64.whl", hash = "sha256:674a12383e3c38a1bcccae7d4f3633b37852230b6047883cb2f4c2d1b36d9bf5", size = 206014712, upload-time = "2026-03-03T05:34:20.843Z" }, - { url = "https://files.pythonhosted.org/packages/67/f4/58e4e91b6919367c7aafb8e36fce9aad1a3047e536bf7e2fd560927d3a4c/nvidia_nccl_cu13-2.29.7-py3-none-manylinux_2_18_x86_64.whl", hash = "sha256:edd81538446786ec3b73972543e53bb43bcaf0bfc8ef76cb679fcc390ffe136d", size = 205976000, upload-time = "2026-03-03T05:36:24.472Z" }, + { url = "https://files.pythonhosted.org/packages/9d/d7/c5383e47c7e9bf1c99d5bd2a8c935af2b6d705ad831a7ec5c97db4d82f4f/nvidia_nvjitlink_cu12-12.6.85-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:eedc36df9e88b682efe4309aa16b5b4e78c2407eac59e8c10a6a47535164369a", size = 19744971, upload-time = "2024-11-20T17:46:53.366Z" }, ] [[package]] -name = "nvidia-nvjitlink" -version = "13.3.33" +name = "nvidia-nvtx-cu12" +version = "12.6.77" source = { registry = "https://pypi.org/simple" } wheels = [ - { url = "https://files.pythonhosted.org/packages/f0/ee/580ca6f29dcab0221db8706badca1bbbb084f1975c4d4e83329c3a7e31f0/nvidia_nvjitlink-13.3.33-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:26a6de7fb4c8fdaa7703d3dad720d6d427ddfea5c48a528fd97c11733ad830e5", size = 40742423, upload-time = "2026-05-26T16:54:51.613Z" }, - { url = "https://files.pythonhosted.org/packages/69/30/45414e35ff2eee7db3da037e5707037ccf9d2b5218ffbdb055ea4d5aa98a/nvidia_nvjitlink-13.3.33-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:ce48b37dfeb3cb1eae4cf85adacb47d7a6539ea2272870c9a3628ce275c2037e", size = 39168635, upload-time = "2026-05-26T16:54:13.906Z" }, -] - -[[package]] -name = "nvidia-nvshmem-cu13" -version = "3.4.5" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/dc/0f/05cc9c720236dcd2db9c1ab97fff629e96821be2e63103569da0c9b72f19/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:6dc2a197f38e5d0376ad52cd1a2a3617d3cdc150fd5966f4aee9bcebb1d68fe9", size = 60215947, upload-time = "2025-09-06T00:32:20.022Z" }, - { url = "https://files.pythonhosted.org/packages/3c/35/a9bf80a609e74e3b000fef598933235c908fcefcef9026042b8e6dfde2a9/nvidia_nvshmem_cu13-3.4.5-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:290f0a2ee94c9f3687a02502f3b9299a9f9fe826e6d0287ee18482e78d495b80", size = 60412546, upload-time = "2025-09-06T00:32:41.564Z" }, -] - -[[package]] -name = "nvidia-nvtx" -version = "13.0.85" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/c2/f3/d86c845465a2723ad7e1e5c36dcd75ddb82898b3f53be47ebd429fb2fa5d/nvidia_nvtx-13.0.85-py3-none-manylinux1_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:4936d1d6780fbe68db454f5e72a42ff64d1fd6397df9f363ae786930fd5c1cd4", size = 148047, upload-time = "2025-09-04T08:29:01.761Z" }, - { url = "https://files.pythonhosted.org/packages/a8/64/3708a90d1ebe202ffdeb7185f878a3c84d15c2b2c31858da2ce0583e2def/nvidia_nvtx-13.0.85-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:cb7780edb6b14107373c835bf8b72e7a178bac7367e23da7acb108f973f157a6", size = 148878, upload-time = "2025-09-04T08:28:53.627Z" }, + { url = "https://files.pythonhosted.org/packages/56/9a/fff8376f8e3d084cd1530e1ef7b879bb7d6d265620c95c1b322725c694f4/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:b90bed3df379fa79afbd21be8e04a0314336b8ae16768b58f2d34cb1d04cd7d2", size = 89276, upload-time = "2024-11-20T17:38:27.621Z" }, + { url = "https://files.pythonhosted.org/packages/9e/4e/0d0c945463719429b7bd21dece907ad0bde437a2ff12b9b12fee94722ab0/nvidia_nvtx_cu12-12.6.77-py3-none-manylinux2014_x86_64.whl", hash = "sha256:6574241a3ec5fdc9334353ab8c479fe75841dbe8f4532a8fc97ce63503330ba1", size = 89265, upload-time = "2024-10-01T17:00:38.172Z" }, ] [[package]] @@ -3961,45 +3865,49 @@ wheels = [ [[package]] name = "torch" -version = "2.13.0" +version = "2.7.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, - { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux'" }, { name = "filelock" }, { name = "fsspec" }, { name = "jinja2" }, { name = "networkx" }, - { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, - { name = "setuptools" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufile-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparselt-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "setuptools", marker = "python_full_version >= '3.12'" }, { name = "sympy" }, - { name = "triton", marker = "python_full_version < '3.15' and sys_platform == 'linux'" }, + { name = "triton", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, { name = "typing-extensions" }, ] wheels = [ - { url = "https://files.pythonhosted.org/packages/5b/fe/cba54dc58523434919b66f13a667e36e436deddd77ca519e96553617d4ec/torch-2.13.0-cp311-cp311-macosx_14_0_arm64.whl", hash = "sha256:e76f9bcecc52b8ff711239a2f7547d5353df95878ab232f0773c1d95928b92f8", size = 111187938, upload-time = "2026-07-08T16:05:17.065Z" }, - { url = "https://files.pythonhosted.org/packages/c2/59/1e3160e18e12aa3038390efab3ce02b36a9d4d6a527ecdd8520dca2e68d8/torch-2.13.0-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:092790c696a760c729fd5722835f50b9d81fd7c8f141571f3f3cf4081a8f664c", size = 427199369, upload-time = "2026-07-08T16:04:51.054Z" }, - { url = "https://files.pythonhosted.org/packages/01/79/1f2d34ad7034ee1c7ffc1cf8bf0f8213af2a81df6ecdb3997ecec107c09d/torch-2.13.0-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:60fcdcb2f3876e21146cb4524ef06397d727ca9ad5f020818547e25075fe3cb7", size = 526574961, upload-time = "2026-07-08T16:04:07.075Z" }, - { url = "https://files.pythonhosted.org/packages/6c/fd/0f2ce40f58aefbdb3392f9acce3c8171940943ae2d661f70558bfa73befb/torch-2.13.0-cp311-cp311-win_amd64.whl", hash = "sha256:a0d8b11f16a48d60e2015d8213aa0390744cbebb98e58b62b3514dddc656e330", size = 122015870, upload-time = "2026-07-08T16:05:27.59Z" }, - { url = "https://files.pythonhosted.org/packages/c4/3a/ed0f4d4d1dcde03bced7aac9a28e800abcdc0cbd06b6775044c9fbd877b7/torch-2.13.0-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:2fe228aba290d14b9f31b049be550dbd469c3fd3013d7a19705b30454da97027", size = 111213045, upload-time = "2026-07-08T16:05:22.997Z" }, - { url = "https://files.pythonhosted.org/packages/df/a9/f6a2a4d763ff1df02e9a64c477029db614295bc9367f4131223791ccc243/torch-2.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:572df8be8ffb4599c88cbd6a0726f1f854f4da65d2e3c09f0e2c2283333cd6d4", size = 427210998, upload-time = "2026-07-08T16:04:37.708Z" }, - { url = "https://files.pythonhosted.org/packages/f3/82/fea946351658e6534db52d2cc12bc53087cbf87f9440c5f180f367c1950b/torch-2.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:796633c4cdf0fe2cdced72d8f88f22e73dbcfce83132763162f6d4bff13b820b", size = 526605292, upload-time = "2026-07-08T16:04:22.81Z" }, - { url = "https://files.pythonhosted.org/packages/21/d6/e8f3c6f7e01f626f77259de9860d2a78bc84c40539e28e79b7e98b0bb659/torch-2.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:024c6cc0c1b085f2f91f20a3dc27b0471d021c31ce84b81be3afdc39f791fd9d", size = 122057313, upload-time = "2026-07-08T16:03:53.43Z" }, - { url = "https://files.pythonhosted.org/packages/0d/fa/c1c10b7aff4a9a3e8956d4f0a5f468fa6db7abc3208805719076772b4833/torch-2.13.0-cp313-cp313-macosx_14_0_arm64.whl", hash = "sha256:33449899ce5496c1b84b4853179d94fd102028ae1407314d9fb956bb79e70d09", size = 111213743, upload-time = "2026-07-08T16:03:28.579Z" }, - { url = "https://files.pythonhosted.org/packages/11/18/9ecb37b56293a0be8d80f810bf672a72fe7e02f8b475d5ef1b9bf8a0d748/torch-2.13.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:1e09d6a722504957c694faceca843acde562786df1144ebcc5a74075ec7f6005", size = 427213008, upload-time = "2026-07-08T16:03:44.106Z" }, - { url = "https://files.pythonhosted.org/packages/d4/5a/7c50ba1b7b713d71d34669c6d13dab0a11531a3eceb0307a5162dbfec0f7/torch-2.13.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:a3a9a21312872af8a26950b2c15680335a386a1f56ed03e780653d78b9607e9e", size = 526602329, upload-time = "2026-07-08T16:03:12.649Z" }, - { url = "https://files.pythonhosted.org/packages/91/3d/e7adcc6aaf36961cd18f56cf8ad0f3058c3a5c84ccf391762176c94581b8/torch-2.13.0-cp313-cp313-win_amd64.whl", hash = "sha256:49b58f1e2c52440abb6f17c28f0335fe6c6d01ad1a7f55b0183b81e4b34d64e6", size = 122057920, upload-time = "2026-07-08T16:03:01.808Z" }, - { url = "https://files.pythonhosted.org/packages/36/76/6dcc7f0c07052102dd36f83cbc5800842a909c8c3fbf1a7f8a5844954de9/torch-2.13.0-cp314-cp314-macosx_14_0_arm64.whl", hash = "sha256:d849b390e07d8d333ce8ecaf91b273c656c598379a19c9acf1318a883f6b391c", size = 111227066, upload-time = "2026-07-08T16:03:33.6Z" }, - { url = "https://files.pythonhosted.org/packages/e9/09/2c10e8cd0e00fa5d23c052df6ce467eaa7182399f5e0f824f1e4ff42ccae/torch-2.13.0-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:a3893dc2da0a972a8ca5d698c85a9f967559ac5f8ee1797b77408aa8734d073c", size = 427226309, upload-time = "2026-07-08T16:02:53.127Z" }, - { url = "https://files.pythonhosted.org/packages/76/c6/22c2102bbef14ca6a6cb4c20e42f088e49c5f812be4e160ae57502e325f9/torch-2.13.0-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:49f1ea385c754e54919408a9bb3b5a72b0b755bbe2c916c1d6f70afbec4908a2", size = 526614507, upload-time = "2026-07-08T16:02:16.441Z" }, - { url = "https://files.pythonhosted.org/packages/2b/0c/7d1deb6bce5bc3e6042caf39100ac768eba3b9a098e1dddd16f75bd6489b/torch-2.13.0-cp314-cp314-win_amd64.whl", hash = "sha256:4f8573e3ce9ebcd53fe922f01077a6085ccdfbe5f12fd215883a9d87d7a744fd", size = 122051871, upload-time = "2026-07-08T16:03:23.521Z" }, - { url = "https://files.pythonhosted.org/packages/f4/ce/aa8b7f9949d32e0f2f624f342bc3b48112c1b8a130288465938bc83bcbf9/torch-2.13.0-cp314-cp314t-macosx_14_0_arm64.whl", hash = "sha256:c28def70706c2f9ecc752574766e8ae4da9b810ab6676b611166761a78a9f1e1", size = 111537025, upload-time = "2026-07-08T16:02:44.28Z" }, - { url = "https://files.pythonhosted.org/packages/69/d1/491e3a0389430946145888b0203f2b6a759ce2a61481b96a85c2da4f2ced/torch-2.13.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:31061ff56ed8fbf26c749806905aeb749ebeb819810fd5d52508aa5afd90dddc", size = 427219769, upload-time = "2026-07-08T16:02:31.18Z" }, - { url = "https://files.pythonhosted.org/packages/9a/1d/38006e045bf0a1fc28ef01e757c554e59e59a8770c284bc4f47b14e60441/torch-2.13.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:cc26eead4cf51d0b544e31e364dcf000846549c273bd148936fe9d24d29acb92", size = 526571320, upload-time = "2026-07-08T16:01:59.348Z" }, - { url = "https://files.pythonhosted.org/packages/56/94/655c91992a882bd5071aa0b5d22a07dbb130d801e872be97c0b627a7c693/torch-2.13.0-cp314-cp314t-win_amd64.whl", hash = "sha256:a7de8a313090dc5c7d7ba4bfe5c3be222528f9a4dba1acc83bddb1157360c4b8", size = 122306773, upload-time = "2026-07-08T16:02:39.832Z" }, + { url = "https://files.pythonhosted.org/packages/11/56/2eae3494e3d375533034a8e8cf0ba163363e996d85f0629441fa9d9843fe/torch-2.7.1-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:236f501f2e383f1cb861337bdf057712182f910f10aeaf509065d54d339e49b2", size = 99093039, upload-time = "2025-06-04T17:39:06.963Z" }, + { url = "https://files.pythonhosted.org/packages/e5/94/34b80bd172d0072c9979708ccd279c2da2f55c3ef318eceec276ab9544a4/torch-2.7.1-cp311-cp311-manylinux_2_28_x86_64.whl", hash = "sha256:06eea61f859436622e78dd0cdd51dbc8f8c6d76917a9cf0555a333f9eac31ec1", size = 821174704, upload-time = "2025-06-04T17:37:03.799Z" }, + { url = "https://files.pythonhosted.org/packages/50/9e/acf04ff375b0b49a45511c55d188bcea5c942da2aaf293096676110086d1/torch-2.7.1-cp311-cp311-win_amd64.whl", hash = "sha256:8273145a2e0a3c6f9fd2ac36762d6ee89c26d430e612b95a99885df083b04e52", size = 216095937, upload-time = "2025-06-04T17:39:24.83Z" }, + { url = "https://files.pythonhosted.org/packages/5b/2b/d36d57c66ff031f93b4fa432e86802f84991477e522adcdffd314454326b/torch-2.7.1-cp311-none-macosx_11_0_arm64.whl", hash = "sha256:aea4fc1bf433d12843eb2c6b2204861f43d8364597697074c8d38ae2507f8730", size = 68640034, upload-time = "2025-06-04T17:39:17.989Z" }, + { url = "https://files.pythonhosted.org/packages/87/93/fb505a5022a2e908d81fe9a5e0aa84c86c0d5f408173be71c6018836f34e/torch-2.7.1-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:27ea1e518df4c9de73af7e8a720770f3628e7f667280bce2be7a16292697e3fa", size = 98948276, upload-time = "2025-06-04T17:39:12.852Z" }, + { url = "https://files.pythonhosted.org/packages/56/7e/67c3fe2b8c33f40af06326a3d6ae7776b3e3a01daa8f71d125d78594d874/torch-2.7.1-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:c33360cfc2edd976c2633b3b66c769bdcbbf0e0b6550606d188431c81e7dd1fc", size = 821025792, upload-time = "2025-06-04T17:34:58.747Z" }, + { url = "https://files.pythonhosted.org/packages/a1/37/a37495502bc7a23bf34f89584fa5a78e25bae7b8da513bc1b8f97afb7009/torch-2.7.1-cp312-cp312-win_amd64.whl", hash = "sha256:d8bf6e1856ddd1807e79dc57e54d3335f2b62e6f316ed13ed3ecfe1fc1df3d8b", size = 216050349, upload-time = "2025-06-04T17:38:59.709Z" }, + { url = "https://files.pythonhosted.org/packages/3a/60/04b77281c730bb13460628e518c52721257814ac6c298acd25757f6a175c/torch-2.7.1-cp312-none-macosx_11_0_arm64.whl", hash = "sha256:787687087412c4bd68d315e39bc1223f08aae1d16a9e9771d95eabbb04ae98fb", size = 68645146, upload-time = "2025-06-04T17:38:52.97Z" }, + { url = "https://files.pythonhosted.org/packages/66/81/e48c9edb655ee8eb8c2a6026abdb6f8d2146abd1f150979ede807bb75dcb/torch-2.7.1-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:03563603d931e70722dce0e11999d53aa80a375a3d78e6b39b9f6805ea0a8d28", size = 98946649, upload-time = "2025-06-04T17:38:43.031Z" }, + { url = "https://files.pythonhosted.org/packages/3a/24/efe2f520d75274fc06b695c616415a1e8a1021d87a13c68ff9dce733d088/torch-2.7.1-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:d632f5417b6980f61404a125b999ca6ebd0b8b4bbdbb5fbbba44374ab619a412", size = 821033192, upload-time = "2025-06-04T17:38:09.146Z" }, + { url = "https://files.pythonhosted.org/packages/dd/d9/9c24d230333ff4e9b6807274f6f8d52a864210b52ec794c5def7925f4495/torch-2.7.1-cp313-cp313-win_amd64.whl", hash = "sha256:23660443e13995ee93e3d844786701ea4ca69f337027b05182f5ba053ce43b38", size = 216055668, upload-time = "2025-06-04T17:38:36.253Z" }, + { url = "https://files.pythonhosted.org/packages/95/bf/e086ee36ddcef9299f6e708d3b6c8487c1651787bb9ee2939eb2a7f74911/torch-2.7.1-cp313-cp313t-macosx_14_0_arm64.whl", hash = "sha256:0da4f4dba9f65d0d203794e619fe7ca3247a55ffdcbd17ae8fb83c8b2dc9b585", size = 68925988, upload-time = "2025-06-04T17:38:29.273Z" }, + { url = "https://files.pythonhosted.org/packages/69/6a/67090dcfe1cf9048448b31555af6efb149f7afa0a310a366adbdada32105/torch-2.7.1-cp313-cp313t-manylinux_2_28_aarch64.whl", hash = "sha256:e08d7e6f21a617fe38eeb46dd2213ded43f27c072e9165dc27300c9ef9570934", size = 99028857, upload-time = "2025-06-04T17:37:50.956Z" }, + { url = "https://files.pythonhosted.org/packages/90/1c/48b988870823d1cc381f15ec4e70ed3d65e043f43f919329b0045ae83529/torch-2.7.1-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:30207f672328a42df4f2174b8f426f354b2baa0b7cca3a0adb3d6ab5daf00dc8", size = 821098066, upload-time = "2025-06-04T17:37:33.939Z" }, + { url = "https://files.pythonhosted.org/packages/7b/eb/10050d61c9d5140c5dc04a89ed3257ef1a6b93e49dd91b95363d757071e0/torch-2.7.1-cp313-cp313t-win_amd64.whl", hash = "sha256:79042feca1c634aaf6603fe6feea8c6b30dfa140a6bbc0b973e2260c7e79a22e", size = 216336310, upload-time = "2025-06-04T17:36:09.862Z" }, + { url = "https://files.pythonhosted.org/packages/b1/29/beb45cdf5c4fc3ebe282bf5eafc8dfd925ead7299b3c97491900fe5ed844/torch-2.7.1-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:988b0cbc4333618a1056d2ebad9eb10089637b659eb645434d0809d8d937b946", size = 68645708, upload-time = "2025-06-04T17:34:39.852Z" }, ] [[package]] @@ -4042,19 +3950,16 @@ wheels = [ [[package]] name = "triton" -version = "3.7.1" +version = "3.3.1" source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "setuptools", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] wheels = [ - { url = "https://files.pythonhosted.org/packages/7b/f9/19d842d06a08559534fa1eaab6ca551b1bcf40f06620bddec1babaa2772d/triton-3.7.1-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4a0e1cd4c4a76370ed74a8432a53cea28716827d19e40ffc732233e35ceb3f6", size = 184664887, upload-time = "2026-06-17T20:03:42.913Z" }, - { url = "https://files.pythonhosted.org/packages/cd/5e/fce69606f7f240297f163e25539906732b199530d486ce67ae319877e821/triton-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6744957e9fd610a29680ec2346057d0c86948ed3812468670719f391e94b44a5", size = 197701306, upload-time = "2026-06-17T19:53:13.673Z" }, - { url = "https://files.pythonhosted.org/packages/94/fa/f856e24deb462d5f18bd4b5a746957862ab9b6ee5834bda60605ec348366/triton-3.7.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9497f2e696ee368862a181a90b2dcc03ca978cc4f602abd67c7d81022a6988e1", size = 184692359, upload-time = "2026-06-17T20:03:48.288Z" }, - { url = "https://files.pythonhosted.org/packages/c4/6f/fb96d15db6f36d6eae4cafb998c2e0353bf59d7c4ea1662d7497f269134a/triton-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:7e40869937a68206ec70d7f25bb7ec6433cb083f9135e1f36dbd318dc449a728", size = 197719725, upload-time = "2026-06-17T19:53:20.419Z" }, - { url = "https://files.pythonhosted.org/packages/00/42/c5089d4d9327fcd1e862c599cc2927f39418f84dd11a84cb2ccff9d4787a/triton-3.7.1-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdbfc09d9ec58bc5e68321525653220de7515c199e7a8097a97c85e62b52cd0a", size = 184694629, upload-time = "2026-06-17T20:03:53.444Z" }, - { url = "https://files.pythonhosted.org/packages/07/42/2c3ac59253ae8892b6f307875263dd23dc875cdf732d3aea40d6d41fb7cb/triton-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:58c0e131da05134a2a4788ccbcc0c1105cf0f54c8e98f19e34cd465396dc15eb", size = 197729241, upload-time = "2026-06-17T19:53:27.801Z" }, - { url = "https://files.pythonhosted.org/packages/40/71/e01aa7ad573883ed9456f130226babdec70b005e098c4d6226a6238e761b/triton-3.7.1-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:fe4ea396a06171f1f1f58cbd39c70b09294398f7dd7c620939bab54ad6f934fa", size = 184705764, upload-time = "2026-06-17T20:03:59.064Z" }, - { url = "https://files.pythonhosted.org/packages/a4/09/5683146fda6a2b569deb78ccfd8fbfea8bfe55f726b081c0a6bb18dd6f28/triton-3.7.1-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:2020153b08280415ec0da6607834e79166442147e78e144df06b508c75b186d2", size = 197729537, upload-time = "2026-06-17T19:53:35.516Z" }, - { url = "https://files.pythonhosted.org/packages/e9/f8/448220c3092019f9fdfab39ec47985968181d67da34b44f6a7f6280a5cbb/triton-3.7.1-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c58e4c61f0c73b5dba3b5d19b4a7093c32f90dc18b2a7f121a7c16ccd31107b7", size = 184814760, upload-time = "2026-06-17T20:04:04.984Z" }, - { url = "https://files.pythonhosted.org/packages/f0/ac/229b7d4589d2e5937310e72c6d46e89599d16a4a12b479ffa1499fee8eb8/triton-3.7.1-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:10ba85fa2cca4a2fbdeb36bf1cb082f2c252bda55bf9fccd74f65ec5bc647e68", size = 197824404, upload-time = "2026-06-17T19:53:42.772Z" }, + { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937, upload-time = "2025-05-29T23:39:44.182Z" }, + { url = "https://files.pythonhosted.org/packages/24/5f/950fb373bf9c01ad4eb5a8cd5eaf32cdf9e238c02f9293557a2129b9c4ac/triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43", size = 155669138, upload-time = "2025-05-29T23:39:51.771Z" }, + { url = "https://files.pythonhosted.org/packages/74/1f/dfb531f90a2d367d914adfee771babbd3f1a5b26c3f5fbc458dee21daa78/triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240", size = 155673035, upload-time = "2025-05-29T23:40:02.468Z" }, + { url = "https://files.pythonhosted.org/packages/28/71/bd20ffcb7a64c753dc2463489a61bf69d531f308e390ad06390268c4ea04/triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42", size = 155735832, upload-time = "2025-05-29T23:40:10.522Z" }, ] [[package]]