"""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()