[feature] support local scheduler

areal/api/scheduler_api.py

@@ -36,10 +36,11 @@ class SchedulingConfig:
 
 
 class Scheduler(abc.ABC):
-    def create_workers(self, worker_key, scheduler_config, *args, **kwargs) -> str:
+    def create_workers(self, worker_key, scheduler_config, *args, **kwargs) -> None:
         """
         Start workers, return job id
         """
+        raise NotImplementedError()
 
     def get_workers(self, worker_key, timeout=None) -> List[Worker]:
         """

areal/api/workflow_api.py

@@ -330,14 +330,11 @@ async def _rollout_thread_async(self):
         try:
             while not self.exiting.is_set():
                 # Check capacity
-                capacity = self.get_capacity()
+                # capacity = self.get_capacity()
+                # self.logger.info(f"Current rollout capacity: {capacity}")
                 # Create new rollout task
                 self.lock.acquire()
-                while (
-                    capacity > 0
-                    and not self.paused.is_set()
-                    and self.input_queue.qsize() > 0
-                ):
+                while not self.paused.is_set() and self.input_queue.qsize() > 0:
                     x = self.input_queue.get_nowait()
                     x: _RolloutTaskInput
                     self.logger.debug(f"Get data from puller: {x.data}")
@@ -357,7 +354,7 @@ async def _rollout_thread_async(self):
                             f"running: {self.rollout_stat.running}, "
                             f"accepted: {self.rollout_stat.accepted}."
                         )
-                    capacity -= 1
+                    # capacity -= 1
                     rid += 1
                 tasks = [x.task for x in rollout_tasks.values()]
                 self.lock.release()

areal/engine/base_hf_engine.py

@@ -73,7 +73,7 @@ def __init__(self, config: TrainEngineConfig):
         )
         self.is_vision_model = is_valid_vision_model(self.model_config.model_type)
 
-        self.world_size = int(os.environ["WORLD_SIZE"])
+        self.world_size: int
 
     def set_version(self, version: int):
         self._version = version

areal/engine/fsdp_engine.py

@@ -121,9 +121,17 @@ def create_process_group(self, parallel_strategy: ParallelStrategy | None = None
         self.dp_head = int(self.world_mesh["sp_tp"].mesh[0].item())
         self.dp_rank = dist.get_rank(self.dp_group)
 
+        self.world_size = int(os.environ["WORLD_SIZE"])
+
         self.logger.info(f"Data parallel head {self.dp_head} and rank {self.dp_rank}")
 
-    def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None):
+    def initialize(
+        self,
+        addr: str | None,
+        ft_spec: FinetuneSpec | None,
+        parallel_strategy: ParallelStrategy | None = None,
+    ):
+        self.create_process_group(parallel_strategy)
         # Initialize distributed enviroments and load model.
         assert addr is None, "FSDPEngine does not support remote initialization."
         assert ft_spec is not None, "FSDPEngine requires FinetuneSpec to initialize."

areal/engine/ppo/actor.py

@@ -67,7 +67,7 @@ def calc_logprobs(logits, input_data):
             aggregate_fn=lambda xs: torch.cat(xs, dim=-1),
         )
 
-    def compute_advantages(self, data: Dict[str, Any]) -> None:
+    def compute_advantages(self, data: Dict[str, Any]) -> Dict[str, Any]:
         bs = data["input_ids"].shape[0]
         max_seqlen = data["input_ids"].shape[1]
         batch_indices = torch.arange(
@@ -162,6 +162,8 @@ def compute_advantages(self, data: Dict[str, Any]) -> None:
         # because we have rolled old_logp by -1
         data["logprobs"] = old_logp
 
+        return data
+
     def ppo_update(self, data: Dict[str, Any]) -> List[Dict[str, float]]:
 
         if self.dynamic_sampling and len(data["rewards"]) % self.group_size == 0:
@@ -286,8 +288,8 @@ def compute_logp(self, *args, **kwargs) -> torch.Tensor | None:
         return self.actor.compute_logp(*args, **kwargs)
 
     @torch.no_grad()
-    def compute_advantages(self, *args, **kwargs) -> None:
-        self.actor.compute_advantages(*args, **kwargs)
+    def compute_advantages(self, *args, **kwargs):
+        return self.actor.compute_advantages(*args, **kwargs)
 
     def ppo_update(self, *args, **kwargs) -> List[Dict[str, float]]:
         return self.actor.ppo_update(*args, **kwargs)

areal/engine/sglang_remote.py

@@ -45,11 +45,8 @@ def __init__(self, config: InferenceEngineConfig):
         self.distributed_weight_update_initialized = False
         self._version = 0
 
-        self.lock = Lock()
-        self.workflow_executor = WorkflowExecutor(
-            config=config,
-            inference_engine=self,
-        )
+        self.lock: Lock
+        self.workflow_executor: WorkflowExecutor
 
     def _wait_for_server(self, address):
         base_url = f"http://{address}"
@@ -74,6 +71,11 @@ def initialize(
         addr: str | List[str] | None = None,
         train_data_parallel_size: int | None = None,
     ):
+        self.lock = Lock()
+        self.workflow_executor = WorkflowExecutor(
+            config=self.config,
+            inference_engine=self,
+        )
         if engine_id is None:
             if dist.is_initialized():
                 engine_id = str(dist.get_rank())

areal/reward/gsm8k_reward.py

@@ -0,0 +1,5 @@
+from areal.reward.math_parser import process_results
+
+
+def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs):
+    return int(process_results(completions, answer)[0])