Address PR review feedback

- Remove dead gc_targets variable and unused n_checked counter
- Rename avg_plies_to_forfeit → avg_plies_completed (completed games
  contribute their full game_length to the average)
- Free all GPU tensors in game completion eval cleanup
- Move chess_engine import to top of trainer.py
- Extract shift_legal_mask() into pawn/data.py to deduplicate the
  np.roll + zero-fill pattern between data.py and trainer.py
- Use math.ceil for fractional CPU counts in cgroup detection

Author: thomas-schweich

pawn/data.py

@@ -240,6 +240,22 @@ def _watchdog():
             step += 1
 
 
+def shift_legal_mask(mask: np.ndarray) -> np.ndarray:
+    """Shift a legal move mask forward by one ply to align with CLM targets.
+
+    The engine's legal masks are indexed by position *before* each move:
+    mask[ply] = legal moves at position ply.  But CLM targets[ply] is the
+    *next* move (= move_ids[ply+1]), so we need legal moves at position
+    ply+1.  This rolls the mask by -1 along the ply axis and zeros the
+    last entry (no next move at the final ply).
+
+    Works for any mask shape (B, T, ...).
+    """
+    shifted = np.roll(mask, -1, axis=1)
+    shifted[:, -1] = 0
+    return shifted
+
+
 def create_validation_set(
     n_games: int, max_ply: int, seed: int,
     discard_ply_limit: bool = False,
@@ -267,13 +283,10 @@ def create_validation_set(
     }
 
     # Compute legal move masks for evaluating legal move rate.
-    # legal_grid[ply] contains legal moves at the position *before* move_ids[ply],
-    # but targets[ply] is the *next* move (= move_ids[ply+1]).  Shift by one so
-    # legal_grid[ply] aligns with targets[ply].
+    # Shift by one ply so legal_grid[ply] aligns with targets[ply]
+    # (see shift_legal_mask docstring).
     legal_grid, _legal_promo = engine.compute_legal_move_masks(move_ids, game_lengths)
-    legal_grid = np.roll(legal_grid, -1, axis=1)
-    legal_grid[:, -1, :] = 0  # last ply has no next move
-    batch["legal_grid"] = torch.from_numpy(legal_grid).long()
+    batch["legal_grid"] = torch.from_numpy(shift_legal_mask(legal_grid)).long()
     batch["game_lengths"] = torch.from_numpy(game_lengths).long()
 
     if no_outcome and prepend_outcome:

pawn/trainer.py

@@ -21,9 +21,10 @@
 import torch.nn.functional as F
 from torch.utils.data import DataLoader
 
+import chess_engine as engine
 from pawn.config import CLMConfig, TrainingConfig
 from pawn.model import PAWNCLM
-from pawn.data import CLMDataset, create_validation_set
+from pawn.data import CLMDataset, create_validation_set, shift_legal_mask
 from pawn.logging import MetricsLogger
 
 from pawn.data_utils import unpack_grid
@@ -260,26 +261,25 @@ def compute_game_completion(
     Returns dict with:
         game_completion_rate: fraction of games with zero illegal moves
         avg_pct_completion: mean fraction of game completed before forfeit
-        avg_plies_to_forfeit: mean plies before first illegal move (inf-free)
+        avg_plies_completed: mean plies completed before first illegal move.
+            Games with no illegal moves contribute their full game_length.
     """
     B, T = preds.shape
 
     with torch.no_grad():
         n_complete = 0
         pct_completions = []
-        plies_to_forfeit = []
+        plies_completed = []
 
         for b in range(B):
             gl = min(int(game_lengths[b].item()), T)
             forfeit_ply = -1
-            n_checked = 0
             for p in range(gl):
                 if not loss_mask[b, p]:
                     continue
                 # Skip plies with no legal moves (end-of-game padding)
                 if not legal_mask[b, p].any():
                     continue
-                n_checked += 1
                 token = int(preds[b, p].item())
                 if token < legal_mask.shape[2] and not legal_mask[b, p, token]:
                     forfeit_ply = p
@@ -291,15 +291,15 @@ def compute_game_completion(
             if forfeit_ply < 0:
                 n_complete += 1
                 pct_completions.append(1.0)
-                plies_to_forfeit.append(float(gl))
+                plies_completed.append(float(gl))
             else:
                 pct_completions.append(forfeit_ply / gl if gl > 0 else 0.0)
-                plies_to_forfeit.append(float(forfeit_ply))
+                plies_completed.append(float(forfeit_ply))
 
     return {
         "game_completion_rate": n_complete / B if B > 0 else 0.0,
         "avg_pct_completion": sum(pct_completions) / len(pct_completions) if pct_completions else 0.0,
-        "avg_plies_to_forfeit": sum(plies_to_forfeit) / len(plies_to_forfeit) if plies_to_forfeit else 0.0,
+        "avg_plies_completed": sum(plies_completed) / len(plies_completed) if plies_completed else 0.0,
     }
 
 
@@ -582,10 +582,8 @@ def evaluate(self) -> dict[str, float]:
         # without picking an illegal move?  Uses a small subset to avoid
         # materializing a large dense (B, T, vocab) token mask.
         if "game_lengths" in self.val_data:
-            import chess_engine as engine_mod
             gc_n = min(64, n)
             gc_input = self.val_data["input_ids"][:gc_n].to(self.device)
-            gc_targets = self.val_data["targets"][:gc_n].to(self.device)
             gc_loss_mask = self.val_data["loss_mask"][:gc_n].to(self.device)
             gc_game_lengths = self.val_data["game_lengths"][:gc_n].to(self.device)
             move_ids = self.val_data["input_ids"][:gc_n].numpy().astype(np.int16)
@@ -598,18 +596,17 @@ def evaluate(self) -> dict[str, float]:
                     gc_logits = model.lm_head(hidden)
                 gc_preds = gc_logits.argmax(dim=-1)
 
-            # Dense legal token mask, shifted to align with targets
-            legal_tokens = engine_mod.compute_legal_token_masks(move_ids, gl_np, vocab_size)
-            legal_tokens = np.roll(legal_tokens, -1, axis=1)
-            legal_tokens[:, -1, :] = False
-            legal_mask_t = torch.from_numpy(legal_tokens).to(self.device)
+            legal_tokens = engine.compute_legal_token_masks(move_ids, gl_np, vocab_size)
+            legal_mask_t = torch.from_numpy(
+                shift_legal_mask(legal_tokens)
+            ).to(self.device)
 
             gc = compute_game_completion(gc_preds, legal_mask_t, gc_loss_mask, gc_game_lengths)
             avg["val/game_completion_rate"] = gc["game_completion_rate"]
             avg["val/avg_pct_completion"] = gc["avg_pct_completion"]
-            avg["val/avg_plies_to_forfeit"] = gc["avg_plies_to_forfeit"]
+            avg["val/avg_plies_completed"] = gc["avg_plies_completed"]
 
-            del legal_mask_t, gc_logits, gc_preds
+            del gc_input, gc_loss_mask, gc_game_lengths, legal_mask_t, gc_logits, gc_preds
             if self.device != "cpu" and torch.cuda.is_available():
                 torch.cuda.empty_cache()
 
@@ -705,7 +702,7 @@ def _graceful_exit(signum, frame):
                     if "val/game_completion_rate" in val_metrics:
                         val_msg += (
                             f" | complete {val_metrics['val/game_completion_rate']:.3f}"
-                            f" | avg_ply {val_metrics['val/avg_plies_to_forfeit']:.0f}"
+                            f" | avg_ply {val_metrics['val/avg_plies_completed']:.0f}"
                         )
 
                     # Compound early stopping

scripts/benchmark.py

@@ -1095,12 +1095,13 @@ def _collect_cpu_cache() -> dict[str, str]:
 
 def _cgroup_cpu_count() -> int | None:
     """Return container CPU limit from cgroups, or None if unconstrained."""
+    import math as _math
     # cgroup v2: cpu.max contains "quota period" (e.g. "200000 100000" = 2 CPUs)
     try:
         text = Path("/sys/fs/cgroup/cpu.max").read_text().strip()
         quota_s, period_s = text.split()
         if quota_s != "max":
-            return max(1, int(int(quota_s) / int(period_s)))
+            return max(1, _math.ceil(int(quota_s) / int(period_s)))
     except (OSError, ValueError):
         pass
 
@@ -1109,7 +1110,7 @@ def _cgroup_cpu_count() -> int | None:
         quota = int(Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us").read_text().strip())
         if quota > 0:
             period = int(Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us").read_text().strip())
-            return max(1, int(quota / period))
+            return max(1, _math.ceil(quota / period))
     except (OSError, ValueError):
         pass