[ci] add the recipe/ directory to pre-commit hooks (#580)

Author: fishcrap

.pre-commit-config.yaml

@@ -43,11 +43,11 @@ repos:
         name: Run Linter Check (Ruff)
         types_or: [ python, pyi, jupyter ]
         args: [ --fix ]
-        files: ^(areal|examples|docs)/
+        files: ^(areal|examples|docs|recipe)/
       - id: ruff-format # Run the formatter.
         name: Run Formatter (Ruff)
         types_or: [ python, pyi, jupyter ]
-        files: ^(areal|examples|docs)/
+        files: ^(areal|examples|docs|recipe)/
 
   # Clean notebook outputs and metadata
   - repo: https://github.com/kynan/nbstripout

recipe/AEnt/actor.py

@@ -1,28 +1,25 @@
 import functools
-from typing import Dict, List, Any
+from typing import Any
 
 import torch
 
+from recipe.AEnt.aent_args import AEntPPOActorConfig
+from recipe.AEnt.functional import gather_logprobs_clamped_entropy
 
-from areal.api.cli_args import MicroBatchSpec, PPOActorConfig
+from areal.api.cli_args import MicroBatchSpec
 from areal.api.engine_api import TrainEngine
 from areal.engine.fsdp_engine import FSDPEngine
 from areal.engine.ppo.actor import PPOActor
 from areal.utils import stats_tracker
 from areal.utils.data import split_padded_tensor_dict_into_mb_list
 from areal.utils.functional import (
     dynamic_sampling,
-    gather_logprobs,
     gather_logprobs_entropy,
     ppo_actor_loss_fn,
-    reward_overlong_penalty,
 )
-from recipe.AEnt.aent_args import AEntPPOActorConfig
-from recipe.AEnt.functional import gather_logprobs_clamped_entropy
 
 
 class AEntPPOActor(PPOActor):
-
     def __init__(self, config: AEntPPOActorConfig, engine: TrainEngine):
         super().__init__(config, engine)
         self.entropy_coeff = config.aent.entropy_coeff
@@ -39,7 +36,7 @@ def __init__(self, config: AEntPPOActorConfig, engine: TrainEngine):
     @stats_tracker.scope_func_wrapper("aent_ppo_actor")
     def aent_ppo_update(
         self, data: dict[str, Any], global_step: int
-    ) -> List[Dict[str, float]]:
+    ) -> list[dict[str, float]]:
         with stats_tracker.scope("dynamic_sampling"):
             if self.dynamic_sampling and len(data["rewards"]) % self.group_size == 0:
                 data, sampling_stat = dynamic_sampling(data, self.group_size)
@@ -156,7 +153,6 @@ def aent_ppo_update(
 
 
 class FSDPAEntPPOActor(FSDPEngine):
-
     def __init__(self, config: AEntPPOActorConfig):
         super().__init__(config)
         self.actor = AEntPPOActor(config, self)
@@ -169,14 +165,14 @@ def compute_logp(self, *args, **kwargs) -> torch.Tensor | None:
     def compute_advantages(self, *args, **kwargs) -> None:
         self.actor.compute_advantages(*args, **kwargs)
 
-    def aent_ppo_update(self, *args, **kwargs) -> List[Dict[str, float]]:
+    def aent_ppo_update(self, *args, **kwargs) -> list[dict[str, float]]:
         return self.actor.aent_ppo_update(*args, **kwargs)
 
 
 # AEnt regularized grpo loss
 def aent_grpo_loss_fn(
     logits: torch.Tensor,
-    input_data: Dict,
+    input_data: dict,
     temperature: float,
     eps_clip: float,
     entropy_coeff: float,

recipe/AEnt/gsm8k_aent_grpo.py

@@ -5,8 +5,11 @@
 import torch.distributed as dist
 from torchdata.stateful_dataloader import StatefulDataLoader
 
+from recipe.AEnt.actor import FSDPAEntPPOActor
+from recipe.AEnt.aent_args import AEntGRPOConfig
+
 from areal.api.alloc_mode import AllocationMode
-from areal.api.cli_args import GRPOConfig, load_expr_config
+from areal.api.cli_args import load_expr_config
 from areal.api.io_struct import FinetuneSpec, StepInfo, WeightUpdateMeta
 from areal.dataset import get_custom_dataset
 from areal.engine.ppo.actor import FSDPPPOActor
@@ -15,7 +18,6 @@
 from areal.utils import seeding, stats_tracker
 from areal.utils.data import (
     broadcast_tensor_container,
-    cycle_dataloader,
     tensor_container_to,
 )
 from areal.utils.device import log_gpu_stats
@@ -25,8 +27,6 @@
 from areal.utils.saver import Saver
 from areal.utils.stats_logger import StatsLogger
 from areal.workflow.rlvr import RLVRWorkflow
-from recipe.AEnt.actor import FSDPAEntPPOActor
-from recipe.AEnt.aent_args import AEntGRPOConfig
 
 
 def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs):