[data] fix: tolerate shared filesystem mkdir races (#3888)

Signed-off-by: jbernloehr <jbernloehr@nvidia.com>
Co-authored-by: Kamran Jafari <kjafarisadeg@nvidia.com>

Author: janbernloehr

src/megatron/bridge/data/builders/finetuning_dataset.py

@@ -339,7 +339,11 @@ def default_pack_path(self) -> Path:
             self.dataset_root / "packed" / f"{tokenizer_model_name}_pad_seq_to_mult{self._pad_seq_to_mult}"
         )
         if not default_pack_path.exists():
-            default_pack_path.mkdir(parents=True, exist_ok=True)
+            try:
+                # Shared filesystems can expose stale parent-dir state despite exist_ok=True.
+                default_pack_path.mkdir(parents=True, exist_ok=True)
+            except (FileExistsError, FileNotFoundError):
+                pass
             logger.info(f"Using default path for packing files: {str(default_pack_path)}")
 
         return default_pack_path

src/megatron/bridge/data/datasets/sft.py

@@ -91,7 +91,11 @@ def get_dataset_root(name: str) -> Path:
         Path: The path to the dataset's root directory.
     """
     output = Path(NEMO_DATASETS_CACHE) / name
-    output.mkdir(parents=True, exist_ok=True)
+    try:
+        # Shared filesystems can expose stale parent-dir state despite exist_ok=True.
+        output.mkdir(parents=True, exist_ok=True)
+    except (FileExistsError, FileNotFoundError):
+        pass
 
     return output
 

tests/unit_tests/data/builders/test_finetuning_dataset.py

@@ -0,0 +1,48 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from pathlib import Path
+from unittest.mock import MagicMock
+
+import pytest
+
+from megatron.bridge.data.builders.finetuning_dataset import FinetuningDatasetBuilder
+from megatron.bridge.data.datasets.packed_sequence import PackedSequenceSpecs
+
+
+@pytest.mark.parametrize("mkdir_error", [FileExistsError, FileNotFoundError])
+def test_default_pack_path_ignores_shared_fs_mkdir_race(tmp_path, monkeypatch, mkdir_error):
+    """Network filesystems can leak mkdir races even with exist_ok=True."""
+    builder = FinetuningDatasetBuilder(
+        dataset_root=tmp_path,
+        tokenizer=MagicMock(),
+        packed_sequence_specs=PackedSequenceSpecs(
+            packed_sequence_size=128,
+            tokenizer_model_name="mock-tokenizer",
+            pad_seq_to_mult=8,
+        ),
+    )
+    expected_path = tmp_path / "packed" / "mock-tokenizer_pad_seq_to_mult8"
+
+    monkeypatch.setattr(Path, "exists", lambda _: False)
+
+    def raise_mkdir(self, parents=False, exist_ok=False):
+        assert self == expected_path
+        assert parents is True
+        assert exist_ok is True
+        raise mkdir_error("stale shared filesystem state")
+
+    monkeypatch.setattr(Path, "mkdir", raise_mkdir)
+
+    assert builder.default_pack_path == expected_path

tests/unit_tests/data/datasets/test_sft.py

@@ -13,11 +13,14 @@
 # limitations under the License.
 
 import json
+from pathlib import Path
 from unittest.mock import MagicMock
 
 import numpy as np
+import pytest
 
-from megatron.bridge.data.datasets.sft import GPTSFTChatDataset, GPTSFTDataset, GPTSFTPackedDataset
+from megatron.bridge.data.datasets import sft as sft_module
+from megatron.bridge.data.datasets.sft import GPTSFTChatDataset, GPTSFTDataset, GPTSFTPackedDataset, get_dataset_root
 
 
 def create_mock_tokenizer():
@@ -98,6 +101,22 @@ def get_gpt_sft(tmp_path, dataset_type="sft"):
     return dataset, num_samples
 
 
+@pytest.mark.parametrize("mkdir_error", [FileExistsError, FileNotFoundError])
+def test_get_dataset_root_ignores_shared_fs_mkdir_race(tmp_path, monkeypatch, mkdir_error):
+    """Network filesystems can leak mkdir races even with exist_ok=True."""
+    monkeypatch.setattr(sft_module, "NEMO_DATASETS_CACHE", tmp_path)
+
+    def raise_mkdir(self, parents=False, exist_ok=False):
+        assert self == tmp_path / "sft"
+        assert parents is True
+        assert exist_ok is True
+        raise mkdir_error("stale shared filesystem state")
+
+    monkeypatch.setattr(Path, "mkdir", raise_mkdir)
+
+    assert get_dataset_root("sft") == tmp_path / "sft"
+
+
 class TestDataGPTSFTDataset:
     def test_build_samples_mapping(self, tmp_path):
         dataset, _ = get_gpt_sft(tmp_path)