fix: use per-rank jsonl instead of file lock in case that NFS does not support it (#513)

Author: rchardx

areal/tools/perf_trace_converter.py

@@ -5,6 +5,7 @@
 import os
 import sys
 from collections.abc import Sequence
+from glob import glob
 from pathlib import Path
 
 
@@ -24,19 +25,41 @@ def _load_events(path: Path) -> list[dict]:
     return events
 
 
+def _resolve_trace_files(source: Path) -> list[Path]:
+    if source.is_file():
+        return [source]
+    if source.is_dir():
+        return sorted(p for p in source.glob("*.jsonl") if p.is_file())
+    matches = [Path(p) for p in glob(str(source), recursive=True)]
+    files = [p for p in matches if p.is_file()]
+    return sorted(files)
+
+
 def convert_jsonl_to_chrome_trace(
     input_path: str | os.PathLike[str],
     output_path: str | os.PathLike[str] | None = None,
     *,
     display_time_unit: str = "ms",
 ) -> dict:
-    """Convert newline-delimited trace events into Chrome Trace JSON."""
+    """Convert newline-delimited trace events into Chrome Trace JSON.
+
+    The ``input_path`` may point to a single JSONL file, a directory containing
+    per-rank JSONL files, or a glob pattern. All matching files are concatenated
+    in lexical order before emitting the Chrome trace payload.
+    """
+
+    sources = _resolve_trace_files(Path(input_path))
+    if not sources:
+        raise FileNotFoundError(f"No trace files matched input path: {input_path}")
+
+    events: list[dict] = []
+    for path in sources:
+        events.extend(_load_events(path))
 
-    source = Path(input_path)
-    if not source.is_file():  # pragma: no cover - defensive guard
-        raise FileNotFoundError(f"Input trace file not found: {source}")
+    events.sort(
+        key=lambda event: (event.get("ts", 0), event.get("pid", 0), event.get("tid", 0))
+    )
 
-    events = _load_events(source)
     chrome_trace = {
         "traceEvents": events,
         "displayTimeUnit": display_time_unit,
@@ -55,7 +78,14 @@ def _parse_args(argv: Sequence[str] | None) -> argparse.Namespace:
     parser = argparse.ArgumentParser(
         description="Convert PerfTracer JSONL output into Chrome Trace JSON format.",
     )
-    parser.add_argument("input", type=str, help="Path to the PerfTracer JSONL file")
+    parser.add_argument(
+        "input",
+        type=str,
+        help=(
+            "Path, directory, or glob pattern for PerfTracer JSONL files "
+            "(per-rank outputs allowed)"
+        ),
+    )
     parser.add_argument(
         "output",
         type=str,

areal/utils/perf_tracer.py

@@ -64,6 +64,13 @@ class PerfTraceCategory(str, Enum):
 _REQUEST_TRACE_FILENAME = "requests.jsonl"
 
 
+def _rank_qualified_filename(filename: str, rank: int | None) -> str:
+    if rank is None:
+        return filename
+    root, ext = os.path.splitext(filename)
+    return f"{root}-r{rank}{ext}"
+
+
 def _maybe_duration(start: float | None, end: float | None) -> float | None:
     if start is None or end is None:
         return None
@@ -92,6 +99,7 @@ def _default_trace_path(
     config: PerfTracerConfig,
     *,
     filename: str = _PERF_TRACE_FILENAME,
+    rank: int | None = None,
 ) -> str:
     base_dir = os.path.join(
         os.path.expanduser(os.path.expandvars(config.fileroot)),
@@ -100,7 +108,7 @@ def _default_trace_path(
         config.experiment_name,
         config.trial_name,
     )
-    return os.path.join(base_dir, filename)
+    return os.path.join(base_dir, _rank_qualified_filename(filename, rank))
 
 
 def _normalize_flush_threshold(config: RequestTracerConfig) -> int:
@@ -315,8 +323,7 @@ def flush(self, force: bool = False) -> None:
             with _acquire_file_lock(self._output_path):
                 with open(self._output_path, "a", encoding="utf-8") as fout:
                     for line in lines:
-                        fout.write(line)
-                        fout.write("\n")
+                        fout.write(f"{line}\n")
                     fout.flush()
                     os.fsync(fout.fileno())
         except OSError as exc:  # pragma: no cover - depends on filesystem
@@ -455,7 +462,7 @@ def __init__(self, config: PerfTracerConfig, *, rank: int) -> None:
         self._origin_ns = time.perf_counter_ns()
         self._thread_meta_emitted: set[int] = set()
         self._process_meta_emitted: set[int] = set()
-        self._output_path = _default_trace_path(config)
+        self._output_path = _default_trace_path(config, rank=rank)
         self._save_interval = _normalize_save_interval(config)
         self._request_tracer: RequestTracer | None = None
         self._configure_request_tracer(config, rank=rank)
@@ -479,7 +486,11 @@ def _configure_request_tracer(self, config: PerfTracerConfig, *, rank: int) -> N
         request_cfg = getattr(config, "request_tracer", None)
         enabled = bool(request_cfg and getattr(request_cfg, "enabled", False))
         if enabled:
-            output_path = _default_trace_path(config, filename=_REQUEST_TRACE_FILENAME)
+            output_path = _default_trace_path(
+                config,
+                filename=_REQUEST_TRACE_FILENAME,
+                rank=rank,
+            )
             if self._request_tracer is None:
                 self._request_tracer = RequestTracer(
                     request_cfg,
@@ -500,7 +511,7 @@ def _configure_request_tracer(self, config: PerfTracerConfig, *, rank: int) -> N
     def apply_config(self, config: PerfTracerConfig, *, rank: int) -> None:
         self._config = config
         self.set_rank(rank)
-        self._output_path = _default_trace_path(config)
+        self._output_path = _default_trace_path(config, rank=rank)
         self.set_enabled(config.enabled)
         self._save_interval = _normalize_save_interval(config)
         self._configure_request_tracer(config, rank=rank)
@@ -601,8 +612,7 @@ def save(self, *, step: int | None = None, force: bool = False) -> None:
                 with _acquire_file_lock(output_path):
                     with open(output_path, "a", encoding="utf-8") as fout:
                         for line in serialized_events:
-                            fout.write(line)
-                            fout.write("\n")
+                            fout.write(f"{line}\n")
                         fout.flush()
                         os.fsync(fout.fileno())
                 self._events = []