chore: format areal/api with ruff (#491)

* format areal/api with ruff

* format

* .

Author: garrett4wade

areal/api/alloc_mode.py

@@ -9,7 +9,6 @@
 import enum
 import math
 from dataclasses import dataclass, field
-from typing import Optional
 
 from lark import Lark, Transformer
 
@@ -216,7 +215,7 @@ def parallelism_eq(this, other):
 class MegatronParallelStrategy(ParallelStrategy):
     """Megatron parallel strategy with additional sequence parallelism and virtual pipeline parallelism."""
 
-    virtual_pipeline_parallel_size: Optional[int] = field(
+    virtual_pipeline_parallel_size: int | None = field(
         default=None,
         metadata={
             "help": "Virtual pipeline parallelism size for megatron modules "
@@ -234,10 +233,7 @@ class MegatronParallelStrategy(ParallelStrategy):
     def parallelism_eq(this, other):
         """Compare Megatron parallelism configurations (excluding sequence parallelism)."""
         return ParallelStrategy.parallelism_eq(this, other) and (
-            (
-                this.virtual_pipeline_parallel_size
-                == other.virtual_pipeline_parallel_size
-            )
+            this.virtual_pipeline_parallel_size == other.virtual_pipeline_parallel_size
         )
 
 
@@ -274,9 +270,9 @@ class AllocationMode:
 
     type_: AllocationType
     gen: ParallelStrategy = field(default_factory=ParallelStrategy)
-    train: Optional[ParallelStrategy] = None
-    gen_backend: Optional[str] = None
-    train_backend: Optional[str] = None
+    train: ParallelStrategy | None = None
+    gen_backend: str | None = None
+    train_backend: str | None = None
 
     @property
     def gen_instance_size(self) -> int:
@@ -407,7 +403,7 @@ class TrainingParallelism:
     and comprehensive validation rules.
     """
 
-    backend: Optional[str] = None
+    backend: str | None = None
     strategy: ParallelStrategy = field(default_factory=lambda: ParallelStrategy())
 
     def __post_init__(self):

areal/api/cli_args.py

@@ -3,21 +3,19 @@
 import os
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
-from typing import Dict, List
 
 import uvloop
 import yaml
-
-from areal.utils.pkg_version import is_version_less
-
-uvloop.install()
 from hydra import compose as hydra_compose
 from hydra import initialize as hydra_init
 from hydra.core.global_hydra import GlobalHydra
 from omegaconf import MISSING, DictConfig, OmegaConf
 
 from areal.platforms import current_platform
 from areal.utils import name_resolve, pkg_version
+from areal.utils.pkg_version import is_version_less
+
+uvloop.install()
 
 
 @dataclass
@@ -129,11 +127,11 @@ class GenerationHyperparameters:
         default=1.0,
         metadata={"help": "Sampling temperature. Higher values increase diversity."},
     )
-    stop_token_ids: List[int] = field(
+    stop_token_ids: list[int] = field(
         default_factory=list,
         metadata={"help": "Stop generation when encountering these token IDs."},
     )
-    stop: List[str] | None = field(
+    stop: list[str] | None = field(
         default=None,
         metadata={
             "help": "One or multiple stop words. Generation will stop if one of these words is sampled."
@@ -232,7 +230,7 @@ class OptimizerConfig:
 class FSDPWrapPolicy:
     """Policy configuration for FSDP model layer wrapping. None defaults to wrapping transformer decoder layers defined by transformers."""
 
-    transformer_layer_cls_to_wrap: List[str] | None = field(
+    transformer_layer_cls_to_wrap: list[str] | None = field(
         default=None,
         metadata={"help": "A list of transformer layer names for FSDP to wrap."},
     )
@@ -310,7 +308,7 @@ class MegatronEngineConfig:
     recompute_method: str | None = "uniform"
     recompute_num_layers: int | None = 1
     distribute_saved_activations: bool | None = None
-    recompute_modules: List[str] | None = None
+    recompute_modules: list[str] | None = None
 
 
 @dataclass
@@ -378,7 +376,7 @@ class TrainEngineConfig:
     )
     lora_rank: int = field(default=32, metadata={"help": "lora rank"})
     lora_alpha: int = field(default=16, metadata={"help": "lora alpha"})
-    target_modules: List[str] = field(
+    target_modules: list[str] = field(
         default_factory=list,
         metadata={"help": "lora target_modules."},
     )
@@ -500,7 +498,7 @@ class PPOActorConfig(TrainEngineConfig):
         default=False,
         metadata={"help": "Log statistics for agent trajectories"},
     )
-    log_agent_stats_keys: List[str] = field(
+    log_agent_stats_keys: list[str] = field(
         default_factory=lambda: [],
         metadata={"help": "Keys for logging agent trajectory statistics"},
     )
@@ -574,7 +572,7 @@ def build_args(
         port,
         dist_init_addr: str | None = None,
     ):
-        args: Dict = conf_as_dict(vllm_config)
+        args: dict = conf_as_dict(vllm_config)
         args = dict(
             host=host,
             port=port,
@@ -608,11 +606,11 @@ def build_cmd(
             if v is None or v is False or v == "":
                 continue
             if v is True:
-                flags.append(f"--{k.replace('_','-')}")
+                flags.append(f"--{k.replace('_', '-')}")
             elif isinstance(v, list):
-                flags.append(f"--{k.replace('_','-')} {' '.join(map(str, v))}")
+                flags.append(f"--{k.replace('_', '-')} {' '.join(map(str, v))}")
             else:
-                flags.append(f"--{k.replace('_','-')} {v}")
+                flags.append(f"--{k.replace('_', '-')} {v}")
         return f"python3 -m areal.thirdparty.vllm.areal_vllm_server {' '.join(flags)}"
 
 
@@ -638,7 +636,7 @@ class SGLangConfig:
     enable_torch_compile: bool = False
     torch_compile_max_bs: int = 32
     cuda_graph_max_bs: int | None = None
-    cuda_graph_bs: List[int] | None = None
+    cuda_graph_bs: list[int] | None = None
     torchao_config: str = ""
     enable_nan_detection: bool = False
     enable_p2p_check: bool = False
@@ -667,8 +665,8 @@ class SGLangConfig:
     # lora
     enable_lora: bool | None = None
     max_lora_rank: int | None = None
-    lora_target_modules: List[str] | None = None
-    lora_paths: List[str] | None = None
+    lora_target_modules: list[str] | None = None
+    lora_paths: list[str] | None = None
     max_loaded_loras: int = 1
     max_loras_per_batch: int = 1
     lora_backend: str = "triton"
@@ -719,11 +717,11 @@ def build_cmd(
             if v is None or v is False or v == "":
                 continue
             if v is True:
-                flags.append(f"--{k.replace('_','-')}")
+                flags.append(f"--{k.replace('_', '-')}")
             elif isinstance(v, list):
-                flags.append(f"--{k.replace('_','-')} {' '.join(map(str, v))}")
+                flags.append(f"--{k.replace('_', '-')} {' '.join(map(str, v))}")
             else:
-                flags.append(f"--{k.replace('_','-')} {v}")
+                flags.append(f"--{k.replace('_', '-')} {v}")
         return f"python3 -m sglang.launch_server {' '.join(flags)}"
 
     @staticmethod
@@ -738,11 +736,12 @@ def build_args(
         node_rank: int = 0,
     ):
         # Map "all-linear" to "all"
-        args: Dict = conf_as_dict(sglang_config)
+        args: dict = conf_as_dict(sglang_config)
         if sglang_config.enable_multithread_load or sglang_config.enable_fast_load:
-            assert pkg_version.is_version_equal(
-                "sglang", "0.5.2"
-            ), f"Customized model loading requires exact SGLang version 0.5.2"
+            if not pkg_version.is_version_equal("sglang", "0.5.2"):
+                raise RuntimeError(
+                    "Customized model loading requires exact SGLang version 0.5.2"
+                )
             model_loader_extra_config = dict(
                 enable_multithread_load=sglang_config.enable_multithread_load,
                 enable_fast_load=sglang_config.enable_fast_load,
@@ -791,7 +790,8 @@ class InferenceEngineConfig:
     max_concurrent_rollouts: None | int = field(
         default=None,
         metadata={
-            "help": "Maximum number of concurrent rollouts to the inference engine. Defaults to consumer_batch_size."
+            "help": "Maximum number of concurrent rollouts to "
+            "the inference engine. Defaults to consumer_batch_size."
         },
     )
     queue_size: None | int = field(
@@ -915,8 +915,8 @@ class WandBConfig:
     job_type: str | None = None
     group: str | None = None
     notes: str | None = None
-    tags: List[str] | None = None
-    config: Dict | None = None
+    tags: list[str] | None = None
+    config: dict | None = None
     id_suffix: str | None = "train"
 
 
@@ -926,7 +926,7 @@ class SwanlabConfig:
 
     project: str | None = None
     name: str | None = None
-    config: Dict | None = None
+    config: dict | None = None
     logdir: str | None = None
     mode: str | None = "disabled"
     api_key: str | None = os.getenv("SWANLAB_API_KEY", None)
@@ -1023,7 +1023,7 @@ class SchedulerConfig:
     endpoint: str = field(default="http://localhost:8081")
     deploy_mode: str = field(default="separation")
     functioncall_service_domain: str = field(default="http://localhost:8080")
-    reward_functioncall_config: Dict = field(default_factory=dict)
+    reward_functioncall_config: dict = field(default_factory=dict)
     reward_model_path: str = field(default="")
     reward_model_service_url: str = field(default="http://localhost:30000/classify")
 
@@ -1076,7 +1076,7 @@ class SlurmLauncherConfig:
         default="--mpi=pmi2 -K --chdir $PWD",
         metadata={"help": "Additional arguments to pass to the srun command."},
     )
-    additional_bash_cmds: List[str] | None = field(
+    additional_bash_cmds: list[str] | None = field(
         default=None,
         metadata={
             "help": "Additional bash commands to setup the container before running "
@@ -1244,7 +1244,7 @@ class PPOConfig(GRPOConfig):
     critic: PPOCriticConfig = field(default_factory=PPOCriticConfig)
 
 
-def parse_cli_args(argv: List[str]):
+def parse_cli_args(argv: list[str]):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--config", help="Path to the main configuration file", required=True
@@ -1277,7 +1277,7 @@ def to_structured_cfg(cfg, config_cls):
     return cfg
 
 
-def load_expr_config(argv: List[str], config_cls):
+def load_expr_config(argv: list[str], config_cls):
     cfg, config_file = parse_cli_args(argv)
     cfg = to_structured_cfg(cfg, config_cls=config_cls)
     cfg = OmegaConf.to_object(cfg)
@@ -1305,7 +1305,7 @@ def save_config(cfg, log_dir):
     os.makedirs(log_dir, exist_ok=True)
     config_save_path = os.path.join(log_dir, "config.yaml")
     with open(config_save_path, "w") as f:
-        config_dict: Dict = asdict(cfg)
+        config_dict: dict = asdict(cfg)
         yaml.dump(
             config_dict,
             f,