infer-stack 0.6.0__py3-none-any.whl

infer_stack/experimental/stress_test_long_context.py

@@ -0,0 +1,397 @@
+#!/usr/bin/env python3
+"""
+Long-context stress test for a vLLM/OpenAI-compatible endpoint.
+
+What it does:
+1. Generates a long synthetic corpus with random "needle" facts embedded throughout.
+2. Uses the Qwen tokenizer chat template to size the prompt near a target token budget.
+3. Sends retrieval and arithmetic questions to the model with requests.
+4. Checks the returned answers and prints PASS/FAIL.
+
+Typical usage:
+
+  python stress_test_long_context.py \
+    --base-url http://127.0.0.1:18000/v1 \
+    --model qwen3.5-122b-a10b-fp8-262k \
+    --max-context 262144 \
+    --reserved-output 512
+
+If your API key is in generated/.env as VLLM_BACKEND_API_KEY, it will be auto-detected.
+Otherwise pass --api-key explicitly.
+
+Requirements:
+  pip install requests transformers
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import random
+import re
+import string
+import sys
+import uuid
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import requests
+from transformers import AutoTokenizer
+
+
+@dataclass
+class Fact:
+    marker: str
+    secret_value: int
+    secret_city: str
+    secret_word: str
+
+
+def load_api_key(cli_value: str | None) -> str | None:
+    if cli_value:
+        return cli_value
+    env_key = os.environ.get('VLLM_BACKEND_API_KEY') or os.environ.get(
+        'OPENAI_API_KEY'
+    )
+    if env_key:
+        return env_key
+    env_path = Path('generated/.env')
+    if env_path.exists():
+        text = env_path.read_text(encoding='utf-8')
+        for line in text.splitlines():
+            if line.startswith('VLLM_BACKEND_API_KEY='):
+                return line.split('=', 1)[1].strip()
+    return None
+
+
+def strip_think_tags(text: str) -> str:
+    # Remove <think>...</think> blocks if present.
+    return re.sub(r'<think>.*?</think>', '', text, flags=re.S).strip()
+
+
+def normalize_answer(text: str) -> str:
+    text = strip_think_tags(text)
+    text = text.strip()
+    # Remove wrapping code fences or quotes.
+    text = text.strip('`').strip()
+    text = text.strip('"').strip("'").strip()
+    return text
+
+
+def random_word(min_len: int = 4, max_len: int = 10) -> str:
+    return ''.join(
+        random.choices(
+            string.ascii_lowercase, k=random.randint(min_len, max_len)
+        )
+    )
+
+
+def filler_paragraph(num_words: int = 120) -> str:
+    return ' '.join(random_word() for _ in range(num_words))
+
+
+def build_facts(n: int) -> list[Fact]:
+    cities = [
+        'Lima',
+        'Oslo',
+        'Accra',
+        'Perth',
+        'Quito',
+        'Hobart',
+        'Dakar',
+        'Tallinn',
+        'Salta',
+        'Cork',
+        'Uppsala',
+        'Medellin',
+        'LuangPrabang',
+    ]
+    facts: list[Fact] = []
+    for _ in range(n):
+        marker = str(uuid.uuid4())[:8].upper()
+        facts.append(
+            Fact(
+                marker=marker,
+                secret_value=random.randint(100000, 999999),
+                secret_city=random.choice(cities),
+                secret_word=random_word(7, 12).upper(),
+            )
+        )
+    return facts
+
+
+def build_corpus(target_sections: int, facts: list[Fact]) -> str:
+    """
+    Build a long boring corpus with fact blocks inserted at spread-out positions.
+    """
+    sections = [
+        f'SECTION {i}\n{filler_paragraph()}\n' for i in range(target_sections)
+    ]
+    # Spread facts through the document.
+    positions = [
+        int(i * target_sections / max(len(facts), 1)) for i in range(len(facts))
+    ]
+    positions = [min(target_sections - 1, max(0, p)) for p in positions]
+
+    for pos, fact in zip(positions, facts):
+        sections[pos] += (
+            f'\nMARKER {fact.marker}\n'
+            f'SECRET_VALUE {fact.secret_value}\n'
+            f'SECRET_CITY {fact.secret_city}\n'
+            f'SECRET_WORD {fact.secret_word}\n'
+            f'END_MARKER {fact.marker}\n'
+        )
+
+    return '\n'.join(sections)
+
+
+def chat_token_count(tokenizer: Any, user_content: str) -> int:
+    ids = tokenizer.apply_chat_template(
+        [{'role': 'user', 'content': user_content}],
+        tokenize=True,
+        add_generation_prompt=True,
+    )
+    return len(ids)
+
+
+def fit_corpus_to_budget(
+    tokenizer: Any,
+    base_instruction: str,
+    corpus_template: str,
+    target_prompt_tokens: int,
+) -> str:
+    """
+    Binary-search the number of repeated filler sections until the whole prompt
+    fits just under target_prompt_tokens.
+    """
+    lo, hi = 1, 4000
+    best = corpus_template
+
+    while lo <= hi:
+        mid = (lo + hi) // 2
+        test_corpus = corpus_template.format(
+            extra_sections=('\nEXTRA\n' + filler_paragraph() + '\n') * mid
+        )
+        user_content = (
+            f'{base_instruction}\n\nBEGIN CORPUS\n{test_corpus}\nEND CORPUS'
+        )
+        n = chat_token_count(tokenizer, user_content)
+        if n <= target_prompt_tokens:
+            best = test_corpus
+            lo = mid + 1
+        else:
+            hi = mid - 1
+
+    return best
+
+
+def make_questions(facts: list[Fact]) -> list[dict[str, Any]]:
+    """
+    Build a small suite of retrieval and aggregation questions.
+    """
+    qs: list[dict[str, Any]] = []
+
+    # Exact value retrieval
+    f0 = facts[0]
+    qs.append(
+        {
+            'kind': 'value',
+            'question': (
+                f'Using only the provided corpus, what is the SECRET_VALUE for marker {f0.marker}? '
+                'Reply with only the integer.'
+            ),
+            'expected': str(f0.secret_value),
+        }
+    )
+
+    # Exact city retrieval
+    f1 = facts[len(facts) // 2]
+    qs.append(
+        {
+            'kind': 'city',
+            'question': (
+                f'Using only the provided corpus, what is the SECRET_CITY for marker {f1.marker}? '
+                'Reply with only the city name.'
+            ),
+            'expected': f1.secret_city,
+        }
+    )
+
+    # Exact word retrieval from later in the corpus
+    f2 = facts[-1]
+    qs.append(
+        {
+            'kind': 'word',
+            'question': (
+                f'Using only the provided corpus, what is the SECRET_WORD for marker {f2.marker}? '
+                'Reply with only the word.'
+            ),
+            'expected': f2.secret_word,
+        }
+    )
+
+    # Arithmetic over multiple markers
+    picks = [facts[1], facts[len(facts) // 3], facts[-2]]
+    total = sum(f.secret_value for f in picks)
+    markers = ', '.join(f.marker for f in picks)
+    qs.append(
+        {
+            'kind': 'sum',
+            'question': (
+                f'Using only the provided corpus, add the SECRET_VALUE values for markers {markers}. '
+                'Reply with only the integer sum.'
+            ),
+            'expected': str(total),
+        }
+    )
+
+    return qs
+
+
+def call_chat_completion(
+    session: requests.Session,
+    base_url: str,
+    api_key: str | None,
+    model: str,
+    prompt: str,
+    question: str,
+    max_tokens: int,
+    temperature: float,
+) -> str:
+    url = base_url.rstrip('/') + '/chat/completions'
+    payload: dict[str, Any] = {
+        'model': model,
+        'messages': [
+            {
+                'role': 'system',
+                'content': (
+                    'Answer using only the provided corpus. '
+                    'If the requested value is not present, reply with NOT_FOUND.'
+                ),
+            },
+            {
+                'role': 'user',
+                'content': f'{prompt}\n\nQUESTION:\n{question}',
+            },
+        ],
+        'max_tokens': max_tokens,
+        'temperature': temperature,
+    }
+    headers = {'Content-Type': 'application/json'}
+    if api_key:
+        headers['Authorization'] = f'Bearer {api_key}'
+
+    response = session.post(
+        url, headers=headers, data=json.dumps(payload), timeout=3600
+    )
+    response.raise_for_status()
+    data = response.json()
+    return data['choices'][0]['message']['content']
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--base-url', default='http://127.0.0.1:18000/v1')
+    parser.add_argument('--model', required=True)
+    parser.add_argument(
+        '--tokenizer', default=None, help='Defaults to --model if omitted.'
+    )
+    parser.add_argument('--api-key', default=None)
+    parser.add_argument('--max-context', type=int, required=True)
+    parser.add_argument('--reserved-output', type=int, default=512)
+    parser.add_argument('--temperature', type=float, default=0.0)
+    parser.add_argument('--num-facts', type=int, default=10)
+    parser.add_argument('--seed', type=int, default=0)
+    parser.add_argument('--save-corpus', default='haystack_generated.txt')
+    parser.add_argument('--save-answers', default='answer_key_generated.json')
+    args = parser.parse_args()
+
+    random.seed(args.seed)
+    api_key = load_api_key(args.api_key)
+    tokenizer_name = args.tokenizer or args.model
+
+    print(f'Loading tokenizer: {tokenizer_name}', file=sys.stderr)
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+
+    facts = build_facts(args.num_facts)
+
+    # Base corpus with distributed facts. We will append extra filler during sizing.
+    corpus_without_extra = build_corpus(target_sections=1800, facts=facts)
+    corpus_template = corpus_without_extra + '\n{extra_sections}'
+
+    base_instruction = (
+        'You will be given a long corpus delimited by BEGIN CORPUS and END CORPUS. '
+        'Use only that corpus to answer the later question.'
+    )
+
+    # Keep a little slack so the specific question still fits.
+    target_prompt_tokens = args.max_context - args.reserved_output - 1024
+    fitted_corpus = fit_corpus_to_budget(
+        tokenizer=tokenizer,
+        base_instruction=base_instruction,
+        corpus_template=corpus_template,
+        target_prompt_tokens=target_prompt_tokens,
+    )
+
+    # Save generated materials for inspection.
+    Path(args.save_corpus).write_text(fitted_corpus, encoding='utf-8')
+    Path(args.save_answers).write_text(
+        json.dumps([fact.__dict__ for fact in facts], indent=2),
+        encoding='utf-8',
+    )
+
+    prompt = f'{base_instruction}\n\nBEGIN CORPUS\n{fitted_corpus}\nEND CORPUS'
+    prompt_tokens = chat_token_count(tokenizer, prompt)
+    print(f'Prompt tokens: {prompt_tokens}')
+    print(f'Reserved output tokens: {args.reserved_output}')
+    print(f'Total budget target: {args.max_context}')
+    print(f'Saved corpus to: {args.save_corpus}')
+    print(f'Saved answer key to: {args.save_answers}')
+
+    questions = make_questions(facts)
+    session = requests.Session()
+
+    passed = 0
+    for i, q in enumerate(questions, start=1):
+        print(f'\n=== Test {i}: {q["kind"]} ===')
+        print('Question:', q['question'])
+        try:
+            raw = call_chat_completion(
+                session=session,
+                base_url=args.base_url,
+                api_key=api_key,
+                model=args.model,
+                prompt=prompt,
+                question=q['question'],
+                max_tokens=args.reserved_output,
+                temperature=args.temperature,
+            )
+        except Exception as ex:
+            print(f'FAIL: request error: {ex}')
+            continue
+
+        got = normalize_answer(raw)
+        expected = str(q['expected']).strip()
+
+        # Gentle normalization for exact-answer tasks.
+        if q['kind'] in {'value', 'sum'}:
+            m = re.search(r'-?\d+', got)
+            got_cmp = m.group(0) if m else got
+        else:
+            got_cmp = got.splitlines()[0].strip()
+
+        ok = got_cmp == expected
+        print('Expected:', expected)
+        print('Got:', got_cmp)
+        print('Raw output:', raw[:1000].replace('\n', '\\n'))
+        print('Result:', 'PASS' if ok else 'FAIL')
+        passed += int(ok)
+
+    print(f'\nPassed {passed}/{len(questions)} tests.')
+    return 0 if passed == len(questions) else 1
+
+
+if __name__ == '__main__':
+    raise SystemExit(main())

infer_stack/hardware.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import csv
+import subprocess
+from typing import Any
+
+
+def _run(cmd: list[str]) -> str:
+    try:
+        out = subprocess.check_output(cmd, text=True, stderr=subprocess.DEVNULL)
+    except Exception:
+        return ''
+    return out
+
+
+def simulate_inventory(spec: str) -> dict[str, Any]:
+    """Build a fake inventory from a spec string like '4x96' (4 GPUs × 96 GiB)."""
+    try:
+        count_str, gib_str = spec.lower().split('x', 1)
+        gpu_count = int(count_str)
+        memory_gib = float(gib_str)
+    except (ValueError, AttributeError):
+        raise ValueError(
+            f'Invalid --simulate-hardware spec {spec!r}. Expected format: NxM (e.g. 4x96, 2x80).'
+        )
+    memory_mib = int(memory_gib * 1024)
+    gpus = [
+        {
+            'index': i,
+            'uuid': f'GPU-simulated-{i:04d}',
+            'name': f'Simulated GPU ({memory_gib:.0f}GiB)',
+            'memory_mib': memory_mib,
+            'memory_gib': memory_gib,
+            'display_active': False,
+        }
+        for i in range(gpu_count)
+    ]
+    return {'gpu_count': gpu_count, 'gpus': gpus}
+
+
+def detect_inventory() -> dict[str, Any]:
+    query = _run(
+        [
+            'nvidia-smi',
+            '--query-gpu=index,uuid,name,memory.total,display_active',
+            '--format=csv,noheader,nounits',
+        ]
+    )
+    gpus: list[dict[str, Any]] = []
+    if query:
+        reader = csv.reader(line for line in query.splitlines() if line.strip())
+        for row in reader:
+            if len(row) < 5:
+                continue
+            idx, uuid, name, mem, display_active = [x.strip() for x in row[:5]]
+            gpus.append(
+                {
+                    'index': int(idx),
+                    'uuid': uuid,
+                    'name': name,
+                    'memory_mib': int(float(mem)),
+                    'memory_gib': round(int(float(mem)) / 1024, 2),
+                    'display_active': display_active.lower()
+                    in {'enabled', 'active', 'on', 'true'},
+                }
+            )
+    return {
+        'gpu_count': len(gpus),
+        'gpus': gpus,
+    }

infer_stack/kubeai_ops.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import os
+import subprocess
+from pathlib import Path
+
+from .config import KUBEAI_GENERATED_SUBDIR, normalized_output
+
+
+class CommandError(RuntimeError):
+    pass
+
+
+def run(cmd: list[str]) -> None:
+    proc = subprocess.run(cmd)
+    if proc.returncode != 0:
+        raise CommandError(
+            f'Command failed with exit code {proc.returncode}: {" ".join(cmd)}'
+        )
+
+
+def deploy_rendered_artifacts(deployment: dict) -> None:
+    cluster = deployment.get('cluster', {})
+    namespace = cluster.get('namespace', 'kubeai')
+    release_name = cluster.get('kubeai_release_name', 'kubeai')
+    chart = cluster.get('kubeai_chart', 'kubeai/kubeai')
+    output_root = Path(
+        normalized_output(deployment.get('output'))['generated_dir']
+    )
+    generated = output_root / KUBEAI_GENERATED_SUBDIR
+    values_file = generated / 'kubeai-values.yaml'
+    namespace_file = generated / 'namespace.yaml'
+    models_file = generated / 'models.yaml'
+    ingress_file = generated / 'ingress.yaml'
+
+    run(['kubectl', 'apply', '-f', str(namespace_file)])
+    run(
+        [
+            'helm',
+            'repo',
+            'add',
+            'kubeai',
+            'https://www.kubeai.org',
+            '--force-update',
+        ]
+    )
+    run(['helm', 'repo', 'update'])
+
+    helm_cmd = [
+        'helm',
+        'upgrade',
+        '--install',
+        release_name,
+        chart,
+        '-n',
+        namespace,
+        '--create-namespace',
+        '-f',
+        str(values_file),
+        '--wait',
+    ]
+    hf_token = os.environ.get('HF_TOKEN', '').strip()
+    if hf_token:
+        helm_cmd.extend(['--set', f'secrets.huggingface.token={hf_token}'])
+    run(helm_cmd)
+
+    run(['kubectl', 'apply', '-f', str(models_file)])
+    if ingress_file.exists():
+        run(['kubectl', 'apply', '-f', str(ingress_file)])
+
+
+def print_status(namespace: str) -> None:
+    run(['kubectl', '-n', namespace, 'get', 'pods'])
+    run(['kubectl', '-n', namespace, 'get', 'svc'])
+    run(['kubectl', '-n', namespace, 'get', 'ingress'])
+    run(['kubectl', '-n', namespace, 'get', 'models'])

infer_stack/paths.py

@@ -0,0 +1,87 @@
+"""CWD-independent locations for the infer-stack config and runtime data.
+
+The CLI used to anchor every path on ``Path.cwd()``, which meant invoking
+``infer-stack`` from a different directory silently changed where config
+was read from, where rendered artifacts landed, and where bind-mount
+state lived. This module replaces that with two stable roots:
+
+* ``config_root()`` — where ``config.yaml`` / ``models.yaml`` /
+  ``kubeai-values.local.yaml`` live. Defaults to
+  ``ub.Path.appdir('infer_stack', type='config')`` (``~/.config/infer_stack``
+  on Linux, respecting ``XDG_CONFIG_HOME``).
+* ``data_root()`` — where ``generated/`` (rendered artifacts) and
+  ``state/`` (hf-cache, postgres volumes, runtime bind-mounts) default
+  to. ``ub.Path.appdir('infer_stack', type='data')``
+  (``~/.local/share/infer_stack`` on Linux, respecting
+  ``XDG_DATA_HOME``). Uses ``data`` and not ``cache`` because the stack
+  hosts persistent state — postgres databases, Open WebUI chat history,
+  and user accounts — that would be silently lost if treated as
+  regenerable cache by a system cleanup tool.
+
+Both can be overridden by env vars (``INFER_STACK_CONFIG_DIR`` /
+``INFER_STACK_DATA_DIR``) or by the CLI flags ``--config-dir`` /
+``--data-dir``. The CLI flags translate into process-wide overrides via
+``set_config_root`` / ``set_data_root``.
+
+``--data-dir`` is the single knob for "put everything I generate in one
+place": at ``setup`` time it is baked into the absolute ``state.*`` and
+``output.generated_dir`` paths written to ``config.yaml``. For a bespoke
+split layout, edit those fields in ``config.yaml`` directly.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import ubelt as ub
+
+
+CONFIG_DIR_ENV = 'INFER_STACK_CONFIG_DIR'
+DATA_DIR_ENV = 'INFER_STACK_DATA_DIR'
+
+_config_root_override: Path | None = None
+_data_root_override: Path | None = None
+
+
+def _default_config_root() -> Path:
+    return Path(ub.Path.appdir('infer_stack', type='config'))
+
+
+def _default_data_root() -> Path:
+    return Path(ub.Path.appdir('infer_stack', type='data'))
+
+
+def config_root() -> Path:
+    if _config_root_override is not None:
+        return _config_root_override
+    env = os.environ.get(CONFIG_DIR_ENV)
+    if env:
+        return Path(env).expanduser()
+    return _default_config_root()
+
+
+def data_root() -> Path:
+    if _data_root_override is not None:
+        return _data_root_override
+    env = os.environ.get(DATA_DIR_ENV)
+    if env:
+        return Path(env).expanduser()
+    return _default_data_root()
+
+
+def set_config_root(path: Path | str | None) -> None:
+    """Override ``config_root()`` for the lifetime of this process.
+
+    Pass ``None`` to clear the override and fall back to env var / default.
+    """
+    global _config_root_override
+    _config_root_override = (
+        Path(path).expanduser() if path is not None else None
+    )
+
+
+def set_data_root(path: Path | str | None) -> None:
+    """Override ``data_root()`` for the lifetime of this process."""
+    global _data_root_override
+    _data_root_override = Path(path).expanduser() if path is not None else None

infer_stack/profile_runtime.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+def vllm_args(service: dict[str, Any]) -> list[str]:
+    args = [
+        f'--served-model-name={service["served_model_name"]}',
+        f'--tensor-parallel-size={service["tensor_parallel_size"]}',
+        f'--data-parallel-size={service["data_parallel_size"]}',
+        f'--max-model-len={service["max_model_len"]}',
+        f'--gpu-memory-utilization={service["gpu_memory_utilization"]}',
+        f'--max-num-batched-tokens={service["max_num_batched_tokens"]}',
+        f'--max-num-seqs={service["max_num_seqs"]}',
+        '--disable-log-requests',
+    ]
+    if service.get('enable_prefix_caching'):
+        args.append('--enable-prefix-caching')
+    if service.get('enable_auto_tool_choice'):
+        args.append('--enable-auto-tool-choice')
+        if service.get('tool_call_parser'):
+            args.append(f'--tool-call-parser={service["tool_call_parser"]}')
+    args.extend(service.get('extra_args', []))
+    return args
+
+
+def default_base_url(
+    deployment: dict[str, Any], *, explicit: str | None = None
+) -> str:
+    if explicit:
+        return explicit.rstrip('/')
+    access = deployment.get('access', {}).get('default') or {}
+    if access.get('base_url'):
+        return str(access['base_url']).rstrip('/')
+    backend = deployment.get('backend', 'compose')
+    if backend == 'kubeai':
+        ingress = deployment.get('cluster', {}).get('ingress', {}) or {}
+        host = ingress.get('host', '')
+        if ingress.get('enabled') and host:
+            return f'http://{host}/openai/v1'
+        return 'http://127.0.0.1:8000/openai/v1'
+    if (deployment.get('gateways', {}).get('litellm') or {}).get('enabled'):
+        return f'http://127.0.0.1:{deployment.get("ports", {}).get("litellm", 14042)}/v1'
+    if (deployment.get('providers', {}).get('ollama') or {}).get('enabled'):
+        return f'http://127.0.0.1:{deployment.get("ports", {}).get("ollama", 11434)}'
+    return f'http://127.0.0.1:{deployment.get("ports", {}).get("litellm", 14042)}/v1'

infer_stack/renderer.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+from .backends import render_compose_artifacts, render_kubeai_artifacts
+
+
+def render_from_lock(lock_data: dict, *, assume_yes: bool = True) -> None:
+    """Render backend artifacts.
+
+    ``assume_yes`` defaults to True so programmatic callers and tests
+    are unaffected. CLI entry points pass ``assume_yes=False`` to surface
+    the per-file diff confirmation prompt.
+    """
+    backend = str(
+        lock_data.get('deployment', {}).get('backend', 'compose')
+    ).lower()
+    if backend == 'kubeai':
+        render_kubeai_artifacts(lock_data, assume_yes=assume_yes)
+        return
+    render_compose_artifacts(lock_data, assume_yes=assume_yes)