infer-stack 0.6.0__py3-none-any.whl

infer_stack/__init__.py

@@ -0,0 +1,2 @@
+__all__ = []
+__version__ = '0.6.0'

infer_stack/backends/__init__.py

@@ -0,0 +1,7 @@
+from .compose_renderer import render_compose_artifacts
+from .kubeai_renderer import render_kubeai_artifacts
+
+__all__ = [
+    'render_compose_artifacts',
+    'render_kubeai_artifacts',
+]

infer_stack/backends/compose_renderer.py

@@ -0,0 +1,243 @@
+from __future__ import annotations
+
+import json
+from importlib.resources import files
+from pathlib import Path
+
+import yaml
+from jinja2 import BaseLoader, Environment
+
+from ..config import normalized_output, normalized_state, DEFAULT_PORTS
+from ..diff_prompt import confirm_writes
+from ..env_utils import ensure_secret, parse_env_file, write_env_file
+
+
+def _template(name: str) -> str:
+    return (
+        files('infer_stack')
+        .joinpath(f'templates/{name}')
+        .read_text(encoding='utf-8')
+    )
+
+
+def _compose_quote(value: object) -> str:
+    """Quote scalars for Compose YAML while preserving env interpolation text."""
+
+    return json.dumps('' if value is None else str(value))
+
+
+def _compose_gpus(value: object, indent: int = 4) -> str:
+    """Render the ``gpus:`` service key for either form Compose accepts.
+
+    Compose understands both a scalar (``all``, ``-1``, or a device count) and
+    a structured list of device-request mappings.  Scalars render inline and
+    quoted; lists/dicts are emitted as indented block YAML so the structured
+    "GPU settings" escape hatch round-trips faithfully instead of collapsing to
+    a quoted Python repr.
+    """
+
+    pad = ' ' * indent
+    if value is None or value == '':
+        return ''
+    if isinstance(value, bool):
+        value = 'true' if value else 'false'
+    if isinstance(value, (str, int, float)):
+        scalar = (
+            _compose_quote(value)
+            if isinstance(value, str)
+            else json.dumps(value)
+        )
+        return f'{pad}gpus: {scalar}'
+    dumped = yaml.safe_dump(
+        {'gpus': value}, default_flow_style=False, sort_keys=False
+    ).rstrip('\n')
+    return '\n'.join(pad + line for line in dumped.splitlines())
+
+
+def render_compose_artifacts(
+    lock_data: dict, *, assume_yes: bool = True
+) -> None:
+    """Render component-aware Compose artifacts for the resolved stack."""
+    deployment = dict(lock_data.get('deployment', {}))
+    deployment['state'] = normalized_state(deployment.get('state', {}))
+    deployment['output'] = normalized_output(deployment.get('output'))
+    generated = Path(deployment['output']['generated_dir'])
+    generated.mkdir(parents=True, exist_ok=True)
+    runtime_dir = Path(deployment['state']['runtime'])
+    runtime_dir.mkdir(parents=True, exist_ok=True)
+
+    env_path = generated / '.env'
+    existing = parse_env_file(env_path)
+    env_values: dict[str, str] = {}
+
+    frontends = deployment.get('frontends', {}) or {}
+    gateways = deployment.get('gateways', {}) or {}
+    providers = deployment.get('providers', {}) or {}
+
+    if (frontends.get('open_webui') or {}).get('enabled'):
+        env_values.update(
+            {
+                'OPENWEBUI_POSTGRES_DB': existing.get(
+                    'OPENWEBUI_POSTGRES_DB', 'openwebui'
+                ),
+                'OPENWEBUI_POSTGRES_USER': existing.get(
+                    'OPENWEBUI_POSTGRES_USER', 'openwebui'
+                ),
+                'OPENWEBUI_POSTGRES_PASSWORD': ensure_secret(
+                    existing, 'OPENWEBUI_POSTGRES_PASSWORD'
+                ),
+                'WEBUI_SECRET_KEY': ensure_secret(existing, 'WEBUI_SECRET_KEY'),
+            }
+        )
+        ldap_defaults = (
+            (frontends.get('open_webui') or {}).get('ldap') or {}
+        ).get('env_defaults') or {}
+        for key, default in ldap_defaults.items():
+            env_values[key] = existing.get(key, str(default))
+
+    if (gateways.get('litellm') or {}).get('enabled'):
+        env_values.update(
+            {
+                'LITELLM_POSTGRES_DB': existing.get(
+                    'LITELLM_POSTGRES_DB', 'litellm'
+                ),
+                'LITELLM_POSTGRES_USER': existing.get(
+                    'LITELLM_POSTGRES_USER', 'litellm'
+                ),
+                'LITELLM_POSTGRES_PASSWORD': ensure_secret(
+                    existing, 'LITELLM_POSTGRES_PASSWORD'
+                ),
+                'LITELLM_MASTER_KEY': ensure_secret(
+                    existing, 'LITELLM_MASTER_KEY', prefix='sk-'
+                ),
+            }
+        )
+
+    if (providers.get('vllm') or {}).get('enabled'):
+        env_values.update(
+            {
+                'VLLM_BACKEND_API_KEY': ensure_secret(
+                    existing, 'VLLM_BACKEND_API_KEY'
+                ),
+                'HF_TOKEN': existing.get('HF_TOKEN', ''),
+            }
+        )
+
+    # Ports: expose configured host ports via environment variables so
+    # docker-compose can reference them and we persist them into `.env`.
+    ports = deployment.get('ports', {}) or {}
+
+    # LiteLLM port
+    if (gateways.get('litellm') or {}).get('enabled'):
+        litellm_port = ports.get('litellm') or DEFAULT_PORTS.get(
+            'litellm', 14042
+        )
+        env_values['INFER_STACK_LITELLM_PORT'] = existing.get(
+            'INFER_STACK_LITELLM_PORT', str(litellm_port)
+        )
+
+    # Open WebUI port
+    if (frontends.get('open_webui') or {}).get('enabled'):
+        open_webui_port = ports.get('open_webui') or DEFAULT_PORTS.get(
+            'open_webui', 13000
+        )
+        env_values['INFER_STACK_OPEN_WEBUI_PORT'] = existing.get(
+            'INFER_STACK_OPEN_WEBUI_PORT', str(open_webui_port)
+        )
+
+    # Reverse proxy ports
+    reverse_proxy = frontends.get('reverse_proxy') or {}
+    if reverse_proxy.get('enabled'):
+        http_port = (
+            reverse_proxy.get('http_port')
+            or ports.get('reverse_proxy_http')
+            or DEFAULT_PORTS.get('reverse_proxy_http', 80)
+        )
+        https_port = (
+            reverse_proxy.get('https_port')
+            or ports.get('reverse_proxy_https')
+            or DEFAULT_PORTS.get('reverse_proxy_https', 443)
+        )
+        env_values['INFER_STACK_REVERSE_PROXY_HTTP_PORT'] = existing.get(
+            'INFER_STACK_REVERSE_PROXY_HTTP_PORT', str(http_port)
+        )
+        env_values['INFER_STACK_REVERSE_PROXY_HTTPS_PORT'] = existing.get(
+            'INFER_STACK_REVERSE_PROXY_HTTPS_PORT', str(https_port)
+        )
+
+    # Ollama port (if publish enabled)
+    if (providers.get('ollama') or {}).get('enabled'):
+        # Ollama host_port is resolved in the deployment; fall back to DEFAULT_PORTS
+        ollama_port = (
+            (providers.get('ollama') or {}).get('host_port')
+            or ports.get('ollama')
+            or DEFAULT_PORTS.get('ollama', 11434)
+        )
+        env_values['INFER_STACK_OLLAMA_PORT'] = existing.get(
+            'INFER_STACK_OLLAMA_PORT', str(ollama_port)
+        )
+
+    # vLLM runtimes: enumerate and export per-runtime host ports (index-based)
+    vllm_runtimes = (providers.get('vllm') or {}).get('runtimes', {}) or {}
+    for idx, (name, svc) in enumerate(vllm_runtimes.items()):
+        host_port = svc.get('host_port') or ports.get('vllm') or (18000 + idx)
+        env_name = f'INFER_STACK_VLLM_{idx}_PORT'
+        env_values[env_name] = existing.get(env_name, str(host_port))
+
+    # Preserve unknown/user-supplied keys, but let managed keys above win.
+    for key, value in existing.items():
+        env_values.setdefault(key, value)
+
+    env = Environment(
+        loader=BaseLoader(),
+        autoescape=False,
+        trim_blocks=True,
+        lstrip_blocks=True,
+    )
+    env.filters['compose_quote'] = _compose_quote
+    env.filters['compose_gpus'] = _compose_gpus
+    normalized_lock = dict(lock_data)
+    normalized_lock['deployment'] = deployment
+
+    reverse_proxy = (deployment.get('frontends') or {}).get(
+        'reverse_proxy'
+    ) or {}
+    if reverse_proxy.get('enabled'):
+        if reverse_proxy.get('config_path'):
+            reverse_proxy['nginx_config_path'] = reverse_proxy['config_path']
+        else:
+            reverse_proxy['nginx_config_path'] = str(runtime_dir / 'nginx.conf')
+        deployment['frontends']['reverse_proxy'] = reverse_proxy
+
+    ctx = {'lock': normalized_lock}
+    compose = (
+        env.from_string(_template('docker-compose.yml.j2')).render(**ctx) + '\n'
+    )
+
+    compose_fpath = generated / 'docker-compose.yml'
+    planned: dict[Path, str] = {compose_fpath: compose}
+
+    lite_llm_config_fpath = runtime_dir / 'litellm_config.yaml'
+    if (gateways.get('litellm') or {}).get('enabled'):
+        litellm_cfg = (
+            env.from_string(_template('litellm_config.yaml.j2')).render(**ctx)
+            + '\n'
+        )
+        planned[lite_llm_config_fpath] = litellm_cfg
+
+    if reverse_proxy.get('enabled') and not reverse_proxy.get('config_path'):
+        nginx_ctx = {'rp': reverse_proxy}
+        planned[Path(reverse_proxy['nginx_config_path'])] = (
+            env.from_string(_template('nginx.conf.j2')).render(**nginx_ctx)
+            + '\n'
+        )
+
+    if not confirm_writes(
+        planned, assume_yes=assume_yes, title='Pending compose render'
+    ):
+        raise SystemExit('Aborted by user; no files were written.')
+
+    write_env_file(env_path, env_values)
+    for path, text in planned.items():
+        path.parent.mkdir(parents=True, exist_ok=True)
+        path.write_text(text, encoding='utf-8')

infer_stack/backends/kubeai_renderer.py

@@ -0,0 +1,202 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from ..config import KUBEAI_GENERATED_SUBDIR, normalized_output
+from ..diff_prompt import confirm_writes
+from ..profile_runtime import vllm_args
+
+
+def _resource_profile_values(plan: dict[str, Any]) -> dict[str, Any]:
+    values_doc = plan.get('deployment', {}).get('resource_profiles_values', {})
+    return values_doc or {'resourceProfiles': {}}
+
+
+def _kubeai_resource_profile(service: dict[str, Any]) -> str:
+    profile = str(service.get('resource_profile', ''))
+    if not profile or ':' in profile:
+        return profile
+    gpu_count = max(
+        1,
+        len(service.get('gpu_indices', [])),
+        int(service.get('tensor_parallel_size', 1) or 1),
+    )
+    return f'{profile}:{gpu_count}'
+
+
+def _kubeai_args(service: dict[str, Any]) -> list[str]:
+    kubeai_service = dict(service)
+    kubeai_service['served_model_name'] = service['profile_public_name']
+    return vllm_args(kubeai_service)
+
+
+def _model_doc(service: dict[str, Any]) -> dict[str, Any]:
+    doc = {
+        'apiVersion': 'kubeai.org/v1',
+        'kind': 'Model',
+        'metadata': {
+            'name': service['kubernetes_name'],
+            'annotations': {
+                'infer-stack/profile-name': service['profile_name'],
+                'infer-stack/public-name': service['profile_public_name'],
+                'infer-stack/logical-model-name': service['logical_model_name'],
+                'infer-stack/protocol-mode': service['protocol_mode'],
+            },
+        },
+        'spec': {
+            'features': service.get('features', ['TextGeneration']),
+            'url': service['model_url'],
+            'engine': service.get('engine', 'VLLM'),
+            'resourceProfile': _kubeai_resource_profile(service),
+            'minReplicas': int(service.get('min_replicas', 0)),
+            'maxReplicas': int(service.get('max_replicas', 1)),
+            'args': _kubeai_args(service),
+        },
+    }
+    if service.get('priority_class_name'):
+        doc['spec']['priorityClassName'] = service['priority_class_name']
+    return doc
+
+
+def render_kubeai_artifacts(
+    lock_data: dict, *, assume_yes: bool = True
+) -> None:
+    """Render the KubeAI backend artifacts.
+
+    When ``assume_yes`` is False, all rendered YAML files are diffed against
+    their on-disk versions and the user is prompted via Rich before any file
+    is written.
+    """
+    deployment = lock_data.get('deployment', {})
+    cluster = deployment.get('cluster', {})
+    namespace = cluster.get('namespace', 'kubeai')
+    output_root = Path(
+        normalized_output(deployment.get('output'))['generated_dir']
+    )
+    generated = output_root / KUBEAI_GENERATED_SUBDIR
+    generated.mkdir(parents=True, exist_ok=True)
+
+    namespace_doc = {
+        'apiVersion': 'v1',
+        'kind': 'Namespace',
+        'metadata': {'name': namespace},
+    }
+    namespace_text = yaml.safe_dump(namespace_doc, sort_keys=False)
+
+    values_doc = _resource_profile_values(lock_data)
+    values_text = yaml.safe_dump(values_doc, sort_keys=False)
+
+    model_docs = [
+        _model_doc(service)
+        for service in (
+            deployment.get('providers', {}).get('vllm', {}).get('runtimes', {})
+            or {}
+        ).values()
+    ]
+    model_text = '---\n'.join(
+        yaml.safe_dump(doc, sort_keys=False) for doc in model_docs
+    )
+
+    ingress = cluster.get('ingress', {}) or {}
+    ingress_path = generated / 'ingress.yaml'
+    ingress_text: str | None = None
+    if ingress.get('enabled'):
+        path_prefix = ingress.get('path_prefix', '/') or '/'
+        ingress_doc: dict[str, Any] = {
+            'apiVersion': 'networking.k8s.io/v1',
+            'kind': 'Ingress',
+            'metadata': {
+                'name': cluster.get('service_name', 'kubeai'),
+                'namespace': namespace,
+            },
+            'spec': {
+                'ingressClassName': ingress.get('class_name', 'traefik'),
+                'rules': [
+                    {
+                        'http': {
+                            'paths': [
+                                {
+                                    'path': path_prefix,
+                                    'pathType': 'Prefix',
+                                    'backend': {
+                                        'service': {
+                                            'name': cluster.get(
+                                                'service_name', 'kubeai'
+                                            ),
+                                            'port': {'number': 80},
+                                        }
+                                    },
+                                }
+                            ]
+                        }
+                    }
+                ],
+            },
+        }
+        if ingress.get('host'):
+            ingress_doc['spec']['rules'][0]['host'] = ingress['host']
+        if ingress.get('tls_secret_name') and ingress.get('host'):
+            ingress_doc['spec']['tls'] = [
+                {
+                    'hosts': [ingress['host']],
+                    'secretName': ingress['tls_secret_name'],
+                }
+            ]
+        ingress_text = yaml.safe_dump(ingress_doc, sort_keys=False)
+
+    readme = f"""# Generated KubeAI artifacts
+
+Namespace: `{namespace}`
+Release: `{cluster.get('kubeai_release_name', 'kubeai')}`
+Chart: `{cluster.get('kubeai_chart', 'kubeai/kubeai')}`
+
+Files:
+- `namespace.yaml`: namespace to apply before the chart and models
+- `kubeai-values.yaml`: custom resource profiles for the KubeAI chart
+- `models.yaml`: KubeAI `Model` objects derived intentionally from the selected serving profile(s)
+- `ingress.yaml`: optional ingress for one stable hostname
+
+Typical flow:
+
+```bash
+kubectl apply -f {generated}/namespace.yaml
+helm repo add kubeai https://www.kubeai.org --force-update
+helm repo update
+helm upgrade --install {cluster.get('kubeai_release_name', 'kubeai')} {cluster.get('kubeai_chart', 'kubeai/kubeai')} \
+  -n {namespace} --create-namespace \
+  -f {generated}/kubeai-values.yaml \
+  --wait
+kubectl apply -f {generated}/models.yaml
+```
+"""
+
+    namespace_path = generated / 'namespace.yaml'
+    values_path = generated / 'kubeai-values.yaml'
+    models_path = generated / 'models.yaml'
+    readme_path = generated / 'README.md'
+
+    planned: dict[Path, str] = {
+        namespace_path: namespace_text,
+        values_path: values_text,
+        models_path: model_text,
+        readme_path: readme,
+    }
+    if ingress_text is not None:
+        planned[ingress_path] = ingress_text
+
+    if not confirm_writes(
+        planned, assume_yes=assume_yes, title='Pending KubeAI render'
+    ):
+        raise SystemExit('Aborted by user; no files were written.')
+
+    namespace_path.write_text(namespace_text, encoding='utf-8')
+    values_path.write_text(values_text, encoding='utf-8')
+    models_path.write_text(model_text, encoding='utf-8')
+    readme_path.write_text(readme, encoding='utf-8')
+    if ingress_text is not None:
+        ingress_path.write_text(ingress_text, encoding='utf-8')
+    elif ingress_path.exists():
+        ingress_path.unlink()

infer_stack/benchmark.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import requests
+
+
+def run_benchmark(
+    base_url: str, api_key: str, model: str, prompts: list[str]
+) -> dict[str, Any]:
+    headers = {'Content-Type': 'application/json'}
+    if api_key:
+        headers['Authorization'] = f'Bearer {api_key}'
+    results = []
+    for prompt in prompts:
+        payload = {
+            'model': model,
+            'messages': [{'role': 'user', 'content': prompt}],
+            'max_tokens': 128,
+        }
+        t0 = time.time()
+        resp = requests.post(
+            f'{base_url.rstrip("/")}/chat/completions',
+            headers=headers,
+            json=payload,
+            timeout=300,
+        )
+        dt = time.time() - t0
+        results.append(
+            {
+                'prompt': prompt,
+                'status_code': resp.status_code,
+                'elapsed_s': dt,
+                'ok': resp.ok,
+            }
+        )
+    return {'results': results}