[data] fix: Propagate MIMO dataset loading errors (#5873)

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/data/megatron_mimo/hf_provider.py

@@ -146,9 +146,11 @@ def _load_hf_dataset(self, split: str) -> Any:
                 ),
             )
             return dataset
-        except ValueError:
-            # Split doesn't exist
-            return None
+        except ValueError as error:
+            message = str(error)
+            if message.startswith('Unknown split "') and ". Should be one of " in message:
+                return None
+            raise
 
     def _build_split_dataset(
         self,

tests/unit_tests/data/megatron_mimo/test_hf_provider.py

@@ -3,6 +3,7 @@
 
 from dataclasses import dataclass
 
+import pytest
 import torch
 
 from megatron.bridge.data.base import DatasetBuildContext
@@ -58,7 +59,7 @@ def fake_load_dataset(path, name=None, split=None, trust_remote_code=None, data_
         del path, name, trust_remote_code, data_files
         calls.load_dataset += 1
         if split == "validation":
-            raise ValueError("missing split")
+            raise ValueError("Unknown split \"validation\". Should be one of ['train', 'test'].")
         return [{"text": "hello", "image": "image_0.jpg"}]
 
     class _AutoImageProcessor:
@@ -95,6 +96,21 @@ def from_pretrained(path, trust_remote_code=None):
     assert calls.is_safe_repo >= 3
 
 
+def test_build_datasets_propagates_non_split_value_error(monkeypatch):
+    def fake_load_dataset(path, name=None, split=None, trust_remote_code=None, data_files=None):
+        del path, name, split, trust_remote_code, data_files
+        raise ValueError("Empty 'data_files': '[]'. It should be either non-empty or None (default).")
+
+    provider = _make_provider()
+    monkeypatch.setattr(provider, "_load_processors", lambda: {})
+    monkeypatch.setattr(provider, "_load_tokenizer", lambda: DummyTokenizer())
+    monkeypatch.setattr("megatron.bridge.data.megatron_mimo.hf_provider.load_dataset", fake_load_dataset)
+
+    context = DatasetBuildContext(train_samples=1, valid_samples=0, test_samples=0)
+    with pytest.raises(ValueError, match="Empty 'data_files'"):
+        provider.build_datasets(context)
+
+
 def test_get_collate_fn_returns_partial():
     provider = _make_provider()
     collate_fn = provider.get_collate_fn()