390 lines
17 KiB
Python
390 lines
17 KiB
Python
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import gzip
|
|
import http.client
|
|
import sys
|
|
import tempfile
|
|
import types
|
|
import shutil
|
|
import unittest
|
|
import zipfile
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
from airfrans_frontier.runtime import remove_pythonpath_entries
|
|
|
|
remove_pythonpath_entries()
|
|
|
|
from airfrans_frontier.raw.bounded_public import prepare_public_airfrans_processed_hf_bounded
|
|
from airfrans_frontier.raw.public import (
|
|
HttpRangeReader,
|
|
_extract_remote_case_members,
|
|
_read_zip_central_directory,
|
|
_remote_archive_case_members,
|
|
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):
|
|
def test_prepare_public_hf_skips_when_dataset_already_published(self) -> None:
|
|
class FakeApi:
|
|
def __init__(self, token=None):
|
|
self.token = token
|
|
|
|
def list_repo_files(self, *, repo_id: str, repo_type: str):
|
|
assert repo_id == "owner/airfrans-processed"
|
|
assert repo_type == "dataset"
|
|
return [
|
|
"processed/full/case_000.npz",
|
|
"processed/full/case_001.npz",
|
|
"processed/full/hf_dataset_manifest.json",
|
|
]
|
|
|
|
fake_module = types.SimpleNamespace(HfApi=FakeApi)
|
|
with tempfile.TemporaryDirectory() as tmp, patch.dict(sys.modules, {"huggingface_hub": fake_module}), patch.dict(
|
|
"os.environ", {"HF_TOKEN": "token"}
|
|
):
|
|
report = ensure_public_airfrans_processed_hf(
|
|
repo_id="owner/airfrans-processed",
|
|
path_in_repo="processed/full",
|
|
work_dir=Path(tmp) / "work",
|
|
output_dir=Path(tmp) / "out",
|
|
min_cases=2,
|
|
)
|
|
|
|
self.assertTrue(report["ok"])
|
|
self.assertEqual(report["phase"], "already_published")
|
|
self.assertEqual(report["npz_file_count"], 2)
|
|
self.assertTrue(report["has_manifest"])
|
|
|
|
def test_extract_of_dataset_finds_public_archive_root(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")
|
|
root = extract_of_dataset(archive, tmp_path / "raw", min_cases=1)
|
|
|
|
self.assertEqual(root.name, "OF_dataset")
|
|
|
|
def test_extract_of_dataset_rejects_zip_slip_paths(self) -> None:
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
tmp_path = Path(tmp)
|
|
archive = tmp_path / "bad.zip"
|
|
with zipfile.ZipFile(archive, "w") as zf:
|
|
zf.writestr("../escape.txt", "bad")
|
|
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_range_streaming_extracts_contiguous_case_with_one_payload_read(self) -> None:
|
|
class CountingRangeReader:
|
|
def __init__(self, path: Path) -> None:
|
|
self._path = path
|
|
self.size = path.stat().st_size
|
|
self.bytes_read = 0
|
|
self.calls: list[tuple[int, int]] = []
|
|
|
|
def read_range(self, start: int, length: int) -> bytes:
|
|
self.calls.append((start, length))
|
|
with self._path.open("rb") as handle:
|
|
handle.seek(start)
|
|
data = handle.read(length)
|
|
self.bytes_read += len(data)
|
|
return data
|
|
|
|
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])
|
|
reader = CountingRangeReader(archive)
|
|
members = _read_zip_central_directory(reader)
|
|
case_members = _remote_archive_case_members(members)[case_name]
|
|
|
|
reader.calls.clear()
|
|
_extract_remote_case_members(reader, case_members, tmp_path / "streaming_raw")
|
|
|
|
self.assertEqual(len(reader.calls), 1)
|
|
self.assertTrue((tmp_path / "streaming_raw" / case_name / "constant" / "transportProperties").is_file())
|
|
|
|
def test_http_range_reader_retries_timeout_before_failing_run(self) -> None:
|
|
class FakeResponse:
|
|
status = 206
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, traceback):
|
|
return False
|
|
|
|
def read(self) -> bytes:
|
|
return b"ok"
|
|
|
|
reader = HttpRangeReader.__new__(HttpRangeReader)
|
|
reader.url = "https://example.test/OF_dataset.zip"
|
|
reader.size = 10
|
|
reader.bytes_read = 0
|
|
|
|
with patch("airfrans_frontier.raw.public.urllib.request.urlopen", side_effect=[TimeoutError("timed out"), FakeResponse()]), patch(
|
|
"airfrans_frontier.raw.public.time.sleep"
|
|
) as sleep:
|
|
data = reader.read_range(2, 2)
|
|
|
|
self.assertEqual(data, b"ok")
|
|
self.assertEqual(reader.bytes_read, 2)
|
|
sleep.assert_called_once_with(2.0)
|
|
|
|
def test_http_range_reader_retries_incomplete_body_before_failing_run(self) -> None:
|
|
class FakeResponse:
|
|
status = 206
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc, traceback):
|
|
return False
|
|
|
|
def read(self) -> bytes:
|
|
return b"ok"
|
|
|
|
reader = HttpRangeReader.__new__(HttpRangeReader)
|
|
reader.url = "https://example.test/OF_dataset.zip"
|
|
reader.size = 10
|
|
reader.bytes_read = 0
|
|
|
|
with patch(
|
|
"airfrans_frontier.raw.public.urllib.request.urlopen",
|
|
side_effect=[http.client.IncompleteRead(b"pa", 2), FakeResponse()],
|
|
), patch("airfrans_frontier.raw.public.time.sleep") as sleep:
|
|
data = reader.read_range(2, 2)
|
|
|
|
self.assertEqual(data, b"ok")
|
|
self.assertEqual(reader.bytes_read, 2)
|
|
sleep.assert_called_once_with(2.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)
|
|
|
|
|
|
def test_bounded_public_hf_uploads_verified_chunks_and_cleans_staging(self) -> None:
|
|
uploaded: dict[str, tuple[int, str, bytes]] = {}
|
|
commits: list[tuple[str, tuple[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:
|
|
self.token = token
|
|
|
|
def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None:
|
|
self.repo_id = repo_id
|
|
|
|
def create_commit(self, *, repo_id: str, repo_type: str, operations, commit_message: str):
|
|
paths: list[str] = []
|
|
for operation in operations:
|
|
payload = Path(operation.path_or_fileobj).read_bytes()
|
|
uploaded[operation.path_in_repo] = (
|
|
len(payload),
|
|
hashlib.sha256(payload).hexdigest(),
|
|
payload,
|
|
)
|
|
paths.append(operation.path_in_repo)
|
|
commits.append((commit_message, tuple(paths)))
|
|
return types.SimpleNamespace(commit_url=f"https://huggingface.co/datasets/{repo_id}/commit/{len(commits)}", oid=str(len(commits)))
|
|
|
|
def repo_info(self, *, repo_id: str, repo_type: str, files_metadata: bool):
|
|
siblings = [
|
|
types.SimpleNamespace(rfilename=path, size=size, lfs={"sha256": sha})
|
|
for path, (size, sha, _payload) in uploaded.items()
|
|
]
|
|
return types.SimpleNamespace(siblings=siblings)
|
|
|
|
def list_repo_files(self, *, repo_id: str, repo_type: str):
|
|
return sorted(uploaded)
|
|
|
|
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 = ["airFoil2D_SST_10.0_5.0_0012", "airFoil2D_SST_11.0_5.0_0012"]
|
|
write_minimal_airfrans_archive(archive, case_names)
|
|
|
|
report = prepare_public_airfrans_processed_hf_bounded(
|
|
repo_id="owner/repo",
|
|
path_in_repo="processed/full",
|
|
work_dir=tmp_path / "work",
|
|
output_dir=tmp_path / "staging",
|
|
source_url=str(archive),
|
|
min_cases=2,
|
|
chunk_max_bytes=1,
|
|
train_cases=1,
|
|
val_cases=1,
|
|
test_cases=0,
|
|
split_seed=123,
|
|
)
|
|
|
|
state = json.loads((tmp_path / "work" / "bounded_prepare_state.json").read_text())
|
|
final_manifest = json.loads((tmp_path / "work" / "chunk_manifests" / "hf_dataset_manifest.json").read_text())
|
|
|
|
self.assertTrue(report["ok"])
|
|
self.assertEqual(report["processed_case_count"], 2)
|
|
self.assertEqual(report["chunks_uploaded"], 2)
|
|
self.assertEqual(state["phase"], "published")
|
|
self.assertEqual(final_manifest["case_count"], 2)
|
|
self.assertEqual(final_manifest["feature_names"][0], "x")
|
|
self.assertEqual(final_manifest["target_names"], ["velocity_x", "velocity_y", "pressure", "turbulent_viscosity"])
|
|
self.assertEqual(final_manifest["split_compatibility"]["train_cases"], 1)
|
|
self.assertFalse(any((tmp_path / "staging").glob("*.npz")))
|
|
self.assertFalse((tmp_path / "work" / "bounded_raw_scratch").exists())
|
|
self.assertIn("processed/full/hf_dataset_manifest.json", uploaded)
|
|
self.assertEqual(sum(1 for path in uploaded if path.endswith(".npz")), 2)
|
|
self.assertEqual(sum(1 for message, _paths in commits if message.startswith("Upload bounded AirfRANS processed chunk")), 2)
|
|
|
|
def test_bounded_public_hf_rejects_remote_checksum_mismatch(self) -> None:
|
|
uploaded: dict[str, tuple[int, 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:
|
|
self.token = token
|
|
|
|
def create_repo(self, *, repo_id: str, repo_type: str, private: bool, exist_ok: bool) -> None:
|
|
pass
|
|
|
|
def create_commit(self, *, repo_id: str, repo_type: str, operations, commit_message: str):
|
|
for operation in operations:
|
|
payload = Path(operation.path_or_fileobj).read_bytes()
|
|
uploaded[operation.path_in_repo] = (len(payload), hashlib.sha256(payload).hexdigest())
|
|
return types.SimpleNamespace(commit_url="https://huggingface.co/datasets/owner/repo/commit/bad", oid="bad")
|
|
|
|
def repo_info(self, *, repo_id: str, repo_type: str, files_metadata: bool):
|
|
siblings = [
|
|
types.SimpleNamespace(rfilename=path, size=size, lfs={"sha256": "0" * 64})
|
|
for path, (size, _sha) in uploaded.items()
|
|
]
|
|
return types.SimpleNamespace(siblings=siblings)
|
|
|
|
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"
|
|
write_minimal_airfrans_archive(archive, ["airFoil2D_SST_10.0_5.0_0012"])
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "HF upload verification failed"):
|
|
prepare_public_airfrans_processed_hf_bounded(
|
|
repo_id="owner/repo",
|
|
path_in_repo="processed/full",
|
|
work_dir=tmp_path / "work",
|
|
output_dir=tmp_path / "staging",
|
|
source_url=str(archive),
|
|
min_cases=1,
|
|
chunk_max_bytes=1024,
|
|
train_cases=1,
|
|
val_cases=0,
|
|
test_cases=0,
|
|
split_seed=123,
|
|
)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|