[OMNIML-6228] Fix two flaky tests: MTP causality tolerance and example-runner port collisions

tests/_test_utils/examples/run_command.py

@@ -28,8 +28,9 @@
 
 MODELOPT_ROOT = Path(__file__).parents[3]
 
-# Substrings that identify a *transient* HuggingFace Hub / dataset access failure
-_HF_TRANSIENT_MARKERS = (
+# Substrings that identify a *transient* failure a rerun clears: a HuggingFace Hub / dataset access
+# error, or a rendezvous port that another process bound between its selection and its use
+_TRANSIENT_MARKERS = (
     "HfHubHTTPError",
     "Too Many Requests",
     "500 Server Error",
@@ -51,6 +52,8 @@
     "Couldn't reach",  # datasets: "Couldn't reach <repo> on the Hub"
     "Temporary failure in name resolution",  # transient DNS
     "Name or service not known",  # transient DNS
+    "EADDRINUSE",  # torch.distributed's TCPStore: "name: EADDRINUSE, message: address already in use"
+    "Address already in use",  # the same from Python's OSError, gloo and NCCL
 )
 _HF_MAX_RETRIES = 1
 _HF_RETRY_DELAY_S = 10
@@ -156,7 +159,7 @@ def run_example_command(
     hf_max_retries: int = _HF_MAX_RETRIES,
     hf_retry_delay_s: int = _HF_RETRY_DELAY_S,
 ) -> str | None:
-    """Run an example command, retrying transient HuggingFace access errors."""
+    """Run an example command, retrying transient errors (see ``_TRANSIENT_MARKERS``) once."""
     print(f"[{example_path}] Running command: {cmd_parts}")
     in_process = _in_process_runner if env is None else None
     if _in_process_runner is not None and env is not None:
@@ -172,8 +175,8 @@ def run_example_command(
             # copy it never sees.
             env["MASTER_PORT"] = str(get_free_port())  # fresh port per attempt
         if in_process is not None:
-            # Inside the loop so in-process steps get the same transient-HuggingFace retries;
-            # these tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
+            # Inside the loop so in-process steps get the same transient-error retries; these
+            # tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
             try:
                 result = in_process(cmd_parts, example_path)
             except Exception as e:
@@ -186,7 +189,7 @@ def run_example_command(
                 captured = getattr(e, "captured_output", "")
                 text = f"{traceback.format_exc()}\n{captured}"
                 if attempt == hf_max_retries or not any(
-                    marker in text for marker in _HF_TRANSIENT_MARKERS
+                    marker in text for marker in _TRANSIENT_MARKERS
                 ):
                     raise
                 returncode, output = 1, text
@@ -197,11 +200,11 @@ def run_example_command(
             returncode, output = _run_capturing(cmd_parts, cwd, env)
         if returncode == 0:
             return output
-        transient = any(marker in output for marker in _HF_TRANSIENT_MARKERS)
+        transient = any(marker in output for marker in _TRANSIENT_MARKERS)
         if not transient or attempt == hf_max_retries:
             raise subprocess.CalledProcessError(returncode, cmd_parts, output=output)
         warnings.warn(
-            f"[{example_path}] transient HuggingFace access error; retrying in "
+            f"[{example_path}] transient error; retrying in "
             f"{hf_retry_delay_s}s (attempt {attempt + 1}/{hf_max_retries})"
         )
         time.sleep(hf_retry_delay_s)

tests/unit/test_example_run_command.py

@@ -23,6 +23,37 @@
 from _test_utils.examples import run_command
 
 
+@pytest.mark.parametrize(
+    ("error", "attempts"),
+    [
+        ("HfHubHTTPError: 503 Server Error: Service Unavailable", 2),
+        ("DistNetworkError: port: 38789, name: EADDRINUSE, message: address already in use", 2),
+        ("AssertionError: the quantized model diverged", 1),
+    ],
+)
+def test_run_example_command_retries_only_transient_errors(monkeypatch, error, attempts):
+    calls = []
+
+    def step_failing_once(cmd_parts, example_path):
+        calls.append(cmd_parts)
+        if len(calls) == 1:
+            raise RuntimeError(error)
+        return "output"
+
+    def run():
+        return run_command.run_example_command(
+            ["torchrun", "step.py"], "hf_ptq", hf_retry_delay_s=0
+        )
+
+    monkeypatch.setattr(run_command, "_in_process_runner", step_failing_once)
+    if attempts == 1:
+        with pytest.raises(RuntimeError, match="diverged"):
+            run()
+    else:
+        assert run() == "output"
+    assert len(calls) == attempts
+
+
 def test_run_capturing_does_not_block_on_a_survivor_holding_the_pipe(
     skip_on_windows, monkeypatch, tmp_path
 ):

tests/unit/torch/models/test_nemotron_h_mtp.py

@@ -246,8 +246,9 @@ def test_mtp_later_attention_blocks_remain_causal(tiny_checkpoint):
         hidden[:, -1] += 10
         embeddings[:, -1] += 10
         actual = mtp(hidden, embeddings)
-    # MTP uses next-token embeddings, so only the last two positions may change.
-    torch.testing.assert_close(actual[:, :-2], expected[:, :-2], rtol=1e-6, atol=1e-9)
+    # MTP uses next-token embeddings, so only the last two positions may change. The tolerance
+    # covers float noise (below 1e-7 over 400 RNG states); a real leak changes them by ~0.1.
+    torch.testing.assert_close(actual[:, :-2], expected[:, :-2], rtol=1e-5, atol=1e-6)
 
 
 def test_fused_weight_quantizers_need_no_activation_forward():