fix(models): default the HF config file lock to HF_HOME instead of ~/.cache/huggingface

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

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()