infer-stack 0.6.0__py3-none-any.whl

infer_stack/contracts.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import Any
+
+from .config import CONFIG_FILE, initial_config, load_yaml
+from .hardware import simulate_inventory
+from .paths import config_root
+from .resolver import resolve
+from .profile_runtime import default_base_url
+
+
+def _default_access(
+    service: dict[str, Any], deployment: dict[str, Any]
+) -> dict[str, Any]:
+    backend = str(deployment.get('backend', 'compose')).lower()
+    if backend == 'kubeai':
+        return {
+            'name': 'kubeai-frontdoor',
+            'kind': 'openai-compatible',
+            'base_url': default_base_url(deployment),
+            'request_model_name': service['profile_public_name'],
+            'auth_env_name': 'KUBEAI_OPENAI_API_KEY',
+            'auth_placeholder': 'EMPTY',
+            'auth_required': False,
+            'notes': [
+                'Use the KubeAI OpenAI-compatible front door for routed requests.'
+            ],
+        }
+    return {
+        'name': 'compose-router',
+        'kind': 'openai-compatible',
+        'base_url': default_base_url(deployment),
+        'request_model_name': service['served_model_name'],
+        'auth_env_name': 'LITELLM_MASTER_KEY',
+        'auth_placeholder': 'SET_LITELLM_MASTER_KEY_IN_ENV',
+        'auth_required': True,
+        'notes': ['Use the LiteLLM router front door for routed requests.'],
+    }
+
+
+def _additional_accesses(
+    service: dict[str, Any], deployment: dict[str, Any]
+) -> list[dict[str, Any]]:
+    benchmark_transport = deepcopy(service.get('benchmark_transport', {}))
+    if not benchmark_transport:
+        return []
+    default = _default_access(service, deployment)
+    default_auth_env = (
+        'VLLM_API_KEY'
+        if benchmark_transport.get('kind') == 'vllm-direct'
+        else default['auth_env_name']
+    )
+    default_auth_placeholder = (
+        'EMPTY'
+        if benchmark_transport.get('kind') == 'vllm-direct'
+        else default['auth_placeholder']
+    )
+    access = {
+        'name': benchmark_transport.get('name')
+        or benchmark_transport.get('kind')
+        or 'compatibility-access',
+        'kind': benchmark_transport.get('kind') or default['kind'],
+        'base_url': benchmark_transport.get('base_url') or default['base_url'],
+        'request_model_name': benchmark_transport.get('request_model_name')
+        or service.get('hf_model_id')
+        or service['served_model_name'],
+        'auth_env_name': benchmark_transport.get('api_key_env')
+        or default_auth_env,
+        'auth_placeholder': benchmark_transport.get('api_key_placeholder')
+        or default_auth_placeholder,
+        'auth_required': bool(
+            benchmark_transport.get(
+                'api_key_required', default_auth_placeholder != 'EMPTY'
+            )
+        ),
+        'notes': [
+            'Optional compatibility access hint retained for external integrations.'
+        ],
+    }
+    dedupe_key = (
+        access['kind'],
+        access['base_url'],
+        access['request_model_name'],
+    )
+    default_key = (
+        default['kind'],
+        default['base_url'],
+        default['request_model_name'],
+    )
+    if dedupe_key == default_key:
+        return []
+    return [access]
+
+
+def build_profile_contract(deployment: dict[str, Any]) -> dict[str, Any]:
+    services = []
+    for service in deployment.get('services', []):
+        services.append(
+            {
+                'profile_name': service['profile_name'],
+                'public_name': service['profile_public_name'],
+                'service_name': service['service_name'],
+                'kubernetes_name': service['kubernetes_name'],
+                'model': {
+                    'model_ref': service['model_ref'],
+                    'hf_model_id': service['hf_model_id'],
+                    'logical_model_name': service['logical_model_name'],
+                    'served_model_name': service['served_model_name'],
+                    'served_aliases': deepcopy(
+                        service.get('served_aliases', [])
+                    ),
+                    'tokenizer_name': service['tokenizer_name'],
+                    'modalities': deepcopy(service.get('modalities', [])),
+                },
+                'protocol': {
+                    'mode': service['protocol_mode'],
+                    'engine': service['engine'],
+                    'features': deepcopy(service.get('features', [])),
+                },
+                'access': {
+                    'default': _default_access(service, deployment),
+                    'additional': _additional_accesses(service, deployment),
+                },
+                'runtime': {
+                    'resource_profile': service['resource_profile'],
+                    'priority_class_name': service.get('priority_class_name'),
+                    'min_replicas': service['min_replicas'],
+                    'max_replicas': service['max_replicas'],
+                    'tensor_parallel_size': service['tensor_parallel_size'],
+                    'data_parallel_size': service['data_parallel_size'],
+                    'max_model_len': service['max_model_len'],
+                    'gpu_memory_utilization': service['gpu_memory_utilization'],
+                    'max_num_batched_tokens': service['max_num_batched_tokens'],
+                    'max_num_seqs': service['max_num_seqs'],
+                    'enable_prefix_caching': service['enable_prefix_caching'],
+                    'extra_args': deepcopy(service.get('extra_args', [])),
+                },
+                'placement': {
+                    'gpu_indices': deepcopy(service.get('gpu_indices', [])),
+                    'placement': deepcopy(service.get('placement', {})),
+                    'topology': deepcopy(service.get('topology', {})),
+                },
+                'notes': deepcopy(service.get('notes', [])),
+                'caveats': deepcopy(service.get('audit_notes', [])),
+            }
+        )
+    return {
+        'schema_version': 1,
+        'kind': 'serving-profile-contract',
+        'backend': deployment['backend'],
+        'profile': {
+            'name': deployment['serving_profile']['name'],
+            'public_name': deployment['serving_profile']['public_name'],
+            'description': deployment['serving_profile'].get('description', ''),
+            'kubernetes_name': deployment['serving_profile'].get(
+                'kubernetes_name', ''
+            ),
+            'service_name': deployment['serving_profile'].get(
+                'service_name', ''
+            ),
+            'protocol_mode': deployment['serving_profile'].get(
+                'protocol_mode', ''
+            ),
+            'engine': deployment['serving_profile'].get('engine', ''),
+            'resource_profile': deployment['serving_profile'].get(
+                'resource_profile', ''
+            ),
+            'notes': deepcopy(deployment['serving_profile'].get('notes', [])),
+            'caveats': deepcopy(
+                deployment['serving_profile'].get('audit_notes', [])
+            ),
+        },
+        'router': {
+            'type': deployment.get('router', {}).get('type', ''),
+            'aliases': deepcopy(
+                deployment.get('router', {}).get('aliases', {})
+            ),
+        },
+        'services': services,
+    }
+
+
+def describe_profile_contract(
+    config: dict[str, Any],
+    *,
+    resolve_fn,
+    profile_name: str | None = None,
+    inventory: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    deployment = resolve_fn(
+        config, inventory=inventory, profile_name=profile_name
+    )
+    return build_profile_contract(deployment)
+
+
+def load_profile_contract(
+    profile_name: str,
+    *,
+    backend: str | None = None,
+    simulate_hardware_spec: str | None = None,
+) -> dict[str, Any]:
+    config_path = config_root() / CONFIG_FILE
+    if config_path.exists():
+        config = load_yaml(config_path)
+    else:
+        config = initial_config()
+    config.setdefault('catalog', {})
+    config['catalog']['builtin_models'] = True
+    config['catalog']['builtin_profiles'] = True
+    if backend is not None:
+        config['backend'] = backend
+    inventory = (
+        simulate_inventory(simulate_hardware_spec)
+        if simulate_hardware_spec
+        else None
+    )
+    return describe_profile_contract(
+        config,
+        resolve_fn=resolve,
+        profile_name=profile_name,
+        inventory=inventory,
+    )

infer_stack/diff_prompt.py

@@ -0,0 +1,117 @@
+"""Render-time diff confirmation.
+
+Builds a unified diff of every file the renderer is about to overwrite, prints
+each one as its own panel using Rich syntax-highlighting, and asks the user to
+accept the changes interactively unless ``assume_yes=True`` is passed.
+"""
+
+from __future__ import annotations
+
+from difflib import unified_diff
+from pathlib import Path
+from typing import Mapping
+
+from rich.console import Console, Group
+from rich.panel import Panel
+from rich.prompt import Confirm
+from rich.syntax import Syntax
+from rich.text import Text
+
+
+def _read(path: Path) -> str:
+    if not path.exists():
+        return ''
+    return path.read_text(encoding='utf-8')
+
+
+def _diff(old: str, new: str, label: Path) -> str:
+    return ''.join(
+        unified_diff(
+            old.splitlines(keepends=True),
+            new.splitlines(keepends=True),
+            fromfile=f'a/{label}',
+            tofile=f'b/{label}',
+            n=3,
+        )
+    )
+
+
+def confirm_writes(
+    planned: Mapping[Path, str],
+    *,
+    assume_yes: bool,
+    console: Console | None = None,
+    title: str = 'Pending render changes',
+) -> bool:
+    """Display a per-file diff and prompt for acceptance.
+
+    Returns True if the writes should proceed, False if the user declined.
+    A planned write whose new content matches the existing file is shown as
+    "no changes" and never blocks acceptance on its own.
+    """
+    console = console or Console()
+
+    panels: list[Panel] = []
+    any_changes = False
+    new_files: list[Path] = []
+    modified_files: list[Path] = []
+    for path, new_content in planned.items():
+        old = _read(path)
+        if old == new_content:
+            panels.append(
+                Panel(
+                    Text('(no changes)', style='dim'),
+                    title=str(path),
+                    border_style='green',
+                    title_align='left',
+                )
+            )
+            continue
+        any_changes = True
+        diff_text = _diff(old, new_content, path)
+        if not diff_text:
+            # New file — synthesize a diff against an empty original.
+            diff_text = _diff('', new_content, path)
+        if not old:
+            new_files.append(path)
+            border = 'yellow'
+            tag = '[bold yellow]NEW[/bold yellow]'
+        else:
+            modified_files.append(path)
+            border = 'cyan'
+            tag = '[bold cyan]MOD[/bold cyan]'
+        body = Syntax(
+            diff_text,
+            'diff',
+            theme='ansi_dark',
+            background_color='default',
+            word_wrap=False,
+            line_numbers=False,
+        )
+        panels.append(
+            Panel(
+                body,
+                title=f'{tag} {path}',
+                border_style=border,
+                title_align='left',
+            )
+        )
+
+    console.rule(f'[bold]{title}[/bold]')
+    console.print(Group(*panels))
+
+    if not any_changes:
+        console.print('[green]No changes to apply.[/green]')
+        return True
+
+    summary = (
+        f'[cyan]{len(modified_files)} modified[/cyan], '
+        f'[yellow]{len(new_files)} new[/yellow]'
+    )
+    console.print(f'Summary: {summary}')
+
+    if assume_yes:
+        console.print('[green]--yes provided; applying changes.[/green]')
+        return True
+
+    return Confirm.ask('Apply these changes?', default=False, console=console)

infer_stack/docker_utils.py

@@ -0,0 +1,230 @@
+from __future__ import annotations
+
+import errno
+import json
+import socket
+import subprocess
+from pathlib import Path
+
+
+class DockerCommandError(RuntimeError):
+    pass
+
+
+class PortInUseError(RuntimeError):
+    """One or more host ports the stack needs are already bound."""
+
+    def __init__(self, busy: list[tuple[str, int, str]]):
+        self.busy = busy
+        lines = ['Cannot start stack: required host ports are already bound.']
+        for service, port, host in busy:
+            lines.append('')
+            lines.append(
+                f'  {service}: {host}:{port} is already in use. Find the owner with:'
+            )
+            lines.append(f"    ss -tlnp 'sport = :{port}'")
+            lines.append(f'    sudo lsof -nP -iTCP:{port} -sTCP:LISTEN')
+            lines.append(f'    docker ps --filter publish={port}')
+        lines.append('')
+        lines.append(
+            'If the conflict is a leftover container from this stack, run '
+            '`infer-stack down` (or `docker stop <name> && docker rm <name>`).'
+        )
+        lines.append(
+            'If a non-stack process owns the port, either stop that process or '
+            'pick different ports: `infer-stack setup --litellm-port N '
+            '--open-webui-port M`, then `infer-stack render --yes`.'
+        )
+        super().__init__('\n'.join(lines))
+
+
+def our_published_ports(
+    compose_cmd: str, compose_file: Path, env_file: Path
+) -> set[int]:
+    """Return host ports currently published by our own compose project.
+
+    Used by the pre-flight check to skip ports that are "in use" only because
+    one of our containers is already publishing them. Requires
+    ``docker compose ps --format json`` (compose v2.6+); on older versions or
+    when the command otherwise fails this returns an empty set, falling back
+    to the strict check.
+    """
+    if not compose_file.exists():
+        return set()
+    cmd = compose_cmd.split() + [
+        '-f',
+        str(compose_file),
+        '--env-file',
+        str(env_file),
+        'ps',
+        '--format',
+        'json',
+    ]
+    try:
+        proc = subprocess.run(
+            cmd, capture_output=True, text=True, check=False, timeout=10
+        )
+    except (subprocess.SubprocessError, OSError):
+        return set()
+    if proc.returncode != 0 or not proc.stdout.strip():
+        return set()
+    ports: set[int] = set()
+    # `docker compose ps --format json` emits either one JSON object per line
+    # (newer versions) or a single JSON array (older versions). Accept both.
+    text = proc.stdout.strip()
+    parsed_any = False
+    if text.startswith('['):
+        try:
+            for entry in json.loads(text):
+                _collect_published_ports(entry, ports)
+                parsed_any = True
+        except json.JSONDecodeError:
+            return set()
+    if not parsed_any:
+        for line in text.splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                entry = json.loads(line)
+            except json.JSONDecodeError:
+                continue
+            _collect_published_ports(entry, ports)
+    return ports
+
+
+def _collect_published_ports(entry: dict, out: set[int]) -> None:
+    for pub in entry.get('Publishers') or []:
+        port = pub.get('PublishedPort')
+        if isinstance(port, int) and port > 0:
+            out.add(port)
+
+
+def check_ports_available(ports: list[tuple[str, int, str]]) -> None:
+    """Pre-flight check: try to bind each (service, port, host) tuple.
+
+    ``host`` is the interface the rendered compose file binds the publication
+    on (e.g. ``"0.0.0.0"`` for litellm / open-webui, ``"127.0.0.1"`` for vllm
+    services). Binding on ``0.0.0.0`` collides with any listener on any
+    interface for that port (including IPv6 ``[::]`` due to
+    IPv4-mapped-IPv6), which is the most common cause of the
+    ``failed to bind port`` errors users hit at ``compose up`` time.
+
+    Raises ``PortInUseError`` listing every conflicting service so the user
+    sees the full picture in one shot, not one failure at a time.
+    """
+    busy: list[tuple[str, int, str]] = []
+    for service, port, host in ports:
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        try:
+            sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            try:
+                sock.bind((host, port))
+            except OSError as ex:
+                if ex.errno in (errno.EADDRINUSE, errno.EACCES):
+                    busy.append((service, port, host))
+                else:
+                    raise
+        finally:
+            sock.close()
+    if busy:
+        raise PortInUseError(busy)
+
+
+def _cmd(
+    compose_cmd: str, compose_file: Path, env_file: Path, *args: str
+) -> list[str]:
+    return compose_cmd.split() + [
+        '-f',
+        str(compose_file),
+        '--env-file',
+        str(env_file),
+        *args,
+    ]
+
+
+def run(cmd: list[str]) -> None:
+    proc = subprocess.run(cmd)
+    if proc.returncode != 0:
+        raise DockerCommandError(
+            f'Command failed with exit code {proc.returncode}: {" ".join(cmd)}'
+        )
+
+
+def compose_up(
+    compose_cmd: str,
+    compose_file: Path,
+    env_file: Path,
+    *,
+    detach: bool = False,
+    remove_orphans: bool = True,
+    force_recreate: bool = False,
+    services: list[str] | None = None,
+) -> None:
+    args = ['up']
+    if detach:
+        args.append('-d')
+    if remove_orphans:
+        args.append('--remove-orphans')
+    if force_recreate:
+        args.append('--force-recreate')
+    if services:
+        args.extend(services)
+    run(_cmd(compose_cmd, compose_file, env_file, *args))
+
+
+def compose_down(compose_cmd: str, compose_file: Path, env_file: Path) -> None:
+    """Stop and remove services. Never removes named volumes."""
+    run(_cmd(compose_cmd, compose_file, env_file, 'down', '--remove-orphans'))
+
+
+def docker_rm_dirs(dirs: list[Path], docker_cmd: str = 'docker') -> None:
+    """Delete host directories that may be root-owned (written by Docker containers).
+
+    Groups paths by parent directory and removes them from inside a temporary
+    Alpine container so that permission errors from user-space rm are avoided.
+    """
+    from collections import defaultdict
+
+    by_parent: dict[Path, list[str]] = defaultdict(list)
+    for d in dirs:
+        if d.exists():
+            by_parent[d.parent].append(d.name)
+
+    for parent, names in by_parent.items():
+        targets = ' '.join(f'/mnt/{name}' for name in names)
+        cmd = [
+            docker_cmd,
+            'run',
+            '--rm',
+            '-v',
+            f'{parent}:/mnt',
+            'alpine',
+            'sh',
+            '-c',
+            f'rm -rf {targets}',
+        ]
+        run(cmd)
+
+
+def compose_recreate_router(
+    compose_cmd: str,
+    compose_file: Path,
+    env_file: Path,
+    *,
+    detach: bool = True,
+) -> None:
+    """Recreate the LiteLLM router container in place.
+
+    Forces LiteLLM to reload its config from the rendered YAML. Open WebUI is
+    deliberately *not* recreated: it re-fetches ``/v1/models`` from LiteLLM
+    on user actions, so a brief stale-cache window is fine, while
+    force-recreating would log every user out of the chat UI.
+    """
+    args = ['up']
+    if detach:
+        args.append('-d')
+    args.extend(
+        ['--remove-orphans', '--force-recreate', '--no-deps', 'litellm']
+    )
+    run(_cmd(compose_cmd, compose_file, env_file, *args))

infer_stack/env_utils.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import secrets
+from pathlib import Path
+
+
+def parse_env_file(path: Path) -> dict[str, str]:
+    """Return a flat dict of KEY=value pairs from a .env file.
+
+    Comments, blank lines, and malformed lines (no '=') are skipped. Both
+    managed and unknown keys are returned together.
+    """
+    if not path.exists():
+        return {}
+    data: dict[str, str] = {}
+    for line in path.read_text(encoding='utf-8').splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith('#') or '=' not in stripped:
+            continue
+        k, v = stripped.split('=', 1)
+        data[k.strip()] = v
+    return data
+
+
+def ensure_secret(
+    env: dict[str, str], key: str, length: int = 32, prefix: str = ''
+) -> str:
+    """Return ``env[key]`` if present (and matches ``prefix``), else a fresh secret.
+
+    The ``prefix`` requirement is checked against any existing value: if the
+    stored value does not start with ``prefix``, a new secret is generated.
+    This matters for keys whose downstream consumer enforces a format
+    (e.g. LiteLLM rejects auth tokens that don't start with ``sk-``).
+    """
+    value = env.get(key, '').strip()
+    if value and (not prefix or value.startswith(prefix)):
+        return value
+    return prefix + secrets.token_urlsafe(length)
+
+
+def _parse_env_lines(text: str) -> list[tuple[str, str | None, str]]:
+    """Tokenize a .env file into ordered records.
+
+    Each record is a tuple ``(kind, key, raw_line)`` where ``kind`` is one of
+    ``"kv"`` (a syntactically valid ``KEY=VALUE`` line) or ``"other"`` (a blank
+    line, comment, or malformed line). For ``"kv"`` records, ``key`` is the
+    variable name; otherwise ``key`` is None. ``raw_line`` is the verbatim
+    line (without trailing newline) to preserve formatting on rewrite.
+    """
+    records: list[tuple[str, str | None, str]] = []
+    for line in text.splitlines():
+        stripped = line.strip()
+        if not stripped or stripped.startswith('#') or '=' not in stripped:
+            records.append(('other', None, line))
+            continue
+        k, _ = stripped.split('=', 1)
+        key = k.strip()
+        if not key or any(c.isspace() for c in key):
+            records.append(('other', None, line))
+            continue
+        records.append(('kv', key, line))
+    return records
+
+
+def write_env_file(path: Path, values: dict[str, str]) -> None:
+    """Write the .env file non-destructively.
+
+    - Existing unknown ``KEY=value`` pairs are preserved verbatim.
+    - Comments and blank lines are preserved in their original positions.
+    - Keys present in ``values`` are updated in place (preserving line order).
+    - Keys in ``values`` not yet present in the file are appended at the end.
+    """
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    existing_text = path.read_text(encoding='utf-8') if path.exists() else ''
+    records = _parse_env_lines(existing_text)
+
+    seen: set[str] = set()
+    out_lines: list[str] = []
+    for kind, key, raw in records:
+        if kind == 'kv' and key in values:
+            out_lines.append(f'{key}={values[key]}')
+            seen.add(key)
+        else:
+            out_lines.append(raw)
+
+    for key, value in values.items():
+        if key in seen:
+            continue
+        out_lines.append(f'{key}={value}')
+
+    text = '\n'.join(out_lines)
+    if text and not text.endswith('\n'):
+        text += '\n'
+
+    print(f'Write .env to {path}')
+    path.write_text(text, encoding='utf-8')