From ac3a0a6b472de275e4ca1fbdbf1a7c89aec1cb2a Mon Sep 17 00:00:00 2001
From: Jake Hillion <jake@hillion.co.uk>
Date: Tue, 9 Dec 2025 12:26:56 +0000
Subject: [PATCH] ci: enable `ruff check` in CI through nix

---
 .github/scripts/bench.py        | 21 ++++++++++-----------
 .github/scripts/build_matrix.py |  1 +
 flake.nix                       |  5 +++++
 justfile                        |  2 +-
 pyproject.toml                  |  2 +-
 5 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/.github/scripts/bench.py b/.github/scripts/bench.py
index 6b4b3ab1..0ba73d58 100644
--- a/.github/scripts/bench.py
+++ b/.github/scripts/bench.py
@@ -15,6 +15,7 @@ from __future__ import annotations
 
 import argparse
 import asyncio
+import contextlib
 import json
 import sys
 import time
@@ -193,7 +194,7 @@ def collect_metrics_snapshot(state: Mapping[str, Any]) -> MetricsSnapshot:
 
     # Count tasks per instance (only Pending and Running exist in state; completed tasks are deleted)
     instance_task_counts: dict[str, dict[str, int]] = {}
-    for instance_id in instances.keys():
+    for instance_id in instances:
         instance_task_counts[instance_id] = {
             "Pending": 0,
             "Running": 0,
@@ -492,7 +493,7 @@ async def wait_for_tasks_drained(api_base: str, timeout_s: int = 600) -> None:
     pending or running tasks remaining.
     """
     print(f"\n{'=' * 80}")
-    print(f"⏳ WAITING FOR ALL TASKS TO DRAIN")
+    print("⏳ WAITING FOR ALL TASKS TO DRAIN")
     print(f"{'=' * 80}")
     start = time.monotonic()
 
@@ -821,10 +822,8 @@ async def monitor_metrics(
             traceback.print_exc()
 
         # Wait for interval or until stopped
-        try:
+        with contextlib.suppress(asyncio.TimeoutError):
             await asyncio.wait_for(stop_event.wait(), timeout=interval_seconds)
-        except asyncio.TimeoutError:
-            pass
 
 
 async def run_stage(
@@ -883,7 +882,7 @@ async def run_stage(
         results = list(await asyncio.gather(*tasks))
 
     # Wait for all tasks in the cluster to be drained
-    print(f"\nHTTP requests completed. Now waiting for cluster tasks to drain...")
+    print("\nHTTP requests completed. Now waiting for cluster tasks to drain...")
     await wait_for_tasks_drained(api_base, timeout_s=600)
 
     stage_completed_at = time.time()
@@ -1026,7 +1025,7 @@ async def run_benchmark(
 
         # Add 30 second delay to allow topology to stabilize before creating instances
         print(
-            f"\nWaiting 30 seconds for topology to stabilize before creating instances..."
+            "\nWaiting 30 seconds for topology to stabilize before creating instances..."
         )
         await asyncio.sleep(30)
         print("Proceeding with instance creation\n")
@@ -1036,7 +1035,7 @@ async def run_benchmark(
 
         model_counts = Counter(model_ids)
 
-        print(f"\nTarget instance counts by model:")
+        print("\nTarget instance counts by model:")
         for model_id, count in model_counts.items():
             print(f"  {model_id}: {count} instance(s)")
         print()
@@ -1088,7 +1087,7 @@ async def run_benchmark(
 
         # Collect all instance IDs for all models
         state = fetch_state(api_base)
-        for model_id in model_counts.keys():
+        for model_id in model_counts:
             ids = get_all_instance_ids_for_model(state, model_id)
             all_instance_ids.extend(ids)
 
@@ -1317,7 +1316,7 @@ async def run_benchmark(
                 print(f"Saving results to: {results_output_path}")
                 with open(results_output_path, "w") as f:
                     json.dump(results_doc, f, indent=2)
-                print(f"Results saved successfully")
+                print("Results saved successfully")
 
             # Cleanup all instances
             for instance_id in all_instance_ids:
@@ -1331,7 +1330,7 @@ async def run_benchmark(
                 "[SECONDARY] Waiting with cluster (primary handles benchmark execution)"
             )
             # Secondary nodes wait until all instances of all models are deleted
-            for model_id in model_counts.keys():
+            for model_id in model_counts:
                 await wait_for_all_instances_deleted(api_base, model_id)
 
         return 0
diff --git a/.github/scripts/build_matrix.py b/.github/scripts/build_matrix.py
index a54cbf7b..2f139350 100644
--- a/.github/scripts/build_matrix.py
+++ b/.github/scripts/build_matrix.py
@@ -2,6 +2,7 @@
 import json
 import os
 from typing import NotRequired, TypedDict, cast
+
 import yaml
 
 
diff --git a/flake.nix b/flake.nix
index 523eac5e..c6141754 100644
--- a/flake.nix
+++ b/flake.nix
@@ -52,6 +52,11 @@
       {
         formatter = treefmtEval.config.build.wrapper;
         checks.formatting = treefmtEval.config.build.check inputs.self;
+        checks.lint = pkgs.runCommand "lint-check" { } ''
+          export RUFF_CACHE_DIR="$TMPDIR/ruff-cache"
+          ${pkgs.ruff}/bin/ruff check ${inputs.self}/
+          touch $out
+        '';
 
         devShells.default = pkgs.mkShell {
           packages =
diff --git a/justfile b/justfile
index 47cd4441..6f4e67e9 100644
--- a/justfile
+++ b/justfile
@@ -2,7 +2,7 @@ fmt:
     nix fmt
 
 lint:
-    uv run ruff check --fix src
+    uv run ruff check --fix
 
 test:
     uv run pytest src
diff --git a/pyproject.toml b/pyproject.toml
index b0c3e18b..d9c4715e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -118,7 +118,7 @@ environments = [
 ###
 
 [tool.ruff]
-extend-exclude = ["shared/protobufs/**"]
+extend-exclude = ["shared/protobufs/**", "*mlx_typings/**", "rust/exo_pyo3_bindings/**"]
 
 [tool.ruff.lint]
 extend-select = ["I", "N", "B", "A", "PIE", "SIM"]