From a8d17b7bf9c7b614bb786d623caf1cecf6d97827 Mon Sep 17 00:00:00 2001 From: linmuchuiyang <6169320+linmuchuiyang@users.noreply.github.com> Date: Fri, 18 Sep 2026 05:14:31 -0700 Subject: [PATCH] fix(models): default the HF config file lock to HF_HOME instead of ~/.cache/huggingface safe_load_config_with_retry() serialises AutoConfig.from_pretrained() behind a filelock. Without MEGATRON_CONFIG_LOCK_DIR the lock file was created under Path.home()/.cache/huggingface, which ignores HF_HOME and, in enroot/pyxis containers, lands on the ephemeral per-container layer. On an 8-node x 4-GPU HF->Megatron import of DeepSeek-V4-Pro every rank timed out 4 x 60 s on '/root/.cache/huggingface/.megatron_config_lock_.lock' with an error that did not mention the override. - Default the lock directory to huggingface_hub.constants.HF_HOME so the lock lives next to the cache it protects (unchanged behaviour when HF_HOME is unset; MEGATRON_CONFIG_LOCK_DIR still wins). - Mention MEGATRON_CONFIG_LOCK_DIR in the retries-exhausted error when the last failure was a filelock.Timeout. - Update the unit tests for the new default and add tests for the HF_HOME default and for the hint. Signed-off-by: linmuchuiyang <6169320+linmuchuiyang@users.noreply.github.com> --- .../hf_pretrained/safe_config_loader.py | 27 ++++++---- .../utils/test_safe_config_loader.py | 52 +++++++++++++++---- 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py b/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py index 5eb6186975..94403a4b6c 100644 --- a/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py +++ b/src/megatron/bridge/models/hf_pretrained/safe_config_loader.py @@ -27,6 +27,7 @@ from typing import Union import filelock +from huggingface_hub import constants as hf_constants from transformers import AutoConfig from transformers.configuration_utils import PretrainedConfig @@ -59,8 +60,10 @@ def safe_load_config_with_retry( Environment Variables: MEGATRON_CONFIG_LOCK_DIR: Override the directory where lock files are created. - Default: ~/.cache/huggingface/ - Useful for multi-node setups where a shared lock directory is needed. + Default: the Hugging Face home directory (``HF_HOME``, normally + ``~/.cache/huggingface``), so the lock lives next to the cache it protects. + Set it to a shared directory when ``HF_HOME`` is not shared across the ranks + of a multi-node job. Example: >>> config = safe_load_config_with_retry("meta-llama/Meta-Llama-3-8B") @@ -88,12 +91,12 @@ def safe_load_config_with_retry( path_hash = hashlib.md5(str(path).encode()).hexdigest() # Allow override of lock directory via environment variable - # This is useful for multi-node setups where a shared lock directory is needed - lock_dir = os.getenv("MEGATRON_CONFIG_LOCK_DIR") - if lock_dir: - lock_file = Path(lock_dir) / f".megatron_config_lock_{path_hash}" - else: - lock_file = Path.home() / ".cache" / "huggingface" / f".megatron_config_lock_{path_hash}" + # (useful for multi-node setups where a shared lock directory is needed). + # Otherwise follow HF_HOME: that is where the config cache being protected + # lives, and on containerised clusters $HOME is often an ephemeral, + # per-container layer that is neither shared nor persistent. + lock_dir = os.getenv("MEGATRON_CONFIG_LOCK_DIR") or hf_constants.HF_HOME + lock_file = Path(lock_dir) / f".megatron_config_lock_{path_hash}" lock_file.parent.mkdir(parents=True, exist_ok=True) @@ -131,8 +134,14 @@ def safe_load_config_with_retry( break # All retries exhausted + hint = "" + if isinstance(last_exception, filelock.Timeout): + hint = ( + " The config file lock could not be acquired; on multi-node jobs set " + "MEGATRON_CONFIG_LOCK_DIR to a directory on a shared filesystem." + ) raise ValueError( f"Failed to load configuration from {path} after {max_retries + 1} attempts. " - f"This might be due to network issues or concurrent access conflicts. " + f"This might be due to network issues or concurrent access conflicts.{hint} " f"Last error: {last_exception}" ) from last_exception diff --git a/tests/unit_tests/utils/test_safe_config_loader.py b/tests/unit_tests/utils/test_safe_config_loader.py index d843a7dbbd..0961346d68 100644 --- a/tests/unit_tests/utils/test_safe_config_loader.py +++ b/tests/unit_tests/utils/test_safe_config_loader.py @@ -23,6 +23,7 @@ from unittest.mock import MagicMock, Mock, patch import pytest +from huggingface_hub import constants as hf_constants from transformers.configuration_utils import PretrainedConfig from megatron.bridge.models.hf_pretrained.safe_config_loader import safe_load_config_with_retry @@ -88,9 +89,7 @@ def test_with_file_locking(self): # Verify file lock was used expected_hash = hashlib.md5(str(self.test_path).encode()).hexdigest() - expected_lock_file = ( - Path.home() / ".cache" / "huggingface" / f".megatron_config_lock_{expected_hash}.lock" - ) + expected_lock_file = Path(hf_constants.HF_HOME) / f".megatron_config_lock_{expected_hash}.lock" mock_filelock.assert_called_once_with(str(expected_lock_file), timeout=60) # Verify context manager was used @@ -280,9 +279,7 @@ def test_lock_file_creation(self): # Verify unique lock file for each path expected_hash = hashlib.md5(str(path).encode()).hexdigest() - expected_lock_file = ( - Path.home() / ".cache" / "huggingface" / f".megatron_config_lock_{expected_hash}.lock" - ) + expected_lock_file = Path(hf_constants.HF_HOME) / f".megatron_config_lock_{expected_hash}.lock" # Check that FileLock was called with the correct path found_call = False @@ -380,13 +377,50 @@ def test_default_lock_directory_when_env_var_not_set(self): # Verify default lock directory was used expected_hash = hashlib.md5(str(self.test_path).encode()).hexdigest() - expected_lock_file = ( - Path.home() / ".cache" / "huggingface" / f".megatron_config_lock_{expected_hash}.lock" - ) + expected_lock_file = Path(hf_constants.HF_HOME) / f".megatron_config_lock_{expected_hash}.lock" mock_filelock.assert_called_once_with(str(expected_lock_file), timeout=60) assert result == self.mock_config + def test_default_lock_directory_follows_hf_home(self, tmp_path): + """Without MEGATRON_CONFIG_LOCK_DIR the lock lives under HF_HOME, not under $HOME.""" + mock_lock = MagicMock() + hf_home = tmp_path / "hf_home" + + with patch.dict(os.environ, {}, clear=True): + with patch("megatron.bridge.models.hf_pretrained.safe_config_loader.hf_constants.HF_HOME", str(hf_home)): + with patch( + "megatron.bridge.models.hf_pretrained.safe_config_loader.filelock.FileLock" + ) as mock_filelock: + mock_filelock.return_value = mock_lock + mock_lock.__enter__ = Mock(return_value=mock_lock) + mock_lock.__exit__ = Mock(return_value=None) + + with patch( + "megatron.bridge.models.hf_pretrained.safe_config_loader.AutoConfig" + ) as mock_auto_config: + mock_auto_config.from_pretrained.return_value = self.mock_config + + result = safe_load_config_with_retry(self.test_path) + + expected_hash = hashlib.md5(str(self.test_path).encode()).hexdigest() + expected_lock_file = hf_home / f".megatron_config_lock_{expected_hash}.lock" + mock_filelock.assert_called_once_with(str(expected_lock_file), timeout=60) + assert hf_home.is_dir() + assert result == self.mock_config + + def test_lock_timeout_error_mentions_override(self): + """A lock timeout on every attempt surfaces the MEGATRON_CONFIG_LOCK_DIR hint.""" + import filelock + + with patch("megatron.bridge.models.hf_pretrained.safe_config_loader.filelock.FileLock") as mock_filelock: + mock_filelock.return_value.__enter__ = Mock(side_effect=filelock.Timeout("/tmp/some.lock")) + mock_filelock.return_value.__exit__ = Mock(return_value=None) + + with patch("megatron.bridge.models.hf_pretrained.safe_config_loader.time.sleep"): + with pytest.raises(ValueError, match="MEGATRON_CONFIG_LOCK_DIR"): + safe_load_config_with_retry(self.test_path, max_retries=1) + def test_custom_lock_directory_with_pathlib_path(self, tmp_path): """Test that custom lock directory works with pathlib.Path inputs.""" mock_lock = MagicMock()