feat(launcher): auto-set OMP_NUM_THREADS based on cpus_per_task

areal/api/cli_args.py

@@ -665,7 +665,7 @@ class SchedulingStrategy:
 @dataclass
 class SchedulingSpec:
     cpu: int = field(
-        default=4, metadata={"help": "Number of CPU cores required per GPU"}
+        default=8, metadata={"help": "Number of CPU cores required per GPU"}
     )
     gpu: int = field(
         default=0,

areal/launcher/local.py

@@ -26,6 +26,8 @@
     JobException,
     JobInfo,
     JobState,
+    get_scheduling_spec,
+    get_thread_env_vars,
     validate_config_for_launcher,
     wait_llm_server_addrs,
 )
@@ -325,19 +327,21 @@ def local_main(config, run_id: int = 0):
         )
 
         # Launch inference servers.
-        try:
-            rollout_env_vars = config.rollout.scheduling_spec[0].env_vars
-        except AttributeError:
-            # In case `scheduling_spec` or `env_vars` is missing
-            rollout_env_vars = {}
+        rollout_spec = get_scheduling_spec(config.rollout)
+        rollout_env_vars = rollout_spec.env_vars
+        rollout_cpus_per_task = rollout_spec.cpu
+        thread_env = get_thread_env_vars(
+            cpus_per_task=rollout_cpus_per_task,
+            existing_env_vars=rollout_env_vars,
+        )
         launcher.submit_array(
             job_name="llm_server",
             cmd=server_cmd,
             count=1,
             gpu=alloc_mode.gen.pp_size
             * alloc_mode.gen.tp_size
             * alloc_mode.gen.dp_size,
-            env_vars={**BASE_ENVIRONS, **rollout_env_vars},
+            env_vars={**BASE_ENVIRONS, **thread_env, **rollout_env_vars},
         )
 
         # Get llm server addresses by name resolve
@@ -374,17 +378,20 @@ def local_main(config, run_id: int = 0):
             _env_vars["NCCL_CUMEM_ENABLE"] = "0"
             _env_vars["NCCL_NVLS_ENABLE"] = "0"
         # All experiment configs should have the `actor` field.
-        try:
-            actor_env_vars = config.actor.scheduling_spec[0].env_vars
-        except AttributeError:
-            # in case `scheduling_spec` or `env_vars` is missing
-            actor_env_vars = {}
+        actor_spec = get_scheduling_spec(config.actor)
+        actor_env_vars = actor_spec.env_vars
+        actor_cpus_per_task = actor_spec.cpu
+        thread_env = get_thread_env_vars(
+            cpus_per_task=actor_cpus_per_task,
+            existing_env_vars=actor_env_vars,
+        )
         launcher.submit(
             job_name="trainer",
             cmd=f"torchrun --nnodes 1 --nproc-per-node {nprocs} --master-addr localhost --master-port {find_free_ports(1, (10000, 50000))[0]} {' '.join(sys.argv[1:])}",
             gpu=gpu,
             env_vars={
                 **BASE_ENVIRONS,
+                **thread_env,
                 **actor_env_vars,
                 **_env_vars,
                 **tms_env_vars,

areal/launcher/ray.py

@@ -17,7 +17,6 @@
 from areal.api.cli_args import (
     ClusterSpecConfig,
     RecoverConfig,
-    SchedulingSpec,
     SGLangConfig,
     parse_cli_args,
     to_structured_cfg,
@@ -30,6 +29,8 @@
     BASE_ENVIRONS,
     JobException,
     JobState,
+    get_scheduling_spec,
+    get_thread_env_vars,
     validate_config_for_distributed_launcher,
     wait_llm_server_addrs,
 )
@@ -372,18 +373,10 @@ def ray_main(config, run_id: int = 0):
     allocation_mode = config.allocation_mode
     allocation_mode = AllocationMode.from_str(allocation_mode)
 
-    try:
-        actor_spec = to_structured_cfg(config.actor.scheduling_spec[0], SchedulingSpec)
-    except AttributeError:
-        actor_spec = SchedulingSpec()
+    actor_spec = get_scheduling_spec(config.actor)
 
     if allocation_mode.gen_backend in ("sglang", "vllm"):
-        try:
-            rollout_spec = to_structured_cfg(
-                config.rollout.scheduling_spec[0], SchedulingSpec
-            )
-        except AttributeError:
-            rollout_spec = SchedulingSpec()
+        rollout_spec = get_scheduling_spec(config.rollout)
 
     if not is_recover_run:
         metadata_file = save_experiment_metadata(
@@ -446,7 +439,14 @@ def sglang_env_hook(
 
             return env_vars
 
-        # launch a task to start all sglang servers in one node
+        # Launch a task to start all sglang servers in one node.
+        # Use full-node CPU allocation for llm_server since one Ray task manages
+        # all GPUs on a node. Thread env vars are set per-node (not per-GPU).
+        sglang_cpus_per_task = rollout_spec.cpu * n_gpus_per_node
+        thread_env = get_thread_env_vars(
+            cpus_per_task=sglang_cpus_per_task,
+            existing_env_vars=rollout_spec.env_vars,
+        )
         launcher.submit_array(
             job_name="llm_server",
             file_path=sglang_entry_point,
@@ -455,9 +455,9 @@ def sglang_env_hook(
             nodes=n_sglang_nodes,
             list_args=sglang_args_list,
             gpus_per_task=n_gpus_per_node,
-            cpus_per_task=rollout_spec.cpu * n_gpus_per_node,
+            cpus_per_task=sglang_cpus_per_task,
             mem_per_task=rollout_spec.mem * 1024 * n_gpus_per_node,
-            env_vars={**BASE_ENVIRONS, **rollout_spec.env_vars},
+            env_vars={**BASE_ENVIRONS, **thread_env, **rollout_spec.env_vars},
             env_hook=(
                 partial(sglang_env_hook, n_sglang_nodes, node_group_size)
                 if cross_nodes
@@ -491,6 +491,11 @@ def sglang_env_hook(
         vllm_entry_point = str(
             pathlib.Path(__file__).resolve().parent.joinpath("vllm_server.py")
         )
+        vllm_cpus_per_task = rollout_spec.cpu * vllm_tp_size
+        thread_env = get_thread_env_vars(
+            cpus_per_task=vllm_cpus_per_task,
+            existing_env_vars=rollout_spec.env_vars,
+        )
         launcher.submit_array(
             job_name="llm_server",
             file_path=vllm_entry_point,
@@ -499,9 +504,9 @@ def sglang_env_hook(
             nodes=n_vllm_nodes,
             list_args=vllm_args_list,
             gpus_per_task=vllm_tp_size,
-            cpus_per_task=rollout_spec.cpu * vllm_tp_size,
+            cpus_per_task=vllm_cpus_per_task,
             mem_per_task=rollout_spec.mem * 1024 * vllm_tp_size,
-            env_vars={**BASE_ENVIRONS, **rollout_spec.env_vars},
+            env_vars={**BASE_ENVIRONS, **thread_env, **rollout_spec.env_vars},
         )
         # Get vllm server addresses via name_resolve
         try:
@@ -570,6 +575,13 @@ def torch_env_hook(n_tasks: int, placement_group: PlacementGroup) -> list[dict]:
             _env_vars["NCCL_CUMEM_ENABLE"] = "0"
             _env_vars["NCCL_NVLS_ENABLE"] = "0"
 
+        # Use per-GPU CPU count for thread env vars since Ray spawns individual
+        # tasks per GPU, each inheriting these env vars. This differs from
+        # llm_server which uses full-node allocation.
+        thread_env = get_thread_env_vars(
+            cpus_per_task=actor_spec.cpu,
+            existing_env_vars=actor_spec.env_vars,
+        )
         launcher.submit_array(
             job_name="trainer",
             file_path=trainer_entry_point,
@@ -582,6 +594,7 @@ def torch_env_hook(n_tasks: int, placement_group: PlacementGroup) -> list[dict]:
             mem_per_task=actor_spec.mem * 1024,
             env_vars={
                 **BASE_ENVIRONS,
+                **thread_env,
                 **actor_spec.env_vars,
                 **_env_vars,
                 **tms_env_vars,

areal/launcher/sglang_server.py

@@ -19,16 +19,26 @@
 )
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names
-from areal.utils.launcher import TRITON_CACHE_PATH, apply_sglang_patch
+from areal.utils.launcher import (
+    TRITON_CACHE_PATH,
+    apply_sglang_patch,
+    get_scheduling_spec,
+)
 from areal.utils.network import find_free_ports, gethostip
 from areal.utils.proc import kill_process_tree
 
 logger = logging.getLogger("SGLangWrapper")
 
 
-def launch_server_cmd(command: list[str]) -> subprocess.Popen:
+def launch_server_cmd(
+    command: list[str], custom_env: dict[str, str] | None = None
+) -> subprocess.Popen:
     """
     Launch inference server in a new process and return its process handle.
+
+    Args:
+        command: The command to execute.
+        custom_env: Custom environment variables to set for the subprocess.
     """
     # Replace newline continuations and split the command string.
     logger.info(f"Launch command: {' '.join(command)}")
@@ -37,6 +47,10 @@ def launch_server_cmd(command: list[str]) -> subprocess.Popen:
     triton_cache_path = _env.get("TRITON_CACHE_PATH", TRITON_CACHE_PATH)
     unique_triton_cache_path = os.path.join(triton_cache_path, str(uuid.uuid4()))
     _env["TRITON_CACHE_PATH"] = unique_triton_cache_path
+
+    if custom_env is not None:
+        _env.update(custom_env)
+
     return subprocess.Popen(
         command,
         env=_env,
@@ -77,13 +91,15 @@ def __init__(
         sglang_config: SGLangConfig,
         allocation_mode: AllocationMode,
         n_gpus_per_node: int,
+        cpu_per_gpu: int | None = None,
     ):
         self.experiment_name = experiment_name
         self.trial_name = trial_name
         self.config = sglang_config
         self.allocation_mode = allocation_mode
         self.server_processes = []  # List to store multiple server processes
         self.n_gpus_per_node = n_gpus_per_node
+        self.cpu_per_gpu = cpu_per_gpu
 
         if self.config.enable_fast_load or self.config.enable_multithread_load:
             apply_sglang_patch()
@@ -208,12 +224,16 @@ def launch_sglang_server(argv):
     allocation_mode = AllocationMode.from_str(allocation_mode)
     assert allocation_mode.gen_backend == "sglang"
 
+    # Get CPU per GPU from rollout scheduling spec
+    rollout_spec = get_scheduling_spec(config.rollout)
+
     sglang_server = SGLangServerWrapper(
         config.experiment_name,
         config.trial_name,
         config.sglang,
         allocation_mode,
         n_gpus_per_node=config.cluster.n_gpus_per_node,
+        cpu_per_gpu=rollout_spec.cpu,
     )
     sglang_server.run()
 

areal/launcher/slurm.py

@@ -11,7 +11,6 @@
 from areal.api.cli_args import (
     ClusterSpecConfig,
     RecoverConfig,
-    SchedulingSpec,
     SGLangConfig,
     parse_cli_args,
     to_structured_cfg,
@@ -25,6 +24,8 @@
     JobException,
     JobInfo,
     JobState,
+    get_scheduling_spec,
+    get_thread_env_vars,
     validate_config_for_distributed_launcher,
     wait_llm_server_addrs,
 )
@@ -414,10 +415,7 @@ def slurm_main(config, run_id: int = 0):
     )
 
     # Get scheduling specs from actor config
-    try:
-        actor_spec = to_structured_cfg(config.actor.scheduling_spec[0], SchedulingSpec)
-    except AttributeError:
-        actor_spec = SchedulingSpec()
+    actor_spec = get_scheduling_spec(config.actor)
 
     launcher = SlurmLauncher(
         experiment_name=config.experiment_name,
@@ -459,12 +457,7 @@ def slurm_main(config, run_id: int = 0):
             random_seed = config.vllm.seed
 
         # Get rollout scheduling spec
-        try:
-            rollout_spec = to_structured_cfg(
-                config.rollout.scheduling_spec[0], SchedulingSpec
-            )
-        except AttributeError:
-            rollout_spec = SchedulingSpec()
+        rollout_spec = get_scheduling_spec(config.rollout)
 
         backend_spec = {
             "sglang": {
@@ -497,7 +490,12 @@ def _build_llm_server_plan(backend: str, spec: dict):
             n_servers_per_node = max(n_backend_servers // n_backend_nodes, 1)
 
             cross_nodes = allocation_mode.gen_instance_size > n_gpus_per_node
-            base_env_vars = {**BASE_ENVIRONS, **rollout_spec.env_vars}
+            rollout_cpus_per_task = rollout_spec.cpu * n_gpus_per_node
+            thread_env = get_thread_env_vars(
+                cpus_per_task=rollout_cpus_per_task,
+                existing_env_vars=rollout_spec.env_vars,
+            )
+            base_env_vars = {**BASE_ENVIRONS, **thread_env, **rollout_spec.env_vars}
             if spec["set_device_env"]:
                 base_env_vars[current_platform.device_control_env_var] = ",".join(
                     list(map(str, range(n_gpus_per_node)))
@@ -622,6 +620,14 @@ def _build_trainer_cmds(
             _env_vars["NCCL_CUMEM_ENABLE"] = "0"
             _env_vars["NCCL_NVLS_ENABLE"] = "0"
 
+        trainer_cpus_per_task = actor_spec.cpu * config.cluster.n_gpus_per_node
+        # Use per-GPU CPU count for thread env vars since torchrun spawns
+        # n_gpus_per_node processes, each inheriting these env vars.
+        # This matches Local and Ray launcher behavior.
+        thread_env = get_thread_env_vars(
+            cpus_per_task=actor_spec.cpu,
+            existing_env_vars=actor_spec.env_vars,
+        )
         launcher.submit_array(
             job_name="trainer",
             cmd=_build_trainer_cmds(
@@ -632,13 +638,14 @@ def _build_trainer_cmds(
             count=trainer_n_nodes,
             nodes=trainer_n_nodes,
             n_gpus_per_node=gpus_per_node,
-            cpus_per_task=actor_spec.cpu * config.cluster.n_gpus_per_node,
+            cpus_per_task=trainer_cpus_per_task,
             mem_per_task=actor_spec.mem * 1024 * config.cluster.n_gpus_per_node,
             container_image=actor_spec.image,
             srun_additional_args=actor_spec.srun_additional_args,
             container_mounts=actor_spec.mount,
             env_vars={
                 **BASE_ENVIRONS,
+                **thread_env,
                 **actor_spec.env_vars,
                 **_env_vars,
                 **tms_env_vars,