feat(BA-4145): Add device partitioning for AUTO_SPLIT + SHARED

[hhoikoo]

resolves #8427 (BA-4145)

Overview

Implements device-based partitioning for multi-agent AUTO_SPLIT mode, enabling proper resource
isolation where devices (CPU/GPU) are exclusively assigned to agents while memory is shared.

Problem Statement

• Multi-agent configurations need exclusive device assignment to prevent resource contention
• Previous slot-based approach didn't partition physical devices between agents
• AUTO_SPLIT mode needed device-centric implementation per BEP-1041 design

Architecture

flowchart TB
    subgraph "Device Discovery"
        LD[_load_resources] --> GDM[GlobalDeviceMap]
        GDM --> GDI[GlobalDeviceInfo per device type]
    end

    subgraph "Partition Types"
        DP[DevicePartition<br/>Exclusive device IDs]
        SP[SlotPartition<br/>Shared slot amounts]
    end

    subgraph "ResourcePartitioner"
        GSA[generate_shared_assignments]
        GASA[generate_autosplit_assignments]
        CDP[_calculate_device_partitions<br/>DiscretePropertyAllocMap]
        CSP[_calculate_slot_partitions<br/>FractionAllocMap]
        GASA --> CDP --> DP
        GASA --> CSP --> SP
    end

    subgraph "Agent Configuration"
        DP --> ADA[_apply_device_assignments]
        SP --> ADA
        ADA --> AC1[Agent1 ComputerContext]
        ADA --> AC2[Agent2 ComputerContext]
    end

Loading

Key Components

ResourcePartitioner: Static methods for generating device assignments

• generate_shared_assignments: All agents see all devices (SHARED mode)
• generate_autosplit_assignments: Exclusive partitioning (AUTO_SPLIT mode)

Partition Types:

• DevicePartition: Exclusive device ID assignment (CPU cores, GPUs)
• SlotPartition: Shared slot amount division (memory)

Partitioning Logic:

• DiscretePropertyAllocMap → DevicePartition: Devices distributed via fill-from-front
• FractionAllocMap → SlotPartition: Slots divided evenly among agents

distribute_devices Function: Implements fill-from-front algorithm with natural sorting

• Given N devices and M agents: first r agents get ⌈N/M⌉ devices, rest get ⌊N/M⌋
• Natural sort ensures cuda0, cuda1, cuda2, cuda10 order (not lexicographic)

---

Checklist: (if applicable)

• Milestone metadata specifying the target backport version
• Mention to the original issue
• Installer updates including:
  • Fixtures for db schema changes
  • New mandatory config options
• Update of end-to-end CLI integration tests in ai.backend.test
• API server-client counterparts (e.g., manager API -> client SDK)
• Test case(s) to:
  • Demonstrate the difference of before/after
  • Demonstrate the flow of abstract/conceptual models with a concrete implementation
• Documentation
  • Contents in the docs directory
  • docstrings in public interfaces and type annotations

[Copilot]

Pull request overview

This PR refactors the resource allocation pipeline to implement device-based partitioning for AUTO_SPLIT mode and introduce a clearer separation between device discovery and allocation, aligning with BEP-1041.

Changes:

• Introduces device partitioning abstractions (DevicePartitioner, WholeDevicePartitioner, SharedDevicePartitioner), a natural device ID sorter, and helpers (distribute_devices, GlobalDeviceInfo, GlobalDeviceMap) to support device-centric allocation.
• Reworks ResourceAllocator.__ainit__ to perform device discovery, compute total/available slots via available_slots(), generate device assignments for agents based on allocation mode, and derive reserved slots from those assignments.
• Extensively updates and extends unit tests around resources and allocation behavior (including AUTO_SPLIT, MANUAL-as-SHARED behavior, device distribution, and the new helpers) to validate the new semantics.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
src/ai/backend/agent/types.py	Adds DeviceAssignments alias and DevicePartitioner protocol to formalize the device partitioning strategy surface.
src/ai/backend/agent/resources.py	Implements natural sorting, device distribution, whole/shared device partitioners, GlobalDeviceInfo/GlobalDeviceMap, reworks ResourceAllocator initialization, total slot calculation via available_slots(), device assignment generation and application, and reserved-slot computation based on assignments.
tests/unit/agent/test_resources.py	Adds tests for GlobalDeviceInfo, _create_global_devices, handling plugins with no devices, and the new _calculate_total_slots behavior using plugin available_slots().
tests/unit/agent/test_resource_allocation.py	Updates existing AUTO_SPLIT tests to device-based semantics, adds coverage for new allocation behaviors (AUTO_SPLIT device exclusivity, MANUAL behaving like SHARED, scaling factor invariants), and tests _natural_sort_key, distribute_devices, and the partitioner classes.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.