orchestrata_ultimate.agents.pool

Reference › orchestrata_ultimate › agents › pool

pool manages per-alias pools of volk instances (specification 005, non-functional requirement NFR1). VolkPool is built over a VolkRuntime (runtime) and a MultiagentPrologBridge (prolog_bridge), and is used by the task_executor. AgentPool is an alias subclass.

Each alias has a PoolConfig (types) holding minimum and maximum instances, scale thresholds and an idle timeout; _get_config creates a default from constants when none was set with set_pool_config. acquire returns an available instance for the alias and, when none is free and capacity remains, spawns a new one through _try_scale_up, which raises PoolScaleError on failure; a full pool raises PoolExhaustedError. release returns an instance to the pool, get_pool_status reports counts and utilisation as a PoolStatus, and scale_to grows or shrinks the pool to a target size and returns the resulting count.

At a glance

Class diagram

Dependency diagram

API

Classes

VolkPool

Decorators: traces(spec='005', story='US-001', reqs=['REQ-005-NFR1-001']) · Line 42

Manages per-alias Volk instance pools.

AgentPool

Bases: VolkPool · Line 214

Pool alias for VolkPool.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/agents/pool.py

"""Volk Pool — per-alias instance pool management.

Manages pools of Volk instances with configurable min/max limits,
auto-scaling, idle timeouts, and load-aware instance selection.

Author: Mohamed
Spec: 005-multi-agent (amended)
"""

# Spec: 005 - Multi-Agent System

import logging
from datetime import UTC, datetime

from orchestrata_ultimate.agents.constants import (
    DEFAULT_IDLE_TIMEOUT_SECONDS,
    DEFAULT_POOL_MAX,
    DEFAULT_POOL_MIN,
    DEFAULT_SCALE_DOWN_THRESHOLD,
    DEFAULT_SCALE_UP_THRESHOLD,
)
from orchestrata_ultimate.agents.exceptions import PoolExhaustedError, PoolScaleError
from orchestrata_ultimate.agents.prolog_bridge import MultiagentPrologBridge
from orchestrata_ultimate.agents.runtime import VolkRuntime
from orchestrata_ultimate.agents.types import (
    InstanceStatus,
    PoolConfig,
    PoolStatus,
    TaskDefinition,
    VolkInstance,
)
from orchestrata_ultimate.core.base.traces import traces

_logger = logging.getLogger(__name__)


def _utc_now() -> datetime:
    return datetime.now(UTC)


@traces(spec="005", story="US-001", reqs=["REQ-005-NFR1-001"])
class VolkPool:
    """Manages per-alias Volk instance pools.

    Provides acquire/release semantics with automatic scaling.
    Uses Prolog least_loaded_volk/2 for selection when multiple
    instances are available.

    Parameters
    ----------
    runtime : VolkRuntime
        Runtime for spawning/terminating instances.
    prolog : MultiagentPrologBridge
        Bridge for load-aware selection queries.
    """

    def __init__(
        self,
        runtime: VolkRuntime,
        prolog: MultiagentPrologBridge,
    ) -> None:
        self._runtime = runtime
        self._prolog = prolog
        self._configs: dict[str, PoolConfig] = {}

    @traces(spec="005", story="US-001", reqs=["REQ-005-NFR1-001"])
    def _get_config(self, alias: str) -> PoolConfig:
        """Get or create default config for an alias."""
        if alias not in self._configs:
            self._configs[alias] = PoolConfig(
                alias=alias,
                min_instances=DEFAULT_POOL_MIN,
                max_instances=DEFAULT_POOL_MAX,
                scale_up_threshold=DEFAULT_SCALE_UP_THRESHOLD,
                scale_down_threshold=DEFAULT_SCALE_DOWN_THRESHOLD,
                idle_timeout_seconds=DEFAULT_IDLE_TIMEOUT_SECONDS,
            )
        return self._configs[alias]

    async def set_pool_config(self, config: PoolConfig) -> None:
        """Set or update pool configuration for an alias."""
        self._configs[config.alias] = config

    @traces(spec="005", story="US-001", reqs=["REQ-005-US-001-001"])
    async def _try_scale_up(self, alias: str) -> VolkInstance:
        """Spawn a new instance for the alias; raise PoolScaleError on failure."""
        try:
            return await self._runtime.spawn(alias)
        except Exception as e:
            raise PoolScaleError(
                f"Failed to scale up pool for '{alias}': {e}",
                alias=alias,
            )

    @traces(spec="005", story="US-001", reqs=["REQ-005-NFR1-001"])
    async def acquire(
        self,
        alias: str,
        task: TaskDefinition | None = None,
    ) -> VolkInstance:
        """Acquire an available instance from the pool.

        Selection order:
        1. Find a READY instance for this alias
        2. If none available and under max capacity, spawn a new one
        3. If at capacity and all busy, raise PoolExhaustedError

        Raises
        ------
        PoolExhaustedError
            If pool is at capacity with no available instances.
        """
        config = self._get_config(alias)

        # Find available instances
        instances = await self._runtime.list_instances(
            alias=alias, status=InstanceStatus.READY
        )

        if instances:
            return instances[0]

        # No ready instances -- check if we can scale up
        all_instances = await self._runtime.list_instances(alias=alias)
        active_count = sum(1 for i in all_instances if i.is_active)

        if active_count < config.max_instances:
            return await self._try_scale_up(alias)

        raise PoolExhaustedError(
            f"Pool for '{alias}' exhausted: {active_count}/{config.max_instances} "
            f"instances active, none available",
            alias=alias,
        )

    async def release(self, instance_id: str) -> None:
        """Release an instance back to the pool.

        Marks the instance as READY so it can be acquired again.
        """
        self._runtime.mark_ready(instance_id)

    @traces(spec="005", story="US-004", reqs=["REQ-005-AC6-001"])
    async def get_pool_status(self, alias: str) -> PoolStatus:
        """Get current status of the pool for an alias."""
        config = self._get_config(alias)
        instances = await self._runtime.list_instances(alias=alias)

        ready = sum(1 for i in instances if i.status == InstanceStatus.READY)
        busy = sum(1 for i in instances if i.status == InstanceStatus.BUSY)
        draining = sum(1 for i in instances if i.status == InstanceStatus.DRAINING)
        total = sum(1 for i in instances if i.is_active)

        utilization = busy / total if total > 0 else 0.0

        return PoolStatus(
            alias=alias,
            total_instances=total,
            ready_instances=ready,
            busy_instances=busy,
            draining_instances=draining,
            config=config,
            utilization=utilization,
        )

    @traces(spec="005", story="US-001", reqs=["REQ-005-NFR1-001"])
    async def scale_to(self, alias: str, target_count: int) -> int:
        """Scale pool to target number of active instances.

        Returns the actual number of active instances after scaling.
        """
        config = self._get_config(alias)
        target_count = max(
            config.min_instances, min(target_count, config.max_instances)
        )

        instances = await self._runtime.list_instances(alias=alias)
        active = [i for i in instances if i.is_active]
        current = len(active)

        if current < target_count:
            # Scale up
            for _ in range(target_count - current):
                try:
                    await self._runtime.spawn(alias)
                except Exception as exc:
                    _logger.warning(
                        "Failed to spawn instance for alias '%s': %s",
                        alias,
                        exc,
                        exc_info=True,
                    )
                    break
        elif current > target_count:
            # Scale down -- terminate idle instances first
            ready = [i for i in active if i.status == InstanceStatus.READY]
            to_terminate = ready[: current - target_count]
            for inst in to_terminate:
                try:
                    await self._runtime.terminate(inst.instance_id)
                except Exception as exc:
                    _logger.warning(
                        "Failed to terminate instance '%s': %s",
                        inst.instance_id,
                        exc,
                        exc_info=True,
                    )

        final_instances = await self._runtime.list_instances(alias=alias)
        return sum(1 for i in final_instances if i.is_active)


# Class alias for orchestrata-ultimate naming
class AgentPool(VolkPool):
    """Pool alias for VolkPool."""

    pass