#!/usr/bin/env python3 """Static and file-layout preflight for the standalone Protenix package.""" from __future__ import annotations import argparse import ast import hashlib import os from pathlib import Path MODEL_SHA256 = "9ea20b0aba42f2256711da1d0cd081510a4b291e64375bff6b70ced70b87a5f1" MODEL_SIZE = 1474265486 REQUIRED_PACKAGE_FILES = [ "README.md", "configuration.json", "configs/inference_config.yaml", "weight/model_v0.5.0.pt", "examples/7r6r.json", "examples/7r6r/msa/1/pairing.a3m", "examples/7r6r/msa/1/non_pairing.a3m", "models/protenix/protenix.py", "models/openfold/primitives.py", "scripts/runner/inference_unified.py", "scripts/run_inference.py", "scripts/train.py", "scripts/finetune.py", ] REQUIRED_DATA_FILES = [ "components.v20240608.cif", "components.v20240608.cif.rdkit_mol.pkl", "seq_to_pdb_index.json", "indices/weightedPDB_indices_before_2021-09-30_wo_posebusters_resolution_below_9.csv.gz", "indices/recentPDB_low_homology_maxtoken1536.csv", "indices/recentPDB_low_homology_maxtoken1024_sample384_pdb_id.txt", "indices/posebusters_indices_mainchain_interface.csv", "mmcif_msa", ] def sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as f: for chunk in iter(lambda: f.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def module_to_path(model_root: Path, module: str) -> Path | None: if not module.startswith(("configs.", "models.", "scripts.")): return None rel = Path(*module.split(".")) package_init = model_root / rel / "__init__.py" module_file = model_root / rel.with_suffix(".py") if package_init.exists(): return package_init if module_file.exists(): return module_file return module_file def check_local_imports(model_root: Path) -> list[str]: missing: list[str] = [] scan_roots = ["configs", "models", "scripts"] py_files = [] for scan_root in scan_roots: py_files.extend((model_root / scan_root).rglob("*.py")) for py_file in sorted(py_files): tree = ast.parse(py_file.read_text(encoding="utf-8"), filename=str(py_file)) for node in ast.walk(tree): module = None if isinstance(node, ast.Import): for alias in node.names: module = alias.name if module == "modules" or module.startswith("modules."): missing.append(f"{py_file.relative_to(model_root)} imports removed package {module}") continue if module == "config" or module.startswith("config."): missing.append(f"{py_file.relative_to(model_root)} imports renamed package {module}") continue target = module_to_path(model_root, module) if target is not None and not target.exists(): missing.append(f"{py_file.relative_to(model_root)} imports missing {module}") elif isinstance(node, ast.ImportFrom) and node.module: module = node.module if node.level == 0: if module == "modules" or module.startswith("modules."): missing.append(f"{py_file.relative_to(model_root)} imports removed package {module}") continue if module == "config" or module.startswith("config."): missing.append(f"{py_file.relative_to(model_root)} imports renamed package {module}") continue target = module_to_path(model_root, module) if target is not None and not target.exists(): missing.append(f"{py_file.relative_to(model_root)} imports missing {module}") return missing def main() -> int: parser = argparse.ArgumentParser(description="Protenix standalone package preflight") parser.add_argument("--model-root", default=".", help="Package root") parser.add_argument( "--data-root", default=os.environ.get("DATA_ROOT_DIR", "../bio_protenix_dataset"), help="Prepared Protenix dataset root", ) parser.add_argument("--strict-weights", action="store_true", help="Validate weight size and LFS pointer state") parser.add_argument("--full-checksum", action="store_true", help="Compute SHA256 for the 1.4GB model file") parser.add_argument("--strict-imports", action="store_true", help="Statically verify local imports and removed modules.* imports") parser.add_argument("--strict-data", action="store_true", help="Require the external dataset files") args = parser.parse_args() model_root = Path(args.model_root).resolve() data_root = Path(args.data_root).resolve() errors: list[str] = [] for rel in REQUIRED_PACKAGE_FILES: path = model_root / rel if not path.exists(): errors.append(f"Missing package file: {path}") config = model_root / "configs/inference_config.yaml" if config.exists(): text = config.read_text(encoding="utf-8") required_fragments = [ 'input_json_path: "./examples/7r6r.json"', 'load_checkpoint_path: "./weight/model_v0.5.0.pt"', 'ccd_components_file: "${DATA_ROOT_DIR}/components.v20240608.cif"', 'ccd_components_rdkit_mol_file: "${DATA_ROOT_DIR}/components.v20240608.cif.rdkit_mol.pkl"', 'pdb_mmseqs_dir: "${DATA_ROOT_DIR}/mmcif_msa"', "${oc.env:PWD}/examples/7r6r/msa/1", ] for fragment in required_fragments: if fragment not in text: errors.append(f"Config missing fragment: {fragment}") checkpoint = model_root / "weight/model_v0.5.0.pt" if args.strict_weights and checkpoint.exists(): size = checkpoint.stat().st_size if size != MODEL_SIZE: errors.append(f"Weight size mismatch: {checkpoint} size={size} expected={MODEL_SIZE}") with checkpoint.open("rb") as f: prefix = f.read(64) if prefix.startswith(b"version https://git-lfs.github.com/spec"): errors.append(f"Weight is a Git LFS pointer, not the real checkpoint: {checkpoint}") if args.full_checksum: digest = sha256(checkpoint) if digest != MODEL_SHA256: errors.append(f"Weight SHA256 mismatch: {checkpoint} sha256={digest}") if args.strict_data: for rel in REQUIRED_DATA_FILES: path = data_root / rel if not path.exists(): errors.append(f"Missing dataset file or directory: {path}") if args.strict_imports: errors.extend(check_local_imports(model_root)) residue_patterns = [ "from modules.", "import modules.", "modules.models.", "modules.runner.", "/public/share/sugonhpcapp01/onestore/" + "onemodels", "/public/home/liuyx19/" + "one" + "science", "checkpoints/" + "model_v0.5.0.pt", "infer_" + "datasets/" + "7r6r", "url" + "retrieve", ] for path in sorted(model_root.rglob("*")): if not path.is_file() or "weight" in path.parts: continue if "modules" in path.relative_to(model_root).parts: continue if path.relative_to(model_root).as_posix() == "scripts/preflight.py": continue try: text = path.read_text(encoding="utf-8") except UnicodeDecodeError: continue for pattern in residue_patterns: if pattern in text: errors.append(f"Residual pattern {pattern!r} in {path.relative_to(model_root)}") print(f"model_root={model_root}") print(f"data_root={data_root}") if errors: print("Preflight failed:") for error in errors: print(f"- {error}") return 1 print("Preflight passed.") return 0 if __name__ == "__main__": raise SystemExit(main())