Fix ruff job

.github/container/nsys_jax/nsys_jax/analyses/Analysis.ipynb

@@ -7,8 +7,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from collections import defaultdict\n",
     "import functools\n",
+    "import os\n",
+    "import pathlib\n",
+    "from collections import defaultdict\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "\n",
     "from nsys_jax import (\n",
     "    align_profiler_data_timestamps,\n",
     "    apply_warmup_heuristics,\n",
@@ -19,11 +25,7 @@
     "    load_profiler_data,\n",
     "    remove_autotuning_detail,\n",
     "    xla_module_metadata,\n",
-    ")\n",
-    "import matplotlib.pyplot as plt\n",
-    "import numpy as np\n",
-    "import os\n",
-    "import pathlib"
+    ")"
    ]
   },
   {
@@ -675,7 +677,7 @@
     "    axs[1].set_xlabel(\"Execution time [ms]\")\n",
     "    axs[1].set_yticks(\n",
     "        np.arange(len(detailed_index)),\n",
-    "        labels=map(lambda idx: f\"{idx[1]} ({idx[0]})\", detailed_index),\n",
+    "        labels=(f\"{idx[1]} ({idx[0]})\" for idx in detailed_index),\n",
     "    )"
    ]
   },

.github/container/nsys_jax/nsys_jax/analyses/communication.py

@@ -5,14 +5,15 @@
 from collections import defaultdict
 from math import sqrt
 
+from prettytable import PrettyTable
+from uncertainties import ufloat  # type: ignore
+
 from nsys_jax import (
     align_profiler_data_timestamps,
     apply_warmup_heuristics,
     ensure_compiled_protos_are_importable,
     load_profiler_data,
 )
-from prettytable import PrettyTable
-from uncertainties import ufloat  # type: ignore
 
 
 def process_communication_data(steady_state):
@@ -96,7 +97,7 @@ def format_bandwidth(data, collective):
         for collective in collective_types
     }
     size_heading = "Size [B]"
-    size_width = max(len(size_heading), max(len(f"{s:,}") for s in summary_data.keys()))
+    size_width = max(len(size_heading), max(len(f"{s:,}") for s in summary_data))
 
     header_log = f"{'':<{size_width}} | Bus bandwidth [GB/s]"
     print(header_log)

.github/container/nsys_jax/nsys_jax/analyses/pgle_costs.py

@@ -1,13 +1,14 @@
 #!/usr/bin/env python
 import argparse
+import pathlib
+
 from nsys_jax import (
     apply_warmup_heuristics,
     ensure_compiled_protos_are_importable,
     load_profiler_data,
     xla_module_metadata,
 )
 from nsys_jax.protobuf import HloProto, HloProtoSet
-import pathlib
 
 
 def get_scheduling_name(module: HloProto, name: str) -> str:

.github/container/nsys_jax/nsys_jax/analyses/summary.py

@@ -1,16 +1,18 @@
 #!/usr/bin/env python
 import argparse
 import math
+import pathlib
+
+from prettytable import PrettyTable
+from uncertainties import ufloat  # type: ignore
+
 from nsys_jax import (
     apply_warmup_heuristics,
     ensure_compiled_protos_are_importable,
     generate_compilation_statistics,
     load_profiler_data,
     remove_autotuning_detail,
 )
-import pathlib
-from prettytable import PrettyTable
-from uncertainties import ufloat  # type: ignore
 
 
 def main():

.github/container/nsys_jax/nsys_jax/analysis.py

@@ -1,13 +1,14 @@
-from collections import defaultdict
 import functools
 import math
-import numpy as np
-import pandas as pd  # type: ignore
 import pathlib
+from collections import defaultdict
 from typing import Any
 
+import numpy as np
+import pandas as pd  # type: ignore
+
 from .protobuf import HloProto, _host_memory_space, xla_module_metadata
-from .utils import make_child_mask, ProfilerData
+from .utils import ProfilerData, make_child_mask
 
 pd.options.mode.copy_on_write = True
 
@@ -28,9 +29,9 @@ def align_profiler_data_timestamps(
     # Error if the communication frame doesn't exist at all, but not if it is empty.
     # Calling this on a profile that does not contain any communication should
     # gracefully yield empty results.
-    assert (
-        frames.communication is not None
-    ), "align_profiler_data_timestamps requires a communication frame"
+    assert frames.communication is not None, (
+        "align_profiler_data_timestamps requires a communication frame"
+    )
     if not len(frames.communication):
         # Nothing to be done, return an empty result
         return frames, {}
@@ -43,9 +44,9 @@ def align_profiler_data_timestamps(
             f"WARNING: cannot align {num_profiled_devices} devices because max collective size is 1"
         )
         return frames, {}
-    assert (
-        num_profiled_devices == max_collective_size
-    ), f"Aligning {num_profiled_devices} using collectives of size {max_collective_size} is not implemented"
+    assert num_profiled_devices == max_collective_size, (
+        f"Aligning {num_profiled_devices} using collectives of size {max_collective_size} is not implemented"
+    )
     # Find the collectives that will be used
     align_df = comm_df[comm_df["CollectiveSize"] == max_collective_size]
     # Calculate the collectives' end times
@@ -189,22 +190,21 @@ def _get_message_size(
 ) -> tuple[int, str, int, float, float]:
     _, inst = module_proto.find_instruction(instruction)
     comm_inst = inst.communication_proto()
-    assert (
-        comm_inst.opcode
-        in {
-            "all-gather",
-            "all-gather-start",
-            "all-reduce",
-            "all-reduce-start",
-            "all-to-all",
-            "collective-broadcast",
-            "collective-permute",
-            "collective-permute-start",
-            "dynamic-slice",
-            "dynamic-update-slice",
-            "reduce-scatter",
-        }
-    ), f"{instruction}: message size calculation for {comm_inst.opcode} has not yet been validated"
+    assert comm_inst.opcode in {
+        "all-gather",
+        "all-gather-start",
+        "all-reduce",
+        "all-reduce-start",
+        "all-to-all",
+        "collective-broadcast",
+        "collective-permute",
+        "collective-permute-start",
+        "dynamic-slice",
+        "dynamic-update-slice",
+        "reduce-scatter",
+    }, (
+        f"{instruction}: message size calculation for {comm_inst.opcode} has not yet been validated"
+    )
 
     def _byte_size(inst) -> int:
         size_bits = math.prod(
@@ -266,21 +266,21 @@ def _byte_size(inst) -> int:
                 mesh = comm_inst.mesh_axes_replica_group_list.mesh
                 axes = comm_inst.mesh_axes_replica_group_list.axes
                 assert len(axes), axes
-                assert not any(
-                    ax.HasField("sub_axis_info") for ax in axes
-                ), f"sub_axis_info not supported: {axes}"
+                assert not any(ax.HasField("sub_axis_info") for ax in axes), (
+                    f"sub_axis_info not supported: {axes}"
+                )
                 collective_size = np.prod(
                     [mesh.axes[ax.mesh_axis_index].size for ax in axes]
                 )
         else:
-            collective_sizes = set(len(group.replica_ids) for group in replica_groups)
-            assert (
-                len(collective_sizes) == 1
-            ), f"Heterogeneous collective {comm_inst} could not be interpreted"
+            collective_sizes = {len(group.replica_ids) for group in replica_groups}
+            assert len(collective_sizes) == 1, (
+                f"Heterogeneous collective {comm_inst} could not be interpreted"
+            )
             collective_size = next(iter(collective_sizes))
-        assert (
-            collective_size > 0
-        ), f"Could not extract collective size from: {comm_inst}"
+        assert collective_size > 0, (
+            f"Could not extract collective size from: {comm_inst}"
+        )
     total_msg_size = 0
     for operand_id in comm_inst.operand_ids:
         _, operand = module_proto.find_instruction_by_id(operand_id)

.github/container/nsys_jax/nsys_jax/data_loaders.py

@@ -1,18 +1,19 @@
-from collections import defaultdict
 import functools
 import itertools
 import lzma
 import multiprocessing
-import numpy as np
 import os
-import pandas as pd  # type: ignore
 import pathlib
 import re
+from collections import defaultdict
+
+import numpy as np
+import pandas as pd  # type: ignore
 
 from .analysis import calculate_collective_metrics
 from .protobuf import _hlo_cache, _remap_program_id, xla_module_metadata
 from .protobuf_utils import ensure_compiled_protos_are_importable
-from .utils import default_data_prefix, make_child_mask, ProfilerData
+from .utils import ProfilerData, default_data_prefix, make_child_mask
 
 pd.options.mode.copy_on_write = True
 
@@ -318,9 +319,7 @@ def _load_nvtx_gpu_proj_trace_single(
             not_last, last = gpu_ops[:-1], gpu_ops[-1]
             if last < np.mean(not_last) - np.std(not_last):
                 print(
-                    "Skipping last occurence of {} because it only had {} GPU operations, compared to {} +/- {} before".format(
-                        mod_name, last, np.mean(not_last), np.std(not_last)
-                    )
+                    f"Skipping last occurence of {mod_name} because it only had {last} GPU operations, compared to {np.mean(not_last)} +/- {np.std(not_last)} before"
                 )
                 mod_id = mod_name_df.index[-1]
                 mod_ids.remove(mod_id)
@@ -779,7 +778,7 @@ def _load_nvtx_pushpop_trace(prefix: pathlib.Path, frames: set[str]) -> pd.DataF
 
 def load_profiler_data(
     prefix: pathlib.Path = default_data_prefix(),
-    frames: set[str] = {"communication", "compile", "module", "thunk"},
+    frames: set[str] | None = None,
 ) -> ProfilerData:
     """
     Load post-processed Nsight Systems traces and prepare them for analysis.
@@ -794,6 +793,8 @@ def load_profiler_data(
      ProfilerData dataclass with members set according to ``frames``
     """
     # Dependency management
+    if frames is None:
+        frames = {"communication", "compile", "module", "thunk"}
     if "communication" in frames:
         frames.add("thunk")
     output = ProfilerData()

.github/container/nsys_jax/nsys_jax/protobuf.py

@@ -1,10 +1,10 @@
-from collections import defaultdict
-from collections.abc import Callable
 import functools
 import lzma
 import pathlib
 import re
 import typing
+from collections import defaultdict
+from collections.abc import Callable
 
 from .utils import default_data_prefix
 
@@ -84,15 +84,13 @@ def _host_operand(i):
                 _, op = wrapped_hlo_proto.find_instruction_by_id(inst.operand_ids[i])
                 return _host_memory_space(op.proto())
 
-            if inst.opcode == "dynamic-slice" and host_dest != _host_operand(0):
-                return True
-            elif (
-                inst.opcode == "dynamic-update-slice"
+            return bool(
+                inst.opcode == "dynamic-slice"
+                and host_dest != _host_operand(0)
+                or inst.opcode == "dynamic-update-slice"
                 and host_dest == _host_operand(0)
                 and host_dest != _host_operand(1)
-            ):
-                return True
-            return False
+            )
 
         if self._proto.opcode in comm_opcodes | comm_start_opcodes:
             self._comm_proto = self._proto
@@ -223,7 +221,7 @@ class HloProtoSet:
     xla_module_metadata with policy="all".
     """
 
-    def __init__(self, protos: dict[typing.Optional[str], HloProto]):
+    def __init__(self, protos: dict[str | None, HloProto]):
         assert len(protos), f"HloProtoSet got {len(protos)} HloProtos"
         self._protos = protos
 
@@ -290,7 +288,6 @@ def _remap_program_id(
     replica: str | None,
     allow_missing_protobuf: bool = False,
 ) -> str:
-    """ """
     # In multi-input mode, we will have something like:
     #   old_id = 1
     #   name = jit_foo
@@ -373,7 +370,7 @@ def xla_module_metadata(
     program_id: str,
     policy: str = "consistent",
     prefix: pathlib.Path = default_data_prefix(),
-) -> typing.Union[HloProto, HloProtoSet]:
+) -> HloProto | HloProtoSet:
     """
     Load the protobuf metadata for module `program_id`. If given, `prefix` is the
     search path. `policy` governs what happens if `nsys-jax-combine` found inconsistent

.github/container/nsys_jax/nsys_jax/protobuf_utils.py

@@ -7,7 +7,6 @@
 import subprocess
 import sys
 import tempfile
-from typing import Optional
 
 from .utils import default_data_prefix
 
@@ -34,7 +33,7 @@ def which(executable: str) -> pathlib.Path:
 def compile_protos(
     proto_dir: str | pathlib.Path,
     output_dir: str | pathlib.Path,
-    output_stub_dir: Optional[str | pathlib.Path] = None,
+    output_stub_dir: str | pathlib.Path | None = None,
 ):
     if not os.path.isdir(proto_dir):
         raise Exception(f"Input: {proto_dir} is not a directory")

.github/container/nsys_jax/nsys_jax/scripts/install_flamegraph.py

@@ -1,5 +1,6 @@
 import argparse
 import os
+
 import requests
 
 

.github/container/nsys_jax/nsys_jax/scripts/install_protoc.py

@@ -1,11 +1,12 @@
 import argparse
-import google.protobuf
 import io
 import os
 import platform
-import requests
 import zipfile
 
+import google.protobuf
+import requests
+
 
 def main():
     # TODO: add a default to (with confirmation) install in the same prefix as this script is installed to
@@ -25,7 +26,7 @@ def main():
     # install a protoc with the same version as google.protobuf. For newer versions, given
     # google.protobuf version X.Y.Z install protoc version Y.Z as described in
     # https://protobuf.dev/support/version-support
-    runtime_version = tuple(map(int, google.protobuf.__version__.split(".")))
+    runtime_version = tuple(int(v) for v in google.protobuf.__version__.split("."))
     if runtime_version < (3, 21):
         # old versioning scheme, try and install a matching protoc version
         protoc_version = runtime_version
@@ -53,14 +54,13 @@ def main():
         if r.status_code == 404:
             # assume this means the architecture is not available
             continue
-    else:
-        r.raise_for_status()
+    r.raise_for_status()
 
     with zipfile.ZipFile(io.BytesIO(r.content)) as z:
         for name in z.namelist():
             if ".." in name:
                 continue
-            if name.startswith("bin/") or name.startswith("include/"):
+            if name.startswith(("bin/", "include/")):
                 z.extract(name, path=args.prefix)
 
     # Make sure the protoc binary is executable