Handle token export mkdir races

Author: samsja

src/prime_rl/trainer/rl/token_export.py

@@ -1,6 +1,7 @@
 import atexit
 import json
 import math
+import time
 from collections.abc import Mapping, Sequence
 from pathlib import Path
 from typing import Any
@@ -12,6 +13,8 @@
 from prime_rl.trainer.rl.loss import compute_importance_ratio_and_mismatch_kl
 
 SCHEMA_VERSION = 1
+MKDIR_RETRY_DELAY_SECONDS = 0.1
+MKDIR_MAX_ATTEMPTS = 5
 
 
 class DisabledTokenExporter:
@@ -88,7 +91,7 @@ def _start_step(self, step: int) -> None:
         self._current_step = step
         self._sequences_this_step = 0
         step_dir = self.output_dir / f"step_{step}"
-        step_dir.mkdir(parents=True, exist_ok=True)
+        _mkdir_existing_dir_ok(step_dir)
         self._file = (step_dir / f"rank_{self.rank}.jsonl").open("w", encoding="utf-8")
 
     def _write(self, record: dict[str, Any]) -> None:
@@ -113,6 +116,19 @@ def setup_token_exporter(
     return exporter
 
 
+def _mkdir_existing_dir_ok(path: Path) -> None:
+    for attempt in range(MKDIR_MAX_ATTEMPTS):
+        try:
+            path.mkdir(parents=True, exist_ok=True)
+            return
+        except FileExistsError:
+            if path.is_dir():
+                return
+            if attempt == MKDIR_MAX_ATTEMPTS - 1:
+                raise
+            time.sleep(MKDIR_RETRY_DELAY_SECONDS)
+
+
 def _export_columns(
     micro_batch: Mapping[str, Any], model_output: Mapping[str, Tensor], loss_config: Any
 ) -> dict[str, list[Any]]:

tests/unit/train/rl/test_token_export.py

@@ -0,0 +1,49 @@
+from pathlib import Path
+
+import pytest
+
+from prime_rl.trainer.rl import token_export
+from prime_rl.trainer.rl.token_export import _mkdir_existing_dir_ok
+
+
+def test_mkdir_existing_dir_ok_retries_transient_file_exists(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    target = tmp_path / "token_exports" / "step_31"
+    original_mkdir = Path.mkdir
+    calls = 0
+
+    def flaky_mkdir(
+        self: Path,
+        mode: int = 0o777,
+        parents: bool = False,
+        exist_ok: bool = False,
+    ) -> None:
+        nonlocal calls
+        if self == target and calls == 0:
+            calls += 1
+            raise FileExistsError(str(self))
+        original_mkdir(self, mode=mode, parents=parents, exist_ok=exist_ok)
+
+    def create_dir_during_retry(_: float) -> None:
+        original_mkdir(target, parents=True, exist_ok=True)
+
+    monkeypatch.setattr(Path, "mkdir", flaky_mkdir)
+    monkeypatch.setattr(token_export.time, "sleep", create_dir_during_retry)
+
+    _mkdir_existing_dir_ok(target)
+
+    assert target.is_dir()
+    assert calls == 1
+
+
+def test_mkdir_existing_dir_ok_raises_when_path_is_file(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    target = tmp_path / "token_exports" / "step_31"
+    target.parent.mkdir(parents=True)
+    target.write_text("not a directory", encoding="utf-8")
+    monkeypatch.setattr(token_export.time, "sleep", lambda _: None)
+
+    with pytest.raises(FileExistsError):
+        _mkdir_existing_dir_ok(target)