[codex] Trigger Hugging Face access during inference - #13
Conversation
Greptile SummaryThis PR moves Hugging Face model-weight preparation into the inference path for both
Confidence Score: 4/5The single-GPU and multi-GPU initialization flows are structurally correct; the one path worth hardening before wide production use is the broken-symlink recovery in The core change — calling Both Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant MONAI as MONAI Bundle
participant prepHF as prepare_huggingface_checkpoint
participant HF as Hugging Face Hub
participant Disk as Local Filesystem
User->>MONAI: monai.bundle run (inference.json)
MONAI->>prepHF: initialize[0]: prepare_huggingface_checkpoint(repo_id, ckpt_file, checkpoint_path, counter_file)
prepHF->>HF: "hf_hub_download(config.json, force_download=True)"
HF-->>prepHF: config.json cached path (HF counter touched)
alt models/model.pt missing
prepHF->>HF: hf_hub_download(vista3d_pretrained_model/model.pt)
HF-->>prepHF: cached .pt path
prepHF->>Disk: symlink_to(cached .pt) OR shutil.copy2
Disk-->>prepHF: models/model.pt ready
end
prepHF-->>MONAI: str(local_path)
MONAI->>MONAI: "set_determinism(seed=123)"
MONAI->>Disk: CheckpointLoader(models/model.pt)
MONAI->>MONAI: evaluator.run()
Note over MONAI,Disk: Multi-GPU: non-rank-0 workers skip prepHF early,<br/>dist.barrier() ensures rank 0 finishes before CheckpointLoader
Reviews (1): Last reviewed commit: "Fix pre-commit lint for inference PR" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| try: | ||
| local_path.symlink_to(checkpoint_path) | ||
| except OSError: | ||
| shutil.copy2(checkpoint_path, local_path) |
There was a problem hiding this comment.
shutil.copy2 fallback
If the HF cache has been cleaned up (e.g. huggingface-cli delete-cache), local_path can be a broken symlink: local_path.exists() returns False (it follows the symlink), so the code proceeds to download and then calls symlink_to, which raises FileExistsError because the symlink inode still exists. The OSError catch then attempts shutil.copy2(checkpoint_path, local_path), but on Linux this writes through the broken symlink to its (now-missing) target — if the target's parent directory was also removed, this raises FileNotFoundError, leaving the broken symlink in place and inference stuck.
Add if local_path.is_symlink() and not local_path.exists(): local_path.unlink() before the symlink_to call to clear the stale symlink first. The same issue exists in NV-Segment-CTMR/scripts/huggingface_download.py.
| def _is_rank_zero() -> bool: | ||
| for name in ("RANK", "LOCAL_RANK", "SLURM_PROCID"): | ||
| value = os.environ.get(name) | ||
| if value not in (None, "", "0"): | ||
| return False | ||
| return True |
There was a problem hiding this comment.
_is_rank_zero returns False if any of RANK, LOCAL_RANK, or SLURM_PROCID is not in (None, "", "0"). In a multi-node SLURM launch that wraps torchrun, it is possible for SLURM_PROCID to be set to the SLURM task ID while RANK/LOCAL_RANK are also set. If those values mismatch (e.g., SLURM assigns non-zero SLURM_PROCID to the PyTorch rank-0 process), the function incorrectly returns False for rank 0, meaning no process performs the HF download, and the barrier in mgpu_inference.json proceeds with models/model.pt still missing, causing all ranks to crash at CheckpointLoader. The same logic is duplicated in the CTMR variant.
1 participant
Summary
NV-Segment-CTandNV-Segment-CTMR.config.jsonon every inference run so cached local weights still register a Hub access.models/model.ptfrom the Hugging Face cache if the local checkpoint path is missing.hf downloadstep.Why
Previously setup downloaded model weights once, so later inference runs used only
models/model.ptand did not touch Hugging Face. This follows the Cosmos-style runtime checkpoint access pattern while preserving the local MONAI checkpoint path expected by existing configs.Validation
Local:
python3 -m py_compile NV-Segment-CT/scripts/huggingface_download.py NV-Segment-CTMR/scripts/huggingface_download.pypython3 -m json.toolfor changed inference, mgpu, and batch configsbash -nfor bothscripts/test_batch_inference.shscriptsDFW:
HEAD/GETfornvidia/NV-Segment-CTMR/config.json.NV-Segment-CT/models/model.ptfrom Hugging Face.DFW output paths:
/lustre/fsw/portfolios/healthcareeng/users/yufanh/experiments/nvseg_ctmr_batch_test/output/s0289/s0289_trans.nii.gz/lustre/fsw/portfolios/healthcareeng/users/yufanh/experiments/nvseg_ct_batch_test/output/spleen_03/spleen_03_trans.nii.gz