| """Environment preflight checks for conversion and golden export.""" |
|
|
| from __future__ import annotations |
|
|
| import importlib.util |
| import json |
| import subprocess |
| import sys |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from typing import Iterable |
|
|
|
|
| @dataclass(frozen=True) |
| class PreflightResult: |
| name: str |
| status: str |
| detail: str |
|
|
|
|
| def _module_available(module_name: str) -> bool: |
| return importlib.util.find_spec(module_name) is not None |
|
|
|
|
| def check_python_modules(modules: Iterable[str]) -> list[PreflightResult]: |
| results: list[PreflightResult] = [] |
| for module_name in modules: |
| if _module_available(module_name): |
| results.append(PreflightResult(module_name, "ok", "module import spec found")) |
| else: |
| results.append(PreflightResult(module_name, "missing", "module import spec not found")) |
| return results |
|
|
|
|
| def check_mlx_import(timeout_seconds: int = 20) -> PreflightResult: |
| code = ( |
| "import mlx.core as mx; " |
| "print({'default_device': str(mx.default_device()), " |
| "'has_metal': hasattr(mx, 'metal')})" |
| ) |
| try: |
| completed = subprocess.run( |
| [sys.executable, "-c", code], |
| check=False, |
| capture_output=True, |
| text=True, |
| timeout=timeout_seconds, |
| ) |
| except subprocess.TimeoutExpired: |
| return PreflightResult("mlx_import", "failed", f"timed out after {timeout_seconds}s") |
| detail = (completed.stdout + completed.stderr).strip() |
| if completed.returncode == 0: |
| return PreflightResult("mlx_import", "ok", detail) |
| return PreflightResult("mlx_import", "failed", detail) |
|
|
|
|
| def check_paths(paths: Iterable[Path]) -> list[PreflightResult]: |
| results: list[PreflightResult] = [] |
| for path in paths: |
| if path.exists(): |
| results.append(PreflightResult(str(path), "ok", "path exists")) |
| else: |
| results.append(PreflightResult(str(path), "missing", "path does not exist")) |
| return results |
|
|
|
|
| def run_preflight(include_mlx_import: bool = True) -> list[PreflightResult]: |
| results = check_python_modules(["mlx", "transformers", "PIL"]) |
| if include_mlx_import: |
| results.append(check_mlx_import()) |
| results.extend( |
| check_paths( |
| [ |
| Path("split_files/text_encoders/qwen_3_06b_base-mlx.safetensors"), |
| Path("split_files/diffusion_models/anima-base-v1.0-mlx.safetensors"), |
| Path("split_files/vae/qwen_image_vae-mlx.safetensors"), |
| Path("tokenizers/qwen25_tokenizer"), |
| Path("tokenizers/t5_tokenizer"), |
| ] |
| ) |
| ) |
| return results |
|
|
|
|
| def write_preflight_report(path: Path, results: list[PreflightResult]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| path.write_text(json.dumps([asdict(result) for result in results], indent=2) + "\n", encoding="utf-8") |
|
|
