fix(dataset): retry HF datasets-server /splits; xfail flaky preprocess tests

The HF datasets-server /splits endpoint intermittently returns transient
5xx (e.g. 503), failing download_hf_dataset_as_jsonl and the megatron
preprocess-data tests that pull nanotron/minipile_100_samples.

- Retry the /splits GET with backoff (1s/2s/4s) on 5xx / connection /
  timeout; fail fast on client errors (e.g. 404 unknown dataset).
- xfail(raises=RuntimeError, strict=False) the 3 dataset-dependent tests
  so a transient outage is allowed to fail without blocking the suite,
  while real assertion failures stay visible.

Signed-off-by: Chenhan Yu <chenhany@nvidia.com>

Author: ChenhanYu

modelopt/torch/utils/dataset_utils.py

@@ -19,6 +19,7 @@
 import json
 import os
 import random
+import time
 from collections.abc import Callable, Iterator
 from contextlib import contextmanager, suppress
 from pathlib import Path
@@ -1291,15 +1292,28 @@ def download_hf_dataset_as_jsonl(
         json_keys = [json_keys]
     jsonl_paths: list[str] = []
 
-    try:
-        response = requests.get(
-            f"https://datasets-server.huggingface.co/splits?dataset={dataset_name}",
-            headers=build_hf_headers(),
-            timeout=10,
-        )
-        response.raise_for_status()
-    except requests.RequestException as e:
-        raise RuntimeError(f"Failed to fetch dataset splits for {dataset_name}: {e}") from e
+    # The HF datasets-server /splits endpoint is intermittently unavailable
+    # (transient 5xx). Retry with backoff so a momentary outage doesn't fail the
+    # whole preprocess run; fail fast on client errors (e.g. 404 unknown dataset).
+    splits_url = f"https://datasets-server.huggingface.co/splits?dataset={dataset_name}"
+    response = None
+    last_exc: Exception | None = None
+    for attempt in range(4):
+        try:
+            response = requests.get(splits_url, headers=build_hf_headers(), timeout=10)
+            response.raise_for_status()
+            break
+        except requests.RequestException as e:
+            last_exc = e
+            status = getattr(getattr(e, "response", None), "status_code", None)
+            if status is not None and status < 500:
+                break  # client error — retrying won't help
+            if attempt < 3:
+                time.sleep(2**attempt)  # 1s, 2s, 4s
+    if response is None or not response.ok:
+        raise RuntimeError(
+            f"Failed to fetch dataset splits for {dataset_name}: {last_exc}"
+        ) from last_exc
 
     response_json = response.json()
     print_rank_0(f"\nFound {len(response_json['splits'])} total splits for {dataset_name}:")

tests/gpu_megatron/torch/utils/plugins/test_megatron_preprocess_data.py

@@ -23,6 +23,15 @@
 from modelopt.torch.utils.plugins.megatron_preprocess_data import megatron_preprocess_data
 
 
+# Depends on the external HF datasets-server, which intermittently returns 5xx
+# (e.g. 503) on the /splits lookup. Allow-fail the transient outage while keeping
+# real assertion failures visible (raises=RuntimeError only, strict=False so a
+# normal pass doesn't xpass-fail). download_hf_dataset_as_jsonl already retries.
+@pytest.mark.xfail(
+    raises=RuntimeError,
+    strict=False,
+    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
+)
 def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     input_jsonl = download_hf_dataset_as_jsonl("nanotron/minipile_100_samples", tmp_path / "raw")
     assert len(input_jsonl) == 1, "Expected 1 JSONL file"
@@ -53,6 +62,12 @@ def test_megatron_preprocess_data_with_jsonl_path(tmp_path):
     assert Path(prefixes[0] + ".idx").stat().st_size > 0, "Index file should not be empty"
 
 
+# Allow-fail transient HF datasets-server 5xx (see the /splits note above).
+@pytest.mark.xfail(
+    raises=RuntimeError,
+    strict=False,
+    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
+)
 @pytest.mark.parametrize(
     ("hf_dataset", "hf_split", "json_keys"),
     [
@@ -197,6 +212,12 @@ def test_megatron_preprocess_data_tool_calls_arguments_normalized(tmp_path):
     )
 
 
+# Allow-fail transient HF datasets-server 5xx (see the /splits note above).
+@pytest.mark.xfail(
+    raises=RuntimeError,
+    strict=False,
+    reason="Flaky: transient HF datasets-server 5xx on the /splits lookup",
+)
 def test_megatron_preprocess_data_hf_streaming_warning(tmp_path):
     # hf_streaming without hf_max_samples_per_split should warn and fall back to non-streaming
     with pytest.warns(UserWarning, match="hf_streaming"):