372 lines
17 KiB
Python
372 lines
17 KiB
Python
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()
|