[4/N] Tiny enable UP ruleset in Ruff (#994)

Author: fzyzcjy

examples/eval/eval_delegate.py

@@ -1,6 +1,7 @@
 import logging
+from collections.abc import Iterable, Mapping, Sequence
 from dataclasses import dataclass, field, fields
-from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence
+from typing import Any, Optional
 
 from omegaconf import OmegaConf
 
@@ -14,7 +15,7 @@ def _first_not_none(*values: Any) -> Any:
     return None
 
 
-def _pick_from_mapping(data: Optional[Mapping[str, Any]], keys: Iterable[str]) -> Any:
+def _pick_from_mapping(data: Mapping[str, Any] | None, keys: Iterable[str]) -> Any:
     if not data:
         return None
     for key in keys:
@@ -28,11 +29,11 @@ class EvalEnvDatasetConfig:
     """Dataset-level generation parameters shared across delegate clients."""
 
     name: str = ""
-    n_samples_per_eval_prompt: Optional[int] = None
-    temperature: Optional[float] = None
-    top_p: Optional[float] = None
-    top_k: Optional[int] = None
-    max_response_len: Optional[int] = None
+    n_samples_per_eval_prompt: int | None = None
+    temperature: float | None = None
+    top_p: float | None = None
+    top_k: int | None = None
+    max_response_len: int | None = None
 
     # TODO: This is ugly, temporarily leave this. We should unify all the config name for dataset, default, and args. (advice from Tom.)
     FIELD_SPECS = {
@@ -75,7 +76,7 @@ def parse(cls, args, dataset_cfg: Mapping[str, Any], defaults: Mapping[str, Any]
                 "Colon in dataset name is not allowed; use `n_samples_per_eval_prompt` to configure samples per prompt."
             )
 
-        values: Dict[str, Any] = {"name": name}
+        values: dict[str, Any] = {"name": name}
         for field_name, spec in cls.FIELD_SPECS.items():
             dataset_value = _pick_from_mapping(dataset_cfg, spec["dataset_keys"])
             default_value = _pick_from_mapping(defaults, spec["default_keys"])
@@ -88,9 +89,9 @@ def parse(cls, args, dataset_cfg: Mapping[str, Any], defaults: Mapping[str, Any]
             obj = cls(**obj)
         return obj
 
-    def to_payload(self) -> Dict[str, Any]:
+    def to_payload(self) -> dict[str, Any]:
         """Return a JSON-serializable payload for this dataset configuration."""
-        payload: Dict[str, Any] = {}
+        payload: dict[str, Any] = {}
         for field_info in fields(self):
             value = getattr(self, field_info.name)
             if value is None:
@@ -104,11 +105,11 @@ class EvalEnvConfig:
     """Environment definition shared across delegate implementations."""
 
     name: str = ""
-    url: Optional[str] = None
+    url: str | None = None
     timeout_secs: int = 3600
     max_retries: int = 1
-    headers: Dict[str, Any] = field(default_factory=dict)
-    defaults: Dict[str, Any] = field(default_factory=dict)
+    headers: dict[str, Any] = field(default_factory=dict)
+    defaults: dict[str, Any] = field(default_factory=dict)
 
     @classmethod
     def parse(cls, raw: Mapping[str, Any], defaults: Mapping[str, Any]) -> "EvalEnvConfig":
@@ -121,9 +122,9 @@ def parse(cls, raw: Mapping[str, Any], defaults: Mapping[str, Any]) -> "EvalEnvC
 
 
 def _rebuild_delegate_config(
-    args, raw_delegate_config: Optional[Sequence[Mapping[str, Any]]], defaults: Optional[Mapping[str, Any]]
-) -> List[EvalEnvConfig]:
-    envs: List[EvalEnvConfig] = []
+    args, raw_delegate_config: Sequence[Mapping[str, Any]] | None, defaults: Mapping[str, Any] | None
+) -> list[EvalEnvConfig]:
+    envs: list[EvalEnvConfig] = []
     defaults = defaults or {}
     for env in raw_delegate_config or []:
         env_name = str(env.get("name", "")).strip().lower()
@@ -151,13 +152,13 @@ class EvalClient:
     def __init__(self, name: str):
         self.name = name
 
-    def evaluate(self, args, rollout_id: int) -> tuple[Dict[str, Any], Dict[str, Any]]:
+    def evaluate(self, args, rollout_id: int) -> tuple[dict[str, Any], dict[str, Any]]:
         raise NotImplementedError("Subclasses must implement this method")
 
 
-def _flatten(result: Dict[str, Any], prefix: Optional[str] = None) -> Dict[str, Any]:
+def _flatten(result: dict[str, Any], prefix: str | None = None) -> dict[str, Any]:
     """Flatten nested metric dicts into slash separated keys."""
-    flattened: Dict[str, Any] = {}
+    flattened: dict[str, Any] = {}
     for key, value in (result or {}).items():
         full_key = f"{prefix}/{key}" if prefix else key
         if isinstance(value, dict):
@@ -174,15 +175,13 @@ def __init__(self, delegates: Sequence[EvalClient]):
         self._delegates = list(delegates)
 
     @classmethod
-    def maybe_create(
-        cls, args, env_configs: Optional[Sequence[EvalEnvConfig]] = None
-    ) -> Optional["EvalDelegateClient"]:
+    def maybe_create(cls, args, env_configs: Sequence[EvalEnvConfig] | None = None) -> Optional["EvalDelegateClient"]:
         env_configs = list(env_configs) if env_configs is not None else getattr(args, "eval_delegate_config", None)
         if not env_configs:
             return None
 
         router_addr = f"http://{args.sglang_router_ip}:{args.sglang_router_port}"
-        delegates: List[EvalClient] = []
+        delegates: list[EvalClient] = []
         for env_cfg in env_configs:
             delegate = cls._create_delegate(env_cfg, router_addr)
             if delegate is not None:
@@ -201,9 +200,9 @@ def _create_delegate(env_cfg: EvalEnvConfig, router_addr: str):
         logger.warning("No delegate client registered for environment: %s", env_name)
         return None
 
-    def evaluate(self, args, rollout_id: int) -> tuple[Dict[str, Any], Dict[str, Any]]:
-        aggregated_metrics: Dict[str, Any] = {}
-        raw_responses: Dict[str, Any] = {}
+    def evaluate(self, args, rollout_id: int) -> tuple[dict[str, Any], dict[str, Any]]:
+        aggregated_metrics: dict[str, Any] = {}
+        raw_responses: dict[str, Any] = {}
         for delegate in self._delegates:
             metrics, response = delegate.evaluate(args, rollout_id)
             if metrics:

examples/eval/eval_delegate_rollout.py

@@ -3,7 +3,7 @@
 import logging
 import os
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from examples.eval.eval_delegate import EvalDelegateClient, _rebuild_delegate_config
 from omegaconf import OmegaConf
@@ -13,7 +13,7 @@
 
 logger = logging.getLogger(__name__)
 
-_DELEGATE_CACHE: dict[str, tuple[Optional[float], Optional[EvalDelegateClient]]] = {}
+_DELEGATE_CACHE: dict[str, tuple[float | None, EvalDelegateClient | None]] = {}
 
 
 def generate_rollout(
@@ -32,7 +32,7 @@ def generate_rollout(
     return result
 
 
-def _get_delegate_client(args) -> Optional[EvalDelegateClient]:
+def _get_delegate_client(args) -> EvalDelegateClient | None:
     config_path = getattr(args, "eval_config", None)
     if not config_path:
         return None
@@ -48,7 +48,7 @@ def _get_delegate_client(args) -> Optional[EvalDelegateClient]:
     return client
 
 
-def _build_delegate_client(args, config_path: str) -> Optional[EvalDelegateClient]:
+def _build_delegate_client(args, config_path: str) -> EvalDelegateClient | None:
     cfg = OmegaConf.load(config_path)
     cfg_dict = OmegaConf.to_container(cfg, resolve=True)
     if not isinstance(cfg_dict, dict):
@@ -70,22 +70,22 @@ def _build_delegate_client(args, config_path: str) -> Optional[EvalDelegateClien
     return EvalDelegateClient.maybe_create(args, env_configs=env_configs)
 
 
-def _safe_mtime(path: str) -> Optional[float]:
+def _safe_mtime(path: str) -> float | None:
     try:
         return os.path.getmtime(path)
     except OSError:
         return None
 
 
-def _log_delegate_metrics(args, rollout_id: int, metrics: dict | None, raw_response: Optional[dict]) -> dict:
+def _log_delegate_metrics(args, rollout_id: int, metrics: dict | None, raw_response: dict | None) -> dict:
     flattened = _flatten_metrics(metrics)
     if raw_response is not None:
         logger.info("External eval raw response for rollout %s: %s", rollout_id, raw_response)
     logger.info("eval %s (external): %s", rollout_id, flattened)
     return flattened
 
 
-def _flatten_metrics(metric_source: Optional[dict]) -> dict:
+def _flatten_metrics(metric_source: dict | None) -> dict:
     flattened_metrics: dict[str, float] = {}
     if not isinstance(metric_source, dict):
         return flattened_metrics

examples/eval/nemo_skills/skills_client.py

@@ -1,6 +1,6 @@
 import logging
 import time
-from typing import Any, Dict, Optional
+from typing import Any
 
 import requests
 from examples.eval.eval_delegate import EvalClient, EvalDelegateError
@@ -28,7 +28,7 @@ def from_config(cls, config: SkillsEvalEnvConfig, router_url: str):
             return None
         return cls(config, router_url)
 
-    def evaluate(self, args, rollout_id: int) -> tuple[Dict[str, Any], Dict[str, Any]]:
+    def evaluate(self, args, rollout_id: int) -> tuple[dict[str, Any], dict[str, Any]]:
         if not self._config.datasets:
             logger.warning("No Skills datasets configured; skipping delegate evaluation.")
             return {}, {}
@@ -38,7 +38,7 @@ def evaluate(self, args, rollout_id: int) -> tuple[Dict[str, Any], Dict[str, Any
         metrics = response["raw_metrics"]
         return metrics, response
 
-    def _build_payload(self, args, rollout_id: int) -> Dict[str, Any]:
+    def _build_payload(self, args, rollout_id: int) -> dict[str, Any]:
         benchmarks = [cfg.to_payload() for cfg in self._config.datasets]
         benchmarks = [cfg for cfg in benchmarks if cfg]
         return {
@@ -47,8 +47,8 @@ def _build_payload(self, args, rollout_id: int) -> Dict[str, Any]:
             "benchmarks": benchmarks,
         }
 
-    def _request(self, payload: Dict[str, Any]) -> Dict[str, Any]:
-        last_error: Optional[Exception] = None
+    def _request(self, payload: dict[str, Any]) -> dict[str, Any]:
+        last_error: Exception | None = None
         for attempt in range(1, self._max_retries + 1):
             try:
                 response = self._session.post(

examples/eval/nemo_skills/skills_config.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
+from collections.abc import Mapping
 from dataclasses import dataclass, field
-from typing import Any, List, Mapping
+from typing import Any
 
 from examples.eval.eval_delegate import EvalEnvConfig, EvalEnvDatasetConfig
 
@@ -35,10 +36,10 @@ def parse(cls, args, dataset_cfg: Mapping[str, Any], defaults: Mapping[str, Any]
 class SkillsEvalEnvConfig(EvalEnvConfig):
     """Environment configuration shared by the Skills client/server."""
 
-    datasets: List[SkillsEvalEnvDatasetConfig] = field(default_factory=list)
+    datasets: list[SkillsEvalEnvDatasetConfig] = field(default_factory=list)
 
     @classmethod
-    def parse(cls, args, raw_env_config: Mapping[str, Any], defaults: Mapping[str, Any]) -> "SkillsEvalEnvConfig":
+    def parse(cls, args, raw_env_config: Mapping[str, Any], defaults: Mapping[str, Any]) -> SkillsEvalEnvConfig:
         base_cfg: SkillsEvalEnvConfig = super().parse(raw_env_config, defaults)
         datasets = raw_env_config.get("datasets") or []
         base_cfg.datasets = [

examples/eval/nemo_skills/skills_server.py

@@ -30,9 +30,10 @@
 import threading
 import time
 import uuid
+from collections.abc import Mapping
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Dict, List, Mapping
+from typing import Any
 
 REPO_ROOT = Path(__file__).resolve().parents[3]
 if str(REPO_ROOT) not in sys.path:
@@ -56,8 +57,8 @@
 class EvalRequestPayload:
     rollout_id: int
     router_url: str
-    defaults: Dict[str, Any] = field(default_factory=dict)
-    benchmarks: List[SkillsEvalEnvDatasetConfig] = field(default_factory=list)
+    defaults: dict[str, Any] = field(default_factory=dict)
+    benchmarks: list[SkillsEvalEnvDatasetConfig] = field(default_factory=list)
 
 
 # ---------------------------------------------------------------------------
@@ -83,8 +84,8 @@ def _hydra_overrides_from_benchmark(
     router_url: str,
     openai_model_name: str,
     max_concurrent_requests: int,
-) -> List[str]:
-    overrides: List[str] = []
+) -> list[str]:
+    overrides: list[str] = []
     for key, hydra_key in HYDRA_OVERRIDE_MAP.items():
         value = getattr(benchmark_cfg, key, None)
         if value is None:
@@ -114,7 +115,7 @@ class ServerConfig:
     max_concurrent_requests: int = 512
 
     @classmethod
-    def from_args(cls, args: argparse.Namespace) -> "ServerConfig":
+    def from_args(cls, args: argparse.Namespace) -> ServerConfig:
         return cls(
             output_root=Path(args.output_root).expanduser().resolve(),
             cluster=args.cluster,
@@ -130,7 +131,7 @@ def __init__(self, config: ServerConfig):
         self._lock = threading.Lock()
         self._config.output_root.mkdir(parents=True, exist_ok=True)
 
-    def evaluate(self, payload: EvalRequestPayload) -> Dict[str, Any]:
+    def evaluate(self, payload: EvalRequestPayload) -> dict[str, Any]:
         if not payload.benchmarks:
             warning_msg = "No benchmarks specified in delegate config; skipping NeMo Skills evaluation."
             logger.warning(warning_msg)
@@ -149,8 +150,8 @@ def evaluate(self, payload: EvalRequestPayload) -> Dict[str, Any]:
         run_dir = self._config.output_root / f"{int(time.time())}-{exp_name}"
         run_dir.mkdir(parents=True, exist_ok=True)
 
-        runs: List[Dict[str, Any]] = []
-        raw_metrics: Dict[str, Any] = {}
+        runs: list[dict[str, Any]] = []
+        raw_metrics: dict[str, Any] = {}
         with self._lock:
             for benchmark in payload.benchmarks:
                 result = self._run_single_benchmark(
@@ -182,7 +183,7 @@ def _run_single_benchmark(
         exp_name: str,
         router_url: str,
         run_dir: Path,
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         name = benchmark.name
         benchmark_run_dir = run_dir / name
         benchmark_run_dir.mkdir(parents=True, exist_ok=True)
@@ -220,7 +221,7 @@ def _build_command(
         run_dir: Path,
         defaults: Mapping[str, Any],
         benchmark_cfg: SkillsEvalEnvDatasetConfig,
-    ) -> List[str]:
+    ) -> list[str]:
         base_cmd = [
             "ns",
             "eval",
@@ -250,29 +251,29 @@ def _build_command(
         )
         return base_cmd + hydra_overrides
 
-    def _build_env(self) -> Dict[str, str]:
+    def _build_env(self) -> dict[str, str]:
         env = os.environ.copy()
         return env
 
     @staticmethod
-    def _run_command(cmd: List[str], *, env: Dict[str, str], log_path: Path):
+    def _run_command(cmd: list[str], *, env: dict[str, str], log_path: Path):
         with open(log_path, "w", encoding="utf-8") as log_file:
             process = subprocess.Popen(cmd, stdout=log_file, stderr=subprocess.STDOUT, env=env)
             retcode = process.wait()
         if retcode != 0:
-            with open(log_path, "r", encoding="utf-8", errors="ignore") as log_file:
+            with open(log_path, encoding="utf-8", errors="ignore") as log_file:
                 tail = "".join(log_file.readlines()[-200:])
             raise RuntimeError(f"`ns eval` failed with exit code {retcode}. See {log_path}\n{tail}")
 
     @staticmethod
-    def _collect_metrics(run_dir: Path, benchmark: str) -> Dict[str, Any]:
+    def _collect_metrics(run_dir: Path, benchmark: str) -> dict[str, Any]:
         benchmark_name = benchmark.split(":")[0]
         metrics_path = run_dir / "eval-results" / benchmark_name / "metrics.json"
         if not metrics_path.exists():
             logger.warning("Metrics file missing for %s at %s", benchmark_name, metrics_path)
             return {}
         try:
-            with open(metrics_path, "r", encoding="utf-8") as fp:
+            with open(metrics_path, encoding="utf-8") as fp:
                 metrics_data = json.load(fp)
         except json.JSONDecodeError as exc:
             logger.warning("Failed to parse %s: %s", metrics_path, exc)

examples/formal_math/single_round/kimina_wrapper.py

@@ -2,7 +2,6 @@
 import os
 import random
 import time
-from typing import List
 
 import ray
 import requests
@@ -25,7 +24,7 @@ async def check(self, *args, **kwargs) -> CheckResponse:
 
 
 class _KiminaClientCluster:
-    def __init__(self, servers: List["_KiminaServerActor"]):
+    def __init__(self, servers: list["_KiminaServerActor"]):
         self._clients = [AsyncKiminaClient(api_url=ray.get(server.get_api_url.remote())) for server in servers]
         self._next_client_index = 0
 
@@ -35,7 +34,7 @@ async def check(self, *args, **kwargs):
         return await client.check(*args, **kwargs)
 
 
-def _create_actor_per_node(actor_cls) -> List:
+def _create_actor_per_node(actor_cls) -> list:
     # for simplicity, we use all available nodes
     nodes = [n for n in ray.nodes() if n.get("Alive")]
     assert len(nodes) > 0