254 lines
10 KiB
Python
254 lines
10 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
import time
|
||
|
|
import urllib.error
|
||
|
|
import urllib.request
|
||
|
|
import zipfile
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from airfrans_frontier.training.data_sources import publish_processed_dataset
|
||
|
|
|
||
|
|
PUBLIC_OF_DATASET_URL = "https://data.isir.upmc.fr/extrality/NeurIPS_2022/OF_dataset.zip"
|
||
|
|
DEFAULT_PUBLIC_WORK_DIR = Path("artifacts/public_airfrans")
|
||
|
|
DEFAULT_PUBLIC_OUTPUT_DIR = Path("artifacts/data_cache/airfrans_processed/processed/full")
|
||
|
|
|
||
|
|
def ensure_public_airfrans_processed_hf(
|
||
|
|
*,
|
||
|
|
repo_id: str,
|
||
|
|
path_in_repo: str = "processed/full",
|
||
|
|
work_dir: str | Path = DEFAULT_PUBLIC_WORK_DIR,
|
||
|
|
output_dir: str | Path = DEFAULT_PUBLIC_OUTPUT_DIR,
|
||
|
|
source_url: str = PUBLIC_OF_DATASET_URL,
|
||
|
|
min_cases: int = 1000,
|
||
|
|
private: bool = False,
|
||
|
|
force: bool = False,
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
if min_cases <= 0:
|
||
|
|
raise ValueError("min_cases must be positive")
|
||
|
|
prefix = path_in_repo.strip("/")
|
||
|
|
started = time.time()
|
||
|
|
existing = _hf_dataset_status(repo_id=repo_id, path_in_repo=prefix)
|
||
|
|
if not force and existing["npz_file_count"] >= min_cases and existing["has_manifest"]:
|
||
|
|
return {
|
||
|
|
"ok": True,
|
||
|
|
"phase": "already_published",
|
||
|
|
"repo_id": repo_id,
|
||
|
|
"repo_type": "dataset",
|
||
|
|
"repo_url": f"https://huggingface.co/datasets/{repo_id}",
|
||
|
|
"path_in_repo": prefix,
|
||
|
|
"min_cases": min_cases,
|
||
|
|
"elapsed_seconds": time.time() - started,
|
||
|
|
**existing,
|
||
|
|
}
|
||
|
|
|
||
|
|
work_root = Path(work_dir).expanduser()
|
||
|
|
output_root = Path(output_dir).expanduser()
|
||
|
|
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)
|
||
|
|
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)
|
||
|
|
publish = publish_processed_dataset(
|
||
|
|
data_root=output_root,
|
||
|
|
repo_id=repo_id,
|
||
|
|
path_in_repo=prefix,
|
||
|
|
private=private,
|
||
|
|
manifest_out=output_root / "hf_dataset_manifest.json",
|
||
|
|
)
|
||
|
|
final = _hf_dataset_status(repo_id=repo_id, path_in_repo=prefix)
|
||
|
|
if final["npz_file_count"] < min_cases:
|
||
|
|
raise RuntimeError(f"Published dataset has {final['npz_file_count']} .npz files under {prefix}; expected at least {min_cases}")
|
||
|
|
if not final["has_manifest"]:
|
||
|
|
raise RuntimeError(f"Published dataset is missing hf_dataset_manifest.json under {prefix}")
|
||
|
|
return {
|
||
|
|
"ok": True,
|
||
|
|
"phase": "published",
|
||
|
|
"repo_id": repo_id,
|
||
|
|
"repo_type": "dataset",
|
||
|
|
"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),
|
||
|
|
"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,
|
||
|
|
"publish": publish,
|
||
|
|
"elapsed_seconds": time.time() - started,
|
||
|
|
**final,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def download_file(url: str, destination: str | Path, *, chunk_size: int = 16 * 1024 * 1024) -> dict[str, Any]:
|
||
|
|
path = Path(destination).expanduser()
|
||
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
expected_size = _remote_content_length(url)
|
||
|
|
existing_size = path.stat().st_size if path.exists() else 0
|
||
|
|
if expected_size is not None and existing_size == expected_size:
|
||
|
|
return {"url": url, "path": str(path), "bytes": existing_size, "resumed": False, "skipped": True}
|
||
|
|
|
||
|
|
headers: dict[str, str] = {}
|
||
|
|
mode = "wb"
|
||
|
|
resumed = False
|
||
|
|
if expected_size is not None and 0 < existing_size < expected_size:
|
||
|
|
headers["Range"] = f"bytes={existing_size}-"
|
||
|
|
mode = "ab"
|
||
|
|
resumed = True
|
||
|
|
|
||
|
|
print(
|
||
|
|
f"download_airfrans_zip url={url} path={path} existing_bytes={existing_size} expected_bytes={expected_size}",
|
||
|
|
flush=True,
|
||
|
|
)
|
||
|
|
request = urllib.request.Request(url, headers=headers)
|
||
|
|
try:
|
||
|
|
response = urllib.request.urlopen(request, timeout=60)
|
||
|
|
except urllib.error.HTTPError as exc:
|
||
|
|
if exc.code == 416 and expected_size is not None and existing_size >= expected_size:
|
||
|
|
return {"url": url, "path": str(path), "bytes": existing_size, "resumed": False, "skipped": True}
|
||
|
|
raise
|
||
|
|
with response:
|
||
|
|
if resumed and getattr(response, "status", None) != 206:
|
||
|
|
mode = "wb"
|
||
|
|
resumed = False
|
||
|
|
existing_size = 0
|
||
|
|
written = existing_size
|
||
|
|
next_report = ((written // 1_000_000_000) + 1) * 1_000_000_000
|
||
|
|
with path.open(mode) as handle:
|
||
|
|
while True:
|
||
|
|
chunk = response.read(chunk_size)
|
||
|
|
if not chunk:
|
||
|
|
break
|
||
|
|
handle.write(chunk)
|
||
|
|
written += len(chunk)
|
||
|
|
if written >= next_report:
|
||
|
|
print(f"downloaded_airfrans_zip_bytes={written}", flush=True)
|
||
|
|
next_report += 1_000_000_000
|
||
|
|
final_size = path.stat().st_size
|
||
|
|
if expected_size is not None and final_size != expected_size:
|
||
|
|
raise RuntimeError(f"Downloaded {final_size} bytes from {url}, expected {expected_size}")
|
||
|
|
return {"url": url, "path": str(path), "bytes": final_size, "resumed": resumed, "skipped": False}
|
||
|
|
|
||
|
|
|
||
|
|
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()
|
||
|
|
root.mkdir(parents=True, exist_ok=True)
|
||
|
|
existing = _find_of_dataset_root(root)
|
||
|
|
if existing is not None and _case_count(existing) >= min_cases:
|
||
|
|
return existing
|
||
|
|
print(f"extract_airfrans_zip archive={archive} root={root}", flush=True)
|
||
|
|
with zipfile.ZipFile(archive) as zf:
|
||
|
|
members = zf.infolist()
|
||
|
|
for index, member in enumerate(members, start=1):
|
||
|
|
_safe_extract_member(zf, member, root)
|
||
|
|
if index % 1000 == 0 or index == len(members):
|
||
|
|
print(f"extracted_airfrans_members={index}/{len(members)}", flush=True)
|
||
|
|
found = _find_of_dataset_root(root)
|
||
|
|
if found is None:
|
||
|
|
raise RuntimeError(f"OF_dataset directory not found after extracting {archive}")
|
||
|
|
case_count = _case_count(found)
|
||
|
|
if case_count < min_cases:
|
||
|
|
raise RuntimeError(f"Extracted AirfRANS OF_dataset has {case_count} cases; expected at least {min_cases}")
|
||
|
|
return found
|
||
|
|
|
||
|
|
|
||
|
|
def _hf_dataset_status(*, repo_id: str, path_in_repo: str) -> dict[str, Any]:
|
||
|
|
try:
|
||
|
|
from huggingface_hub import HfApi
|
||
|
|
except ModuleNotFoundError as exc:
|
||
|
|
raise RuntimeError("huggingface_hub is required for AirfRANS public data preparation") from exc
|
||
|
|
token = _optional_secret("HF_TOKEN")
|
||
|
|
api = HfApi(token=token)
|
||
|
|
try:
|
||
|
|
files = api.list_repo_files(repo_id=repo_id, repo_type="dataset")
|
||
|
|
except Exception:
|
||
|
|
files = []
|
||
|
|
prefix = path_in_repo.strip("/")
|
||
|
|
base = f"{prefix}/" if prefix else ""
|
||
|
|
npz_count = sum(1 for item in files if item.startswith(base) and item.endswith(".npz"))
|
||
|
|
has_manifest = any(item == f"{base}hf_dataset_manifest.json" for item in files)
|
||
|
|
return {
|
||
|
|
"file_count": len(files),
|
||
|
|
"npz_file_count": npz_count,
|
||
|
|
"has_manifest": has_manifest,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def _remote_content_length(url: str) -> int | None:
|
||
|
|
request = urllib.request.Request(url, method="HEAD")
|
||
|
|
try:
|
||
|
|
with urllib.request.urlopen(request, timeout=60) as response:
|
||
|
|
raw = response.headers.get("Content-Length")
|
||
|
|
except Exception:
|
||
|
|
return None
|
||
|
|
if raw is None:
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
return int(raw)
|
||
|
|
except ValueError:
|
||
|
|
return 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}")
|
||
|
|
if member.is_dir():
|
||
|
|
target.mkdir(parents=True, exist_ok=True)
|
||
|
|
return
|
||
|
|
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():
|
||
|
|
return direct
|
||
|
|
for candidate in root.glob("*/OF_dataset"):
|
||
|
|
if candidate.is_dir():
|
||
|
|
return candidate
|
||
|
|
if _case_count(root) > 0:
|
||
|
|
return root
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _case_count(root: Path) -> int:
|
||
|
|
return sum(1 for path in root.iterdir() if path.is_dir() and path.name.startswith("airFoil2D_")) if root.is_dir() else 0
|
||
|
|
|
||
|
|
|
||
|
|
def _optional_secret(name: str) -> str | None:
|
||
|
|
value = os.environ.get(name)
|
||
|
|
if value:
|
||
|
|
return value
|
||
|
|
for path in (Path(".env") / name, Path(".env") / f"{name}.txt"):
|
||
|
|
if path.is_file():
|
||
|
|
text = path.read_text().strip()
|
||
|
|
if text:
|
||
|
|
return text
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
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)
|
||
|
|
report_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|