feat(BA-3024): Add custom resource alloc for agents in config

This change adds configuration for partitioning resources rather than
every agent always seeing the full resource pool. This
prevents unintended over-allocation that could crash kernels.

SHARED mode allows all agents to see full resources (useful for
stress testing). This is the same behavior as before.
AUTO_SPLIT automatically divides resources equally among agents.
MANUAL mode lets users specify exact per-agent allocations for all
resources.

Single-agent deployments remain unaffected and retain access to all
available hardware resources.

Author: hhoikoo

changes/6724.feature.md

@@ -0,0 +1 @@
+Add custom resource allocation in agent server config

configs/agent/sample.toml

@@ -143,6 +143,14 @@
 # If agents field is populated, this field indicates the default values for all
 # agents.
 [resource]
+  # Hard CPU allocation for this agent (e.g., 8 cores).
+  # Only used in MANUAL allocation mode.
+  # All agents must specify this value when allocation-mode is MANUAL.
+  ## allocated-cpu = 8
+  # Hard memory allocation for this agent (e.g., "32G").
+  # Only used in MANUAL allocation mode.
+  # All agents must specify this value when allocation-mode is MANUAL.
+  ## allocated-mem = "32G"
   # The number of CPU cores reserved for the operating system and the agent
   # service.
   reserved-cpu = 1
@@ -156,6 +164,12 @@
   # Currently this value is unused. In future releases, it may be used to preserve
   # the minimum disk space from the scratch disk allocation via loopback files.
   reserved-disk = "8G"
+  # Resource allocation mode for multi-agent scenarios.
+  # - `shared`: All agents share the full resource pool (default, backward
+  # compatible).
+  # - `auto-split`: Automatically divide resources equally (1/N) among all agents.
+  # - `manual`: Manually specify per-agent resource allocations via config.
+  allocation-mode = "shared"
   # The alignment of the reported main memory size to absorb tiny deviations from
   # per-node firmware/hardware settings. Recommended to be multiple of the
   # page/hugepage size (e.g., 2 MiB).
@@ -165,6 +179,10 @@
   # Affinity policy
   affinity-policy = "INTERLEAVED"
 
+  # Device-specific per-slot resource allocations.
+  # Only used in MANUAL allocation mode.
+  [resource.allocated-devices]
+
 # Pyroscope configuration
 [pyroscope]
   # Whether to enable Pyroscope profiling
@@ -409,24 +427,15 @@
 
   # Resource config overrides for the individual agent
   [agents.resource]
-    # The number of CPU cores reserved for the operating system and the agent
-    # service.
-    reserved-cpu = 1
-    # The memory space reserved for the operating system and the agent service. It
-    # is subtracted from the reported main memory size and not available for user
-    # workload allocation. Depending on the memory-align-size option and system
-    # configuration, this may not be the exact value but have slightly less or more
-    # values within the memory-align-size.
-    reserved-mem = 1073741824
-    # The disk space reserved for the operating system and the agent service.
-    # Currently this value is unused. In future releases, it may be used to preserve
-    # the minimum disk space from the scratch disk allocation via loopback files.
-    reserved-disk = 8589934592
-    # The alignment of the reported main memory size to absorb tiny deviations from
-    # per-node firmware/hardware settings. Recommended to be multiple of the
-    # page/hugepage size (e.g., 2 MiB).
-    memory-align-size = 16777216
-    # Resource allocation order
-    allocation-order = [ "cuda", "rocm", "tpu", "cpu", "mem",]
-    # Affinity policy
-    affinity-policy = 1
+    # Hard CPU allocation for this agent (e.g., 8 cores).
+    # Only used in MANUAL allocation mode.
+    # All agents must specify this value when allocation-mode is MANUAL.
+    ## allocated-cpu = 8
+    # Hard memory allocation for this agent (e.g., "32G").
+    # Only used in MANUAL allocation mode.
+    # All agents must specify this value when allocation-mode is MANUAL.
+    ## allocated-mem = "32G"
+
+    # Device-specific per-slot resource allocations.
+    # Only used in MANUAL allocation mode.
+    [agents.resource.allocated-devices]

src/ai/backend/agent/config/unified.py

@@ -11,6 +11,7 @@
 import os
 import sys
 import textwrap
+from decimal import Decimal
 from pathlib import Path
 from typing import (
     Any,
@@ -50,6 +51,7 @@
     BinarySizeField,
     ResourceGroupType,
     ServiceDiscoveryType,
+    SlotName,
 )
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.logging.config import LoggingConfig
@@ -80,6 +82,12 @@ class ScratchType(enum.StrEnum):
     K8S_NFS = "k8s-nfs"
 
 
+class ResourceAllocationMode(enum.StrEnum):
+    SHARED = "shared"
+    AUTO_SPLIT = "auto-split"
+    MANUAL = "manual"
+
+
 class AgentConfigValidationContext(BaseConfigValidationContext):
     is_invoked_subcommand: bool
 
@@ -862,7 +870,7 @@ class ContainerConfig(CommonContainerConfig, OverridableContainerConfig):
     pass
 
 
-class ResourceConfig(BaseConfigSchema):
+class CommonResourceConfig(BaseConfigSchema):
     reserved_cpu: int = Field(
         default=1,
         description="The number of CPU cores reserved for the operating system and the agent service.",
@@ -894,6 +902,18 @@ class ResourceConfig(BaseConfigSchema):
         validation_alias=AliasChoices("reserved-disk", "reserved_disk"),
         serialization_alias="reserved-disk",
     )
+    allocation_mode: ResourceAllocationMode = Field(
+        default=ResourceAllocationMode.SHARED,
+        description=textwrap.dedent("""
+        Resource allocation mode for multi-agent scenarios.
+        - `shared`: All agents share the full resource pool (default, backward compatible).
+        - `auto-split`: Automatically divide resources equally (1/N) among all agents.
+        - `manual`: Manually specify per-agent resource allocations via config.
+        """),
+        examples=[item.value for item in ResourceAllocationMode],
+        validation_alias=AliasChoices("allocation-mode", "allocation_mode"),
+        serialization_alias="allocation-mode",
+    )
     memory_align_size: BinarySizeField = Field(
         default=BinarySize.finite_from_str("16M"),
         description=(
@@ -936,6 +956,64 @@ def _parse_affinity_policy(cls, v: Any) -> AffinityPolicy:
         return v
 
 
+class OverridableResourceConfig(BaseConfigSchema):
+    allocated_cpu: Optional[int] = Field(
+        default=None,
+        description=textwrap.dedent("""
+        Hard CPU allocation for this agent (e.g., 8 cores).
+        Only used in MANUAL allocation mode.
+        All agents must specify this value when allocation-mode is MANUAL.
+        """),
+        examples=[8, 16],
+        validation_alias=AliasChoices("allocated-cpu", "allocated_cpu"),
+        serialization_alias="allocated-cpu",
+    )
+    allocated_mem: Optional[BinarySizeField] = Field(
+        default=None,
+        description=textwrap.dedent("""
+        Hard memory allocation for this agent (e.g., "32G").
+        Only used in MANUAL allocation mode.
+        All agents must specify this value when allocation-mode is MANUAL.
+        """),
+        examples=["32G", "64G"],
+        validation_alias=AliasChoices("allocated-mem", "allocated_mem"),
+        serialization_alias="allocated-mem",
+    )
+    allocated_devices: Mapping[SlotName, Decimal] = Field(
+        default_factory=dict,
+        description=textwrap.dedent("""
+        Device-specific per-slot resource allocations.
+        Only used in MANUAL allocation mode.
+        """),
+        examples=[{"cuda.mem": "0.3", "cuda.shares": "0.5"}],
+        validation_alias=AliasChoices("allocated-devices", "allocated_devices"),
+        serialization_alias="allocated-devices",
+    )
+
+    model_config = ConfigDict(
+        extra="allow",
+        arbitrary_types_allowed=True,
+    )
+
+    @model_validator(mode="after")
+    def validate_values_are_positive(self) -> Self:
+        if self.allocated_cpu is not None and self.allocated_cpu < 0:
+            raise ValueError(
+                f"Allocated cpu must not be a negative value, but given {self.allocated_cpu}"
+            )
+        if self.allocated_mem is not None and self.allocated_mem < 0:
+            raise ValueError(
+                f"Allocated mem must not be a negative value, but given {self.allocated_mem}"
+            )
+        if any(value < 0 for value in self.allocated_devices.values()):
+            raise ValueError("All allocated device resource values must not be a negative value")
+        return self
+
+
+class ResourceConfig(CommonResourceConfig, OverridableResourceConfig):
+    pass
+
+
 class EtcdConfig(BaseConfigSchema):
     namespace: str = Field(
         description="Etcd namespace",
@@ -1165,7 +1243,7 @@ class AgentOverrideConfig(BaseConfigSchema):
         default=None,
         description="Container config overrides for the individual agent",
     )
-    resource: ResourceConfig | None = Field(
+    resource: OverridableResourceConfig | None = Field(
         default=None,
         description="Resource config overrides for the individual agent",
     )
@@ -1228,6 +1306,10 @@ def agent_configs(self) -> Sequence[AgentUnifiedConfig]:
     def agent_ids(self) -> Sequence[AgentId]:
         return [AgentId(agent_config.agent.id) for agent_config in self.agent_configs]
 
+    @property
+    def resource_common(self) -> CommonResourceConfig:
+        return self.resource
+
     def with_updates(
         self,
         *,
@@ -1308,3 +1390,34 @@ def _validate_docker_config(self) -> Self:
                 DockerExtraConfig.model_validate(config.container.model_dump())
 
         return self
+
+    @model_validator(mode="after")
+    def _validate_resource_allocation_mode(self) -> Self:
+        match self.resource.allocation_mode:
+            case ResourceAllocationMode.SHARED | ResourceAllocationMode.AUTO_SPLIT:
+                for config in self.agent_configs:
+                    resource = config.resource
+                    if any([
+                        resource.allocated_cpu is not None,
+                        resource.allocated_mem is not None,
+                        resource.allocated_devices,
+                    ]):
+                        raise ValueError(
+                            "On non-MANUAL mode, config must not specify manual resource allocations"
+                        )
+
+            case ResourceAllocationMode.MANUAL:
+                for config in self.agent_configs:
+                    resource = config.resource
+                    if any([resource.allocated_cpu is None, resource.allocated_mem is None]):
+                        raise ValueError(
+                            "On MANUAL mode, config must specify cpu and mem resource allocations"
+                        )
+
+                slot_names = [
+                    set(config.resource.allocated_devices.keys()) for config in self.agent_configs
+                ]
+                if not all(slot_name == slot_names[0] for slot_name in slot_names):
+                    raise ValueError("All agents must have the same slots defined in the devices!")
+
+        return self