Ray Scheduler Implementation for Single Controller

areal/controller/train_controller.py

@@ -4,7 +4,6 @@
 from pathlib import Path
 from typing import Any
 
-import aiohttp
 import torch.distributed as dist
 from torchdata.stateful_dataloader import StatefulDataLoader
 
@@ -579,25 +578,14 @@ def update_weights(self, meta: WeightUpdateMeta):
             raise ValueError(f"Unknown weight update type {meta.type}")
 
     async def _async_clear_batches(self, *targets: dict[str, RTensor]):
-        """Extract shard IDs and call /data/clear on each worker."""
+        """Extract shard IDs and clear tensors on each worker."""
         shards_by_node = RTensor.collect_shards(targets)
 
         if not shards_by_node:
             return
 
-        async def clear_node(node_addr, shard_ids):
-            async with aiohttp.ClientSession() as session:
-                async with session.delete(
-                    f"http://{node_addr}/data/clear", json={"shard_ids": shard_ids}
-                ) as resp:
-                    if resp.status == 200:
-                        result = await resp.json()
-                        logger.info(
-                            f"Cleared {result.get('cleared_count', 0)} shards on {node_addr}"
-                        )
-
         await asyncio.gather(
-            *[clear_node(addr, sids) for addr, sids in shards_by_node.items()],
+            *[RTensor.clear_node(addr, sids) for addr, sids in shards_by_node.items()],
             return_exceptions=True,
         )
 

areal/core/remote_inf_engine.py

@@ -14,6 +14,7 @@
 from typing import Any, Protocol
 
 import aiohttp
+import ray
 import requests
 import torch.distributed as dist
 import uvloop
@@ -949,6 +950,13 @@ def launch_server(self, server_args: dict[str, Any]) -> LocalInfServerInfo:
         try:
             self._wait_for_server(address)
             self.local_server_processes.append(server_info)
+            if ray.is_initialized():
+                # do not return with process for ray as it is not picklable
+                return LocalInfServerInfo(
+                    host=server_args["host"],
+                    port=server_args["port"],
+                    process=None,
+                )
             return server_info
         except TimeoutError:
             logger.warning(

areal/experimental/trainer/rl.py

@@ -36,7 +36,7 @@
 from areal.engine.sglang_remote import RemoteSGLangEngine
 from areal.engine.vllm_remote import RemotevLLMEngine
 from areal.platforms import current_platform
-from areal.scheduler import LocalScheduler
+from areal.scheduler import LocalScheduler, RayScheduler
 from areal.utils import logging, perf_tracer, seeding, stats_tracker
 from areal.utils.dataloader import create_dataloader
 from areal.utils.environ import is_single_controller
@@ -395,6 +395,8 @@ def _init_scheduler(self) -> Scheduler:
         cfg = self.config.scheduler
         if cfg.type == "local":
             return LocalScheduler(exp_config=self.config)
+        elif cfg.type == "ray":
+            return RayScheduler(exp_config=self.config)
         raise NotImplementedError(f"Unknown scheduler type: {cfg.type}")
 
     def _create_dataloader(

areal/scheduler/__init__.py

@@ -1,5 +1,4 @@
 from .local import LocalScheduler
+from .ray import RayScheduler
 
-__all__ = [
-    "LocalScheduler",
-]
+__all__ = ["LocalScheduler", "RayScheduler"]