style: fix ruff lint errors and apply ruff format

Fix 39 ruff lint violations (F401 unused imports, D208/D209 docstring
formatting, D205 missing blank line after summary) and reformat 28
files with ruff format.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

Author: ZhiyuLi-Nvidia

examples/run_grpo.py

@@ -108,6 +108,7 @@ def main() -> None:
 
         def _make_policy(**kwargs):
             return TQPolicy(**kwargs, dp_cfg=_dp_cfg)
+
         _policy_factory = _make_policy
     else:
         _policy_factory = None  # setup() defaults to plain Policy
@@ -124,7 +125,10 @@ def _make_policy(**kwargs):
         grpo_state,
         master_config,
     ) = setup(
-        config, tokenizer, dataset, val_dataset,
+        config,
+        tokenizer,
+        dataset,
+        val_dataset,
         policy_factory=_policy_factory,
     )
 

nemo_rl/algorithms/grpo_sync.py

@@ -97,12 +97,19 @@ def _apply_dynamic_sampling(
     max_gen_batches: int,
     dp_client: DataPlaneClient,
 ) -> tuple[
-    Optional[KVBatchMeta], Optional[_DSlice],
-    list[torch.Tensor], bool, dict[str, Any], Optional[torch.Tensor],
+    Optional[KVBatchMeta],
+    Optional[_DSlice],
+    list[torch.Tensor],
+    bool,
+    dict[str, Any],
+    Optional[torch.Tensor],
 ]:
-    """One iteration. Returns (pending_meta, pending_slice, pending_rewards,
+    """One iteration.
+
+    Returns (pending_meta, pending_slice, pending_rewards,
     is_complete, ds_metrics, unfiltered_for_log). When complete, the returned
-    pending_* IS the training batch."""
+    pending_* IS the training batch.
+    """
     # Cumulative unfiltered total_reward for legacy metrics["reward"]
     # parity. Reference-only append (no copy) — slice tensors are
     # produced fresh per iteration, not aliased to TQ-owned bulk.
@@ -145,7 +152,9 @@ def _apply_dynamic_sampling(
         )
         pending_meta = pending_meta.slice(0, train_prompts_size)
         pending_slice = pending_slice.slice(0, train_prompts_size)
-        ds_metrics["dynamic_sampling_num_discarded_valid_samples"] = n - train_prompts_size
+        ds_metrics["dynamic_sampling_num_discarded_valid_samples"] = (
+            n - train_prompts_size
+        )
 
     unfiltered_for_log = torch.cat(pending_unfiltered_rewards)[:train_prompts_size]
     return pending_meta, pending_slice, [], True, ds_metrics, unfiltered_for_log
@@ -404,9 +413,7 @@ def grpo_train_sync(
                 # partition exists with the expected schema.
                 policy.prepare_step(
                     num_samples=int(repeated_batch.size),
-                    group_size=master_config["grpo"][
-                        "num_generations_per_prompt"
-                    ],
+                    group_size=master_config["grpo"]["num_generations_per_prompt"],
                 )
 
                 # ── Rollout 1-hop put: actor runs rollout + flatten +
@@ -462,11 +469,13 @@ def grpo_train_sync(
                 # touched by any of these ops).
                 with timer.time("reward_calculation"):
                     slice_data = scale_rewards(
-                        slice_data, master_config["grpo"]["reward_scaling"],
+                        slice_data,
+                        master_config["grpo"]["reward_scaling"],
                     )
                     if master_config["grpo"]["reward_shaping"]["enabled"]:
                         slice_data = apply_reward_shaping(
-                            slice_data, master_config["grpo"]["reward_shaping"],
+                            slice_data,
+                            master_config["grpo"]["reward_shaping"],
                         )
                     if master_config["grpo"]["overlong_filtering"]:
                         lm = slice_data["loss_multiplier"].clone()
@@ -495,9 +504,11 @@ def grpo_train_sync(
                             * master_config["grpo"]["num_generations_per_prompt"]
                         )
                         (
-                            pending_meta, pending_slice,
+                            pending_meta,
+                            pending_slice,
                             pending_unfiltered_rewards,
-                            is_complete, ds_metrics,
+                            is_complete,
+                            ds_metrics,
                             unfiltered_rewards_for_logging,
                         ) = _apply_dynamic_sampling(
                             meta=meta,
@@ -571,7 +582,8 @@ def grpo_train_sync(
                         "skip_reference_policy_logprobs_calculation"
                     ):
                         _ref_lp = policy.get_reference_policy_logprobs_from_meta(
-                            meta, timer=timer,
+                            meta,
+                            timer=timer,
                         )
                         reference_policy_logprobs = _ref_lp["reference_logprobs"]
                     else:
@@ -582,7 +594,8 @@ def grpo_train_sync(
                     # output_ids, attention_mask, position_ids) stays in
                     # TQ — workers will fetch it via ``train_presharded``.
                     extras_bdd = read_columns(
-                        policy._dp_client, meta,
+                        policy._dp_client,
+                        meta,
                         select_fields=["generation_logprobs", "token_mask"],
                         pad_value_dict=_pad_dict,
                     )
@@ -658,7 +671,8 @@ def grpo_train_sync(
                 # sample_mask under the same meta.keys so workers fetch
                 # the union via train_presharded.
                 write_columns(
-                    policy._dp_client, meta,
+                    policy._dp_client,
+                    meta,
                     fields={
                         "advantages": advantages,
                         "sample_mask": sample_mask,
@@ -696,20 +710,27 @@ def grpo_train_sync(
                         # mask / adv columns added later are irrelevant
                         # here.
                         _calib_fields = [
-                            f for f in (meta.fields or [])
-                            if f not in (
-                                "generation_logprobs", "token_mask",
-                                "sample_mask", "prev_logprobs",
-                                "reference_policy_logprobs", "advantages",
+                            f
+                            for f in (meta.fields or [])
+                            if f
+                            not in (
+                                "generation_logprobs",
+                                "token_mask",
+                                "sample_mask",
+                                "prev_logprobs",
+                                "reference_policy_logprobs",
+                                "advantages",
                             )
                         ]
                         calibration_data = read_columns(
-                            policy._dp_client, meta,
+                            policy._dp_client,
+                            meta,
                             select_fields=_calib_fields,
                             pad_value_dict=_pad_dict,
                         )
                         kv_scales_cache = policy.calibrate_qkv_fp8_scales(
-                            calibration_data, include_q=True,
+                            calibration_data,
+                            include_q=True,
                         )["layers"]
                         POLICY_GENERATION_STALE = True
 
@@ -726,15 +747,18 @@ def grpo_train_sync(
                     if "content" in (meta.fields or []):
                         _log_select.append("content")
                     _log_extras = read_columns(
-                        policy._dp_client, meta, select_fields=_log_select,
+                        policy._dp_client,
+                        meta,
+                        select_fields=_log_select,
                         pad_value_dict=_pad_dict,
                     )
                     _log_input_ids = _log_extras["input_ids"]
                     _log_content = _log_extras.get("content")
 
                 # ── Step-end TQ cleanup ────────────────────────────────
                 policy._dp_client.kv_clear(
-                    keys=meta.keys, partition_id=meta.partition_id,
+                    keys=meta.keys,
+                    partition_id=meta.partition_id,
                 )
 
                 is_last_step = total_steps + 1 >= max_num_steps
@@ -779,9 +803,7 @@ def grpo_train_sync(
 
                 # advantages and token_mask are in scope from the
                 # advantage / masking blocks above. No need to re-fetch.
-                response_advantages = torch.masked_select(
-                    advantages, token_mask.bool()
-                )
+                response_advantages = torch.masked_select(advantages, token_mask.bool())
 
                 memory_tracker.snapshot_start_of_stage("Metrics", dir())
                 metrics = {
@@ -1033,7 +1055,9 @@ def grpo_train_sync(
             print(f"  • Generation KL Error: {metrics['gen_kl_error']:.4f}")
             if master_config["grpo"]["use_dynamic_sampling"]:
                 print(f"  • Avg Filtered Reward: {np.mean(rewards.numpy()):.4f}")
-                print(f"  • Avg Total Reward: {np.mean(unfiltered_rewards.numpy()):.4f}")
+                print(
+                    f"  • Avg Total Reward: {np.mean(unfiltered_rewards.numpy()):.4f}"
+                )
             else:
                 print(f"  • Avg Reward: {np.mean(rewards.numpy()):.4f}")
             print(

nemo_rl/data_plane/adapters/transfer_queue.py

@@ -110,8 +110,9 @@ def _mooncake_transport_config() -> dict:
 
 
 def _connect_existing() -> None:
-    """Worker-process path: connect this process's client to the
-    already-running named controller actor in the Ray cluster. Mirrors
+    """Worker-process path: connect this process's client to the Ray cluster.
+
+    Connects to the already-running named controller actor. Mirrors
     rl-arena/arena/dataplane_client.py's `tq.init()` (no args) call.
     """
     _tq().init()
@@ -121,9 +122,10 @@ def _connect_existing() -> None:
 
 
 def _patch_tq_actor_runtime_env() -> None:
-    """Inject Ray ``runtime_env={"pip": ["TransferQueue==0.1.6"]}`` into the
-    ``.options()`` calls on TQ's internal actor classes (``SimpleStorageUnit``,
-    ``TransferQueueController``).
+    """Inject Ray ``runtime_env`` into TQ's internal actor class ``.options()`` calls.
+
+    Injects ``{"pip": ["TransferQueue==0.1.6"]}`` into ``.options()`` for
+    ``SimpleStorageUnit`` and ``TransferQueueController``.
 
     **Why**: TQ spawns these actors via ``Cls.options(...).remote(...)`` with
     no runtime_env. They inherit the *job-level* runtime_env that the driver
@@ -317,6 +319,7 @@ def _to_wire(td: TensorDict) -> TensorDict:
     # metadata-recorded shape. materialize squeezes the trailing 1
     # back on read so consumers see (N,).
     from nemo_rl.data_plane.codec import _KV_PROMOTE_1D as _promote_1d
+
     if _promote_1d:
         new_dict: dict[str, torch.Tensor] = {}
         changed = False
@@ -391,6 +394,7 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None:
                 os.environ["MC_TCP_BIND_ADDRESS"] = local_ip
             os.environ.setdefault("MC_STORE_MEMCPY", "0")
             from nemo_rl.data_plane.codec import set_kv_promote_1d
+
             set_kv_promote_1d(True)
 
         if bootstrap:

nemo_rl/data_plane/codec.py

@@ -13,28 +13,28 @@
 # limitations under the License.
 """Wire <-> trainer codec — jagged-on-the-wire bridge.
 
-  * Writer side: variable-length fields are encoded as
-    ``torch.nested.nested_tensor`` with ``layout=torch.jagged`` before
-    ``kv_batch_put``. Padding tax is paid only when a consumer needs a
-    rectangular tensor.
-
-  * Reader side: :func:`materialize` accepts the wire TensorDict and,
-    when ``layout='padded'``, calls
-    :func:`torch.nested.to_padded_tensor` on any nested leaves using
-    the per-field padding value supplied in ``pad_value_dict``. Trainer
-    code consumes the padded BatchedDataDict unchanged.
-
-  * Worker write-backs that produce ``response``-shaped outputs use
-    :func:`response_from_nested` to extract the response slice from a
-    (prompt+response) nested tensor.
-
-  * Non-tensor object fields (verl-style ``np.ndarray(dtype=object)``)
-    ride the same wire as variable-length tensors: each row is pickled
-    to ``bytes`` and packed into a jagged uint8 nested tensor via
-    :func:`pack_object_array`. Reader unpacks via
-    :func:`unpack_object_array` and emits the field as an object array
-    in the materialized BatchedDataDict. Backends see only tensors —
-    no per-backend non-tensor support required.
+* Writer side: variable-length fields are encoded as
+``torch.nested.nested_tensor`` with ``layout=torch.jagged`` before
+``kv_batch_put``. Padding tax is paid only when a consumer needs a
+rectangular tensor.
+
+* Reader side: :func:`materialize` accepts the wire TensorDict and,
+when ``layout='padded'``, calls
+:func:`torch.nested.to_padded_tensor` on any nested leaves using
+the per-field padding value supplied in ``pad_value_dict``. Trainer
+code consumes the padded BatchedDataDict unchanged.
+
+* Worker write-backs that produce ``response``-shaped outputs use
+:func:`response_from_nested` to extract the response slice from a
+(prompt+response) nested tensor.
+
+* Non-tensor object fields (verl-style ``np.ndarray(dtype=object)``)
+ride the same wire as variable-length tensors: each row is pickled
+to ``bytes`` and packed into a jagged uint8 nested tensor via
+:func:`pack_object_array`. Reader unpacks via
+:func:`unpack_object_array` and emits the field as an object array
+in the materialized BatchedDataDict. Backends see only tensors —
+no per-backend non-tensor support required.
 """
 
 from __future__ import annotations
@@ -101,8 +101,10 @@ def to_nested_by_length(
 
 
 def set_kv_promote_1d(enabled: bool) -> None:
-    """Adapter hook: when True, writer unsqueezes 1D bulk fields to
-    (N, 1) and reader squeezes the trailing 1 in :func:`materialize`.
+    """Adapter hook: enable/disable 1D→(N,1) promotion for bulk fields.
+
+    When True, writer unsqueezes 1D bulk fields to (N, 1) and reader
+    squeezes the trailing 1 in :func:`materialize`.
 
     Required by backends that go through TQ's KVStorageManager path
     (mooncake_cpu) — see ``_KV_PROMOTE_1D`` above for the schema/data
@@ -157,9 +159,7 @@ def pack_object_array(arr: "np.ndarray | list[Any]") -> torch.Tensor:
     """
     if isinstance(arr, np.ndarray):
         if arr.dtype != object:
-            raise TypeError(
-                f"pack_object_array expects dtype=object; got {arr.dtype}"
-            )
+            raise TypeError(f"pack_object_array expects dtype=object; got {arr.dtype}")
         items: list[Any] = list(arr)
     elif isinstance(arr, list):
         items = arr
@@ -173,9 +173,7 @@ def pack_object_array(arr: "np.ndarray | list[Any]") -> torch.Tensor:
         b = pickle.dumps(item, protocol=pickle.HIGHEST_PROTOCOL)
         # np.frombuffer + .copy() avoids the "non-writable buffer" warning
         # and severs the lifetime tie to the bytes object.
-        rows.append(
-            torch.from_numpy(np.frombuffer(b, dtype=np.uint8).copy())
-        )
+        rows.append(torch.from_numpy(np.frombuffer(b, dtype=np.uint8).copy()))
     return torch.nested.as_nested_tensor(rows, layout=torch.jagged)
 
 
@@ -263,9 +261,7 @@ def response_from_nested(
     response_list = []
     for resp_len, seq_offset in zip(response_lens, offsets[1:], strict=True):
         # left-shift output by one token for log_probs / values
-        response_list.append(
-            values[seq_offset - resp_len - 1 : seq_offset - 1]
-        )
+        response_list.append(values[seq_offset - resp_len - 1 : seq_offset - 1])
     return torch.nested.as_nested_tensor(response_list, layout=torch.jagged)
 
 

nemo_rl/data_plane/driver_io.py

@@ -96,12 +96,8 @@ def write_columns(
     from nemo_rl.data_plane.codec import maybe_pack_jagged, pack_object_array
 
     seq_lens = meta.sequence_lengths
-    lengths = (
-        torch.tensor(seq_lens, dtype=torch.long) if seq_lens is not None else None
-    )
-    registered_objects = set(
-        (meta.extra_info or {}).get(META_OBJECT_FIELDS, ())
-    )
+    lengths = torch.tensor(seq_lens, dtype=torch.long) if seq_lens is not None else None
+    registered_objects = set((meta.extra_info or {}).get(META_OBJECT_FIELDS, ()))
 
     packed: dict[str, torch.Tensor] = {}
     for k, v in fields.items():
@@ -127,5 +123,7 @@ def write_columns(
 
     td = TensorDict(packed, batch_size=[len(meta.keys)])
     dp_client.kv_batch_put(
-        keys=meta.keys, partition_id=meta.partition_id, fields=td,
+        keys=meta.keys,
+        partition_id=meta.partition_id,
+        fields=td,
     )

nemo_rl/data_plane/interfaces.py

@@ -112,7 +112,9 @@ def _replace(
             task_name=self.task_name,
             keys=list(keys),
             fields=self.fields,
-            sequence_lengths=list(sequence_lengths) if sequence_lengths is not None else None,
+            sequence_lengths=list(sequence_lengths)
+            if sequence_lengths is not None
+            else None,
             extra_info=dict(self.extra_info or {}),
         )