allocator.py

"""Space allocator \u2014 places bots on HostSpace instances by tier and bot volume.

Capacity model (2026-07-05 rewrite)
===================================

The old "1 user = 1 slot = 0.4 CPU" packer is gone.  In the new model:

* **Users do not bind to Spaces.**  A user picks a plan_tier (``free``
  or ``paid``).  Their ``space_id`` field is a soft UI hint, not an
  allocator constraint.

* **Spaces are pooled by tier.**  ``HostSpace.tier`` says which pool it
  belongs to.  A free user's bot will only ever land on a ``tier ==
  'free'`` Space.

* **Space capacity is measured in bot CPU/RAM, not user count.**
  Multiple users' bots can share a Space as long as the total reserved
  CPU/RAM stays under the limit.

* **Every Space permanently reserves a baseline** (``baseline_cpu_reserved``,
  default 0.1 CPU + 128 MB RAM) for its own orchestrator process and the
  control-plane channel back to the panel.  Subtracted from capacity
  BEFORE any bot reservation, so the orchestrator can never be starved.

* **New Spaces need admin approval** by default
  (``PANEL_CONFIG.require_space_approval``).  The allocator inserts a
  row in ``PENDING_APPROVAL``; an admin clicks Approve via the admin
  panel, which flips status to ``PROVISIONING`` and triggers HF API.
  When the panel operator sets the env-var to ``False``, allocator
  skips straight to ``PROVISIONING`` \u2014 useful for fully-autonomous
  installs.

Public API
----------

* :meth:`SpaceAllocator.assign_bot` \u2014 place a bot on a Space matching
  its owner's tier; creates a new one (with or without approval) if
  none fits.
* :meth:`SpaceAllocator.release_bot` \u2014 inverse.
* :meth:`SpaceAllocator.request_space` \u2014 system-initiated request,
  may end up in PENDING_APPROVAL.
* :meth:`SpaceAllocator.approve_space` / :meth:`reject_space` \u2014 admin
  flow.
* :meth:`SpaceAllocator.create_space_manual` \u2014 admin-initiated,
  bypasses approval.
* :meth:`SpaceAllocator.assign_user` \u2014 DEPRECATED shim, kept only so
  legacy signup code does not break; returns the user's ``space_id``
  without allocating anything new.

Errors
------

:meth:`SpaceAllocator.assign_bot` raises :class:`AllocationError` (or
:class:`SpacePendingApprovalError` when the only way to fit the bot
is to provision a new Space that requires admin OK first \u2014 callers
should surface a "admin will provision a Space soon" message).
"""
from __future__ import annotations

import asyncio
import logging
import uuid
from dataclasses import dataclass
from datetime import datetime
from typing import Optional

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from config import PANEL_CONFIG
from hf_auth import get_write_token, hf_write_api
from models import (
    PLAN_TIER_FREE,
    PLAN_TIERS,
    PLAN_TIER_PAID,
    BotInstance,
    DeploymentMode,
    HostSpace,
    SpaceStatus,
    User,
)

log = logging.getLogger("hosting-panel.allocator")


# ---- Port pool -------------------------------------------------------------

BOT_PORT_MIN: int = 19_000
BOT_PORT_MAX: int = 19_999   # 1000 concurrent bots per Space is plenty


def _ready_states() -> frozenset[SpaceStatus]:
    """Status values that allow new bot placements."""
    return frozenset({SpaceStatus.READY})


# ---- Errors ---------------------------------------------------------------

class AllocationError(Exception):
    """Raised when the allocator cannot satisfy a placement request."""


class SpacePendingApprovalError(AllocationError):
    """Raised when bot placement would require provisioning a new Space
    but admin approval is required and not yet granted.  Carries the
    :class:`HostSpace` row that was inserted in ``PENDING_APPROVAL`` so
    the caller can show the admin something to look at.
    """

    def __init__(self, message: str, space: HostSpace) -> None:
        super().__init__(message)
        self.space = space


# ---- Capacity maths --------------------------------------------------------

@dataclass
class Capacity:
    cpu: float
    ram_mb: int
    disk_mb: int = 0

    def fits_in(self, used: "Capacity", cap: "Capacity") -> bool:
        return (
            used.cpu + self.cpu <= cap.cpu
            and used.ram_mb + self.ram_mb <= cap.ram_mb
            and used.disk_mb + self.disk_mb <= cap.disk_mb
        )


# ---- Hard-limit helpers ----------------------------------------------------

def user_hard_cpu_limit(quota_cores: float) -> float:
    """Return the actual hard CPU limit enforced at runtime per bot.

    The user-facing quota is the *soft* plan; the hard limit is a small
    underbooking (default 5 %) so the host's CPU scheduler can always
    service the orchestrator.  **Never surface this number to end-users.**
    """
    return round(quota_cores * PANEL_CONFIG.user_hard_cpu_factor, 4)


def space_allocable_capacity(cap_cores: float, cap_ram_mb: int,
                             cap_disk_mb: int,
                             baseline_cpu: float = 0.0,
                             baseline_ram_mb: int = 0) -> "Capacity":
    """Capacity actually available for placing bots on a Space.

    Subtracted ``baseline_cpu`` and ``baseline_ram_mb`` are the
    orchestrator reservations; everything that lands in this object is
    available for ``assign_bot`` to draw from.
    """
    return Capacity(
        cpu=max(0.0, cap_cores - baseline_cpu),
        ram_mb=max(0, cap_ram_mb - baseline_ram_mb),
        disk_mb=cap_disk_mb,
    )


# ---- Allocator -------------------------------------------------------------

class SpaceAllocator:
    """Stateless service that owns placement decisions."""

    def __init__(self) -> None:
        self._lock = asyncio.Lock()

    # ------------------------------------------------------------------
    # DEPRECATED — legacy user-to-space binding.  Kept so existing
    # signup / admin code does not break; does NOT create a new Space.
    # ------------------------------------------------------------------
    async def assign_user(
        self,
        db: AsyncSession,
        user: User,
    ) -> Optional[HostSpace]:
        """Compatibility shim.

        .. deprecated::
            The new allocator binds bots, not users, to Spaces.  This
            method only reads ``user.space_id`` and returns the existing
            row (if any).  It does NOT create a new Space.  Use
            :meth:`assign_bot` instead.
        """
        log.debug(
            "assign_user called for %s \u2014 deprecated, no allocation performed",
            user.username,
        )
        if not user.space_id:
            return None
        space = await db.get(HostSpace, user.space_id)
        if not space:
            return None
        if space.status not in _ready_states():
            return None
        return space

    # ------------------------------------------------------------------
    # Bot placement \u2014 the new hot path.
    # ------------------------------------------------------------------
    async def assign_bot(
        self,
        db: AsyncSession,
        bot: BotInstance,
        owner: User,
    ) -> HostSpace:
        """Reserve CPU/RAM on a Space matching ``owner.plan_tier`` and
        allocate a port for the bot.  Creates a new Space if none fit.

        Raises:
            AllocationError: user quota exceeded, no Space fits, or HF
                API rejected the new Space.
            SpacePendingApprovalError: only way to fit is a Space that
                needs admin approval first; carries the pending Space.
        """
        if not bot.cpu_cores or not bot.ram_mb:
            raise AllocationError("Bot must have cpu_cores and ram_mb set")

        # Hard-cap the bot's reservation to the user's hard limit so a
        # bot never consumes more than the user's slice (admin can still
        # override per-bot from the bot detail page; this is just the
        # default ceiling).
        bot_cpu_reserved = min(
            float(bot.cpu_cores),
            user_hard_cpu_limit(float(owner.cpu_quota_cores)),
        )
        bot_ram_reserved = int(bot.ram_mb)

        need = Capacity(cpu=bot_cpu_reserved, ram_mb=bot_ram_reserved, disk_mb=0)

        async with self._lock:
            # 1) user-level quota check (sum of reservations across
            #    owner's other bots).
            used_cpu = sum(
                (b.cpu_cores or 0) for b in owner.bots
                if b.id != bot.id and b.deployment_mode == DeploymentMode.MULTITENANT
            )
            used_ram = sum(
                (b.ram_mb or 0) for b in owner.bots
                if b.id != bot.id and b.deployment_mode == DeploymentMode.MULTITENANT
            )
            if used_cpu + bot_cpu_reserved > owner.cpu_quota_cores:
                raise AllocationError(
                    f"User CPU quota exceeded: have {owner.cpu_quota_cores}, "
                    f"used {used_cpu}, need {bot_cpu_reserved}"
                )
            if used_ram + bot_ram_reserved > owner.ram_quota_mb:
                raise AllocationError(
                    f"User RAM quota exceeded: have {owner.ram_quota_mb} MB, "
                    f"used {used_ram}, need {bot_ram_reserved} MB"
                )

            # 2) Find a Space in the owner's tier with room for the bot.
            tier = owner.plan_tier if owner.plan_tier in PLAN_TIERS else PLAN_TIER_FREE
            candidates = await self._candidate_spaces(db, tier=tier)
            target = None
            for space in candidates:
                if space.status not in _ready_states():
                    continue
                used = Capacity(
                    space.cpu_used_cores,
                    space.ram_used_mb,
                    space.disk_used_mb,
                )
                alloc = space_allocable_capacity(
                    space.cpu_capacity_cores,
                    space.ram_capacity_mb,
                    space.disk_capacity_mb,
                    baseline_cpu=space.baseline_cpu_reserved or 0.0,
                    baseline_ram_mb=space.baseline_ram_reserved_mb or 0,
                )
                if need.fits_in(used, alloc):
                    target = space
                    break
            if target is None:
                # 3) Try to create one.  Honors require_space_approval.
                await self._ensure_capacity_for(db, owner=owner, need_cpu=need.cpu, need_ram_mb=need.ram_mb)
                # Re-search after provisioning triggered.  If we still
                # see no fit, the new Space is probably pending
                # approval \u2014 raise that variant.
                candidates = await self._candidate_spaces(db, tier=tier)
                for space in candidates:
                    if space.status not in _ready_states():
                        continue
                    used = Capacity(
                        space.cpu_used_cores, space.ram_used_mb, space.disk_used_mb,
                    )
                    alloc = space_allocable_capacity(
                        space.cpu_capacity_cores, space.ram_capacity_mb,
                        space.disk_capacity_mb,
                        baseline_cpu=space.baseline_cpu_reserved or 0.0,
                        baseline_ram_mb=space.baseline_ram_reserved_mb or 0,
                    )
                    if need.fits_in(used, alloc):
                        target = space
                        break
                if target is None:
                    # Was the create placed in PENDING_APPROVAL?
                    pending = (await db.execute(
                        select(HostSpace).where(
                            HostSpace.status == SpaceStatus.PENDING_APPROVAL,
                            HostSpace.tier == tier,
                        ).order_by(HostSpace.requested_at.desc()).limit(1)
                    )).scalar_one_or_none()
                    if pending:
                        raise SpacePendingApprovalError(
                            "Bot start requires a new Space, but admin "
                            "approval is pending. An admin must approve "
                            f"space '{pending.name}' before this bot can run.",
                            space=pending,
                        )
                    raise AllocationError(
                        "No Space available in tier "
                        f"'{tier}' for bot {bot.slug} "
                        f"(need {need.cpu:.2f} CPU / {need.ram_mb} MB)"
                    )

            # 4) Allocate a free port in target.
            taken_ports = set(
                p for (p,) in (await db.execute(
                    select(BotInstance.port).where(
                        BotInstance.space_id == target.id,
                        BotInstance.port.is_not(None),
                    )
                )).all() if p is not None
            )
            free_port = None
            for p in range(BOT_PORT_MIN, BOT_PORT_MAX + 1):
                if p not in taken_ports:
                    free_port = p
                    break
            if free_port is None:
                raise AllocationError(
                    f"No free port in [{BOT_PORT_MIN}, {BOT_PORT_MAX}] "
                    f"for space {target.name}"
                )

            # 5) Commit.
            bot.space_id = target.id
            bot.port = free_port
            bot.deployment_mode = DeploymentMode.MULTITENANT
            target.cpu_used_cores += need.cpu
            target.ram_used_mb += need.ram_mb
            await db.commit()
            log.info(
                "Placed bot %s on %s tier=%s port=%d "
                "(+%.2f CPU / +%d MB RAM; baseline=%.2f CPU reserved)",
                bot.slug, target.name, target.tier, free_port,
                need.cpu, need.ram_mb,
                target.baseline_cpu_reserved,
            )
            return target

    async def _candidate_spaces(self, db: AsyncSession, *, tier: str):
        """All Spaces in ``tier`` ordered by packable CPU descending
        (we want bots to land on the roomiest Space first)."""
        result = await db.execute(
            select(HostSpace)
            .where(HostSpace.tier == tier)
            .order_by((HostSpace.cpu_capacity_cores - HostSpace.cpu_used_cores).desc())
        )
        return result.scalars().all()

    async def _ensure_capacity_for(
        self,
        db: AsyncSession,
        *,
        owner: User,
        need_cpu: float,
        need_ram_mb: int,
    ) -> Optional[HostSpace]:
        """Make sure the tier pool has enough packable capacity for a
        bot that needs ``need_cpu``/``need_ram_mb``.  Provisions a new
        Space when it doesn't.

        Returns the new Space if one was created, ``None`` if the
        existing pool already had room (which we missed \u2014 shouldn't
        happen, defensive).
        """
        tier = owner.plan_tier if owner.plan_tier in PLAN_TIERS else PLAN_TIER_FREE
        # Hardware tier matters: free bots need free_tier_hardware,
        # paid bots need paid_tier_hardware.
        hardware = (
            PANEL_CONFIG.paid_tier_hardware
            if tier == PLAN_TIER_PAID
            else PANEL_CONFIG.free_tier_hardware
        )
        existing = await self._candidate_spaces(db, tier=tier)
        for space in existing:
            if space.status in _ready_states():
                # The caller (assign_bot) already filtered these, so if
                # we got here it means none fit \u2014 skip.
                continue
        # Fire a request \u2014 either PENDING_APPROVAL or directly
        # PROVISIONING depending on policy.
        space = await self.request_space(
            db,
            tier=tier,
            hardware=hardware,
            requested_by_user_id=None,
        )
        return space

    # ------------------------------------------------------------------
    # Space lifecycle \u2014 approval-driven provisioning
    # ------------------------------------------------------------------
    async def request_space(
        self,
        db: AsyncSession,
        *,
        tier: str = PLAN_TIER_FREE,
        hardware: Optional[str] = None,
        requested_by_user_id: Optional[int] = None,
        cpu_cores: Optional[float] = None,
        ram_mb: Optional[int] = None,
        disk_mb: Optional[int] = None,
        operator_note: Optional[str] = None,
    ) -> HostSpace:
        """System-initiated request for a new Space.

        If ``PANEL_CONFIG.require_space_approval`` is True (default),
        the row is inserted in ``PENDING_APPROVAL`` and **no** HF API
        call is made until :meth:`approve_space` runs.  Otherwise the
        row goes straight to ``PROVISIONING`` and triggers HF API in
        the background.
        """
        if tier not in PLAN_TIERS:
            tier = PLAN_TIER_FREE
        if hardware is None:
            hardware = (
                PANEL_CONFIG.paid_tier_hardware
                if tier == PLAN_TIER_PAID
                else PANEL_CONFIG.free_tier_hardware
            )
        cpu = float(cpu_cores or PANEL_CONFIG.new_space_cpu_cores)
        ram = int(ram_mb or PANEL_CONFIG.new_space_ram_mb)
        disk = int(disk_mb or PANEL_CONFIG.new_space_disk_mb)

        now = datetime.utcnow()
        name = self._next_space_name(tier)
        owner = PANEL_CONFIG.hf_owner or "auto"
        repo_id = f"{owner}/{name}" if owner != "auto" else name

        initial_status = (
            SpaceStatus.PENDING_APPROVAL
            if PANEL_CONFIG.require_space_approval
            else SpaceStatus.PROVISIONING
        )

        space = HostSpace(
            name=repo_id,
            hf_owner=owner,
            tier=tier,
            requested_at=now,
            requested_by_user_id=requested_by_user_id,
            operator_note=operator_note,
            hardware_tier=hardware,
            cpu_capacity_cores=cpu,
            ram_capacity_mb=ram,
            disk_capacity_mb=disk,
            baseline_cpu_reserved=PANEL_CONFIG.space_baseline_cpu_reserved,
            baseline_ram_reserved_mb=PANEL_CONFIG.space_baseline_ram_reserved_mb,
            status=initial_status,
        )
        db.add(space)
        await db.commit()
        await db.refresh(space)
        log.info(
            "Space requested: %s tier=%s status=%s by_user=%s",
            space.name, tier, initial_status.value, requested_by_user_id,
        )

        if not PANEL_CONFIG.require_space_approval:
            # Auto-approve path: trigger HF API immediately.
            token = get_write_token()
            if token:
                asyncio.create_task(self._create_remote_space(space.name, token))
            else:
                log.error(
                    "No HF write token configured; Space %s will stay in PROVISIONING until admin forces it.",
                    space.name,
                )
        return space

    async def approve_space(
        self,
        db: AsyncSession,
        space_id: int,
        approved_by_user_id: int,
    ) -> HostSpace:
        """Approve a pending Space \u2014 transitions to PROVISIONING and
        triggers HF API."""
        space = await db.get(HostSpace, space_id)
        if not space:
            raise AllocationError(f"Space {space_id} not found")
        if space.status != SpaceStatus.PENDING_APPROVAL:
            raise AllocationError(
                f"Space {space.name} is in status '{space.status.value}'; "
                "only PENDING_APPROVAL spaces can be approved."
            )
        space.status = SpaceStatus.PROVISIONING
        space.approved_at = datetime.utcnow()
        space.approved_by_user_id = approved_by_user_id
        await db.commit()
        await db.refresh(space)
        token = get_write_token()
        if token:
            asyncio.create_task(self._create_remote_space(space.name, token))
        else:
            log.error(
                "No HF write token for Space %s; admin needs to set HF_TOKEN_WRITE.",
                space.name,
            )
        log.info(
            "Space %s approved by user=%s",
            space.name, approved_by_user_id,
        )
        return space

    async def reject_space(
        self,
        db: AsyncSession,
        space_id: int,
        *,
        reason: str = "",
    ) -> HostSpace:
        """Mark a pending Space as ERROR and drop its bot reservations."""
        space = await db.get(HostSpace, space_id)
        if not space:
            raise AllocationError(f"Space {space_id} not found")
        if space.status != SpaceStatus.PENDING_APPROVAL:
            raise AllocationError(
                f"Space {space.name} is in status '{space.status.value}'; "
                "only PENDING_APPROVAL spaces can be rejected."
            )
        space.status = SpaceStatus.ERROR
        if reason:
            space.operator_note = (space.operator_note or "") + f"\n[REJECTED] {reason}"
        await db.commit()
        await db.refresh(space)
        log.info("Space %s rejected (reason=%r)", space.name, reason)
        return space

    async def create_space_manual(
        self,
        db: AsyncSession,
        *,
        tier: str = PLAN_TIER_FREE,
        hardware: Optional[str] = None,
        requested_by_user_id: Optional[int] = None,
        cpu_cores: Optional[float] = None,
        ram_mb: Optional[int] = None,
        disk_mb: Optional[int] = None,
        operator_note: Optional[str] = None,
    ) -> HostSpace:
        """Admin manually creates a Space.  Status starts as
        ``PROVISIONING`` and HF API fires immediately \u2014 admin
        approval is implicit in the manual act.
        """
        if tier not in PLAN_TIERS:
            tier = PLAN_TIER_FREE
        if hardware is None:
            hardware = (
                PANEL_CONFIG.paid_tier_hardware
                if tier == PLAN_TIER_PAID
                else PANEL_CONFIG.free_tier_hardware
            )
        cpu = float(cpu_cores or PANEL_CONFIG.new_space_cpu_cores)
        ram = int(ram_mb or PANEL_CONFIG.new_space_ram_mb)
        disk = int(disk_mb or PANEL_CONFIG.new_space_disk_mb)

        now = datetime.utcnow()
        name = self._next_space_name(tier)
        owner = PANEL_CONFIG.hf_owner or "auto"
        repo_id = f"{owner}/{name}" if owner != "auto" else name

        space = HostSpace(
            name=repo_id,
            hf_owner=owner,
            tier=tier,
            requested_at=now,
            approved_at=now,
            approved_by_user_id=requested_by_user_id,
            requested_by_user_id=requested_by_user_id,
            operator_note=operator_note,
            hardware_tier=hardware,
            cpu_capacity_cores=cpu,
            ram_capacity_mb=ram,
            disk_capacity_mb=disk,
            baseline_cpu_reserved=PANEL_CONFIG.space_baseline_cpu_reserved,
            baseline_ram_reserved_mb=PANEL_CONFIG.space_baseline_ram_reserved_mb,
            status=SpaceStatus.PROVISIONING,
        )
        db.add(space)
        await db.commit()
        await db.refresh(space)
        log.info(
            "Space manually created by admin=%s: %s tier=%s",
            requested_by_user_id, space.name, tier,
        )
        token = get_write_token()
        if token:
            asyncio.create_task(self._create_remote_space(space.name, token))
        else:
            log.error(
                "No HF write token for manual Space %s.",
                space.name,
            )
        return space

    # ------------------------------------------------------------------
    # Bot release \u2014 inverse of assign_bot
    # ------------------------------------------------------------------
    async def release_bot(
        self,
        db: AsyncSession,
        bot: BotInstance,
    ) -> None:
        """Return the bot's CPU/RAM to the Space; clear port."""
        async with self._lock:
            if bot.space_id and bot.cpu_cores:
                space = await db.get(HostSpace, bot.space_id)
                if space:
                    bot_cpu = float(bot.cpu_cores or 0.0)
                    space.cpu_used_cores = max(
                        0.0, float(space.cpu_used_cores or 0.0) - bot_cpu,
                    )
                    space.ram_used_mb = max(
                        0,
                        int(space.ram_used_mb or 0) - int(bot.ram_mb or 0),
                    )
            bot.port = None
            bot.pid = None
            await db.commit()

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------
    @staticmethod
    def _next_space_name(tier: str) -> str:
        suffix = uuid.uuid4().hex[:8]
        return f"hostspace-{tier}-{suffix}"

    async def _create_remote_space(self, repo_id: str, token: str) -> None:
        api = hf_write_api()
        try:
            api.create_repo(
                repo_id=repo_id,
                repo_type="space",
                space_sdk="docker",
                private=True,
            )
        except Exception as exc:
            log.error("Failed to provision HF Space %s: %s", repo_id, exc)
            from database import async_session
            async with async_session() as s:
                sp = (await s.execute(
                    select(HostSpace).where(HostSpace.name == repo_id)
                )).scalar_one_or_none()
                if sp:
                    sp.status = SpaceStatus.ERROR
                    await s.commit()
            return

        from database import async_session
        async with async_session() as s:
            sp = (await s.execute(
                select(HostSpace).where(HostSpace.name == repo_id)
            )).scalar_one_or_none()
            if sp:
                sp.status = SpaceStatus.READY
                sp.hf_space_url = (
                    f"https://{repo_id.replace('/', '-').replace('_', '-').lower()}.hf.space"
                )
                await s.commit()


# Module-level singleton.
ALLOCATOR = SpaceAllocator()