polyserve 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (109) hide show
  1. polyserve-0.1.0/LICENSE +21 -0
  2. polyserve-0.1.0/PKG-INFO +87 -0
  3. polyserve-0.1.0/README.md +49 -0
  4. polyserve-0.1.0/polyserve/__init__.py +6 -0
  5. polyserve-0.1.0/polyserve/backends/__init__.py +36 -0
  6. polyserve-0.1.0/polyserve/backends/base.py +314 -0
  7. polyserve-0.1.0/polyserve/backends/llamacpp.py +298 -0
  8. polyserve-0.1.0/polyserve/backends/sglang.py +151 -0
  9. polyserve-0.1.0/polyserve/backends/vllm.py +285 -0
  10. polyserve-0.1.0/polyserve/backends/vllm_cpu.py +84 -0
  11. polyserve-0.1.0/polyserve/bench/__init__.py +19 -0
  12. polyserve-0.1.0/polyserve/bench/ablation.py +326 -0
  13. polyserve-0.1.0/polyserve/bench/baselines.py +82 -0
  14. polyserve-0.1.0/polyserve/bench/compare.py +358 -0
  15. polyserve-0.1.0/polyserve/bench/references.py +159 -0
  16. polyserve-0.1.0/polyserve/bench/report.py +305 -0
  17. polyserve-0.1.0/polyserve/cache.py +108 -0
  18. polyserve-0.1.0/polyserve/calibrate/__init__.py +7 -0
  19. polyserve-0.1.0/polyserve/calibrate/datasets.py +129 -0
  20. polyserve-0.1.0/polyserve/calibrate/measure.py +593 -0
  21. polyserve-0.1.0/polyserve/calibrate/objectives.py +287 -0
  22. polyserve-0.1.0/polyserve/calibrate/search.py +914 -0
  23. polyserve-0.1.0/polyserve/calibrate/tail.py +30 -0
  24. polyserve-0.1.0/polyserve/calibrate/tokens.py +82 -0
  25. polyserve-0.1.0/polyserve/calibrate/workload.py +305 -0
  26. polyserve-0.1.0/polyserve/cli.py +873 -0
  27. polyserve-0.1.0/polyserve/disagg.py +714 -0
  28. polyserve-0.1.0/polyserve/gguf.py +239 -0
  29. polyserve-0.1.0/polyserve/hardware.py +483 -0
  30. polyserve-0.1.0/polyserve/hfconfig.py +148 -0
  31. polyserve-0.1.0/polyserve/layout.py +387 -0
  32. polyserve-0.1.0/polyserve/memcal.py +322 -0
  33. polyserve-0.1.0/polyserve/memlog.py +213 -0
  34. polyserve-0.1.0/polyserve/memory.py +110 -0
  35. polyserve-0.1.0/polyserve/models.py +336 -0
  36. polyserve-0.1.0/polyserve/pipeline.py +548 -0
  37. polyserve-0.1.0/polyserve/power.py +366 -0
  38. polyserve-0.1.0/polyserve/predict.py +483 -0
  39. polyserve-0.1.0/polyserve/quantized.py +169 -0
  40. polyserve-0.1.0/polyserve/router.py +41 -0
  41. polyserve-0.1.0/polyserve/selector.py +75 -0
  42. polyserve-0.1.0/polyserve/serve/__init__.py +4 -0
  43. polyserve-0.1.0/polyserve/serve/proxy.py +115 -0
  44. polyserve-0.1.0/polyserve/serve/supervisor.py +155 -0
  45. polyserve-0.1.0/polyserve/speculative.py +127 -0
  46. polyserve-0.1.0/polyserve.egg-info/PKG-INFO +87 -0
  47. polyserve-0.1.0/polyserve.egg-info/SOURCES.txt +107 -0
  48. polyserve-0.1.0/polyserve.egg-info/dependency_links.txt +1 -0
  49. polyserve-0.1.0/polyserve.egg-info/entry_points.txt +2 -0
  50. polyserve-0.1.0/polyserve.egg-info/requires.txt +16 -0
  51. polyserve-0.1.0/polyserve.egg-info/top_level.txt +1 -0
  52. polyserve-0.1.0/pyproject.toml +67 -0
  53. polyserve-0.1.0/setup.cfg +4 -0
  54. polyserve-0.1.0/tests/test_ablation.py +115 -0
  55. polyserve-0.1.0/tests/test_ablation_shrink.py +22 -0
  56. polyserve-0.1.0/tests/test_all_levels.py +31 -0
  57. polyserve-0.1.0/tests/test_attention_check.py +31 -0
  58. polyserve-0.1.0/tests/test_baselines.py +81 -0
  59. polyserve-0.1.0/tests/test_break_even.py +54 -0
  60. polyserve-0.1.0/tests/test_budget.py +65 -0
  61. polyserve-0.1.0/tests/test_cache_and_pipeline.py +117 -0
  62. polyserve-0.1.0/tests/test_cli.py +44 -0
  63. polyserve-0.1.0/tests/test_combinations.py +157 -0
  64. polyserve-0.1.0/tests/test_compare.py +117 -0
  65. polyserve-0.1.0/tests/test_compare_heldout.py +26 -0
  66. polyserve-0.1.0/tests/test_compare_repeats.py +76 -0
  67. polyserve-0.1.0/tests/test_confirm.py +70 -0
  68. polyserve-0.1.0/tests/test_contenders.py +114 -0
  69. polyserve-0.1.0/tests/test_cpu_quota.py +32 -0
  70. polyserve-0.1.0/tests/test_dataset_workloads.py +85 -0
  71. polyserve-0.1.0/tests/test_early_stop_check.py +40 -0
  72. polyserve-0.1.0/tests/test_explore.py +79 -0
  73. polyserve-0.1.0/tests/test_fit_per_machine.py +62 -0
  74. polyserve-0.1.0/tests/test_gguf_and_backends.py +129 -0
  75. polyserve-0.1.0/tests/test_hardware_and_config.py +75 -0
  76. polyserve-0.1.0/tests/test_hub_search.py +41 -0
  77. polyserve-0.1.0/tests/test_int8_kv.py +53 -0
  78. polyserve-0.1.0/tests/test_load_pick_order.py +27 -0
  79. polyserve-0.1.0/tests/test_make_dolly_prompts.py +29 -0
  80. polyserve-0.1.0/tests/test_make_oasst_prompts.py +28 -0
  81. polyserve-0.1.0/tests/test_measure_and_proxy.py +226 -0
  82. polyserve-0.1.0/tests/test_measure_levels.py +102 -0
  83. polyserve-0.1.0/tests/test_memcal.py +174 -0
  84. polyserve-0.1.0/tests/test_memory.py +88 -0
  85. polyserve-0.1.0/tests/test_nvlink.py +67 -0
  86. polyserve-0.1.0/tests/test_objectives_and_search.py +219 -0
  87. polyserve-0.1.0/tests/test_oom_retry.py +60 -0
  88. polyserve-0.1.0/tests/test_phases.py +418 -0
  89. polyserve-0.1.0/tests/test_power.py +428 -0
  90. polyserve-0.1.0/tests/test_predict.py +150 -0
  91. polyserve-0.1.0/tests/test_process_group.py +33 -0
  92. polyserve-0.1.0/tests/test_progress_line.py +42 -0
  93. polyserve-0.1.0/tests/test_quant_default.py +58 -0
  94. polyserve-0.1.0/tests/test_report.py +112 -0
  95. polyserve-0.1.0/tests/test_rescore_ttft.py +43 -0
  96. polyserve-0.1.0/tests/test_search_budget.py +83 -0
  97. polyserve-0.1.0/tests/test_selector.py +48 -0
  98. polyserve-0.1.0/tests/test_sglang_python.py +31 -0
  99. polyserve-0.1.0/tests/test_spec_methods.py +50 -0
  100. polyserve-0.1.0/tests/test_strategies.py +631 -0
  101. polyserve-0.1.0/tests/test_strategy_report.py +98 -0
  102. polyserve-0.1.0/tests/test_supervisor.py +123 -0
  103. polyserve-0.1.0/tests/test_tail.py +66 -0
  104. polyserve-0.1.0/tests/test_task_quality.py +29 -0
  105. polyserve-0.1.0/tests/test_ttft_percentile.py +52 -0
  106. polyserve-0.1.0/tests/test_vllm_cli.py +39 -0
  107. polyserve-0.1.0/tests/test_w8a8.py +87 -0
  108. polyserve-0.1.0/tests/test_workload_file.py +70 -0
  109. polyserve-0.1.0/tests/test_workloads.py +96 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Aagam Sachin Bothara
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,87 @@
1
+ Metadata-Version: 2.4
2
+ Name: polyserve
3
+ Version: 0.1.0
4
+ Summary: Autotuner for LLM serving: measures vLLM, SGLang and llama.cpp settings on your GPU and your traffic, then serves the fastest
5
+ Author-email: Aagam Sachin Bothara <mweeb19@gmail.com>
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/Aagam-Bothara/polyserve
8
+ Project-URL: Documentation, https://github.com/Aagam-Bothara/polyserve/blob/main/docs/usage.md
9
+ Project-URL: Benchmarks, https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md
10
+ Project-URL: Issues, https://github.com/Aagam-Bothara/polyserve/issues
11
+ Keywords: llm,inference,vllm,llama.cpp,serving,autotuning
12
+ Classifier: Development Status :: 3 - Alpha
13
+ Classifier: Operating System :: POSIX :: Linux
14
+ Classifier: Programming Language :: Python :: 3
15
+ Classifier: Programming Language :: Python :: 3.10
16
+ Classifier: Programming Language :: Python :: 3.11
17
+ Classifier: Programming Language :: Python :: 3.12
18
+ Classifier: Programming Language :: Python :: 3.13
19
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
20
+ Requires-Python: <3.14,>=3.10
21
+ Description-Content-Type: text/markdown
22
+ License-File: LICENSE
23
+ Requires-Dist: pydantic>=2.0
24
+ Requires-Dist: typer>=0.12
25
+ Requires-Dist: rich>=13.0
26
+ Requires-Dist: httpx>=0.27
27
+ Requires-Dist: fastapi>=0.110
28
+ Requires-Dist: uvicorn>=0.29
29
+ Requires-Dist: psutil>=5.9
30
+ Requires-Dist: huggingface_hub>=0.23
31
+ Provides-Extra: nvml
32
+ Requires-Dist: nvidia-ml-py>=12.535.133; extra == "nvml"
33
+ Provides-Extra: dev
34
+ Requires-Dist: pytest>=7.0; extra == "dev"
35
+ Requires-Dist: pytest-asyncio>=0.23; extra == "dev"
36
+ Requires-Dist: ruff==0.15.0; extra == "dev"
37
+ Dynamic: license-file
38
+
39
+ # PolyServe
40
+
41
+ [![tests](https://github.com/Aagam-Bothara/polyserve/actions/workflows/tests.yml/badge.svg)](https://github.com/Aagam-Bothara/polyserve/actions/workflows/tests.yml)
42
+
43
+ **PolyServe is an autotuner for LLM serving: it finds the fastest configuration for your GPU and your traffic, then serves it.** Give it a model and, ideally, a sample of your prompts. It knows which settings each engine offers on each card (vLLM, SGLang and llama.cpp; weight precision, batch size, KV-cache type, speculative decoding and more), drops every configuration that will not fit in memory, measures the rest on your prompts, and serves the fastest one that meets your latency target behind an OpenAI-compatible API. What pays is what it tries and how it measures, not the order it tries things in. The winning settings changed from card to card and from one set of prompts to another, so PolyServe's pick beat a fixed rule of thumb (fp8 weights and KV cache, a short context, batch 256) by 15–97% on the three cards where both ran; but random sampling of the same settings, given the same time, did as well as PolyServe's staged order on two of those cards and better on the third. Calibration measures speed and latency; what quantization costs in answer quality is checked separately.
44
+
45
+ ## Results
46
+
47
+ Measured on rented GPUs against stock `vllm serve` and SGLang defaults:
48
+
49
+ - **The pick holds on prompts it never saw.** Llama 3.1 8B, calibrated on 300 Dolly-15k prompts and measured on 300 others: **+93%** on an A40, **+59%** on an A100, **+71%** on an RTX 4090 over stock SGLang (stock vLLM does not start there), and **+24%** on an H100 over stock vLLM with fp8 weights.
50
+ - **Your traffic decides what pays.** On one A40 with Qwen2.5-3B the pick changed with the prompts: suffix decoding on Dolly prompts (**+122%** on held-out prompts), a draft model on news-article extraction (+66.5%), an fp8 KV cache on real chat (+10%).
51
+ - **Calibration pays for itself within hours.** It took 22–56 minutes per model and workload, repaid by 0.5–2.6 hours of busy serving on held-out prompts. Measuring the busiest load first skips 29–45% of that time without changing a pick.
52
+
53
+ Every result, with methods, ablations, quality checks and limits: [docs/benchmarks.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md).
54
+
55
+ ## Quickstart
56
+
57
+ ```bash
58
+ pip install git+https://github.com/Aagam-Bothara/polyserve.git # not on PyPI yet
59
+ pip install "vllm==0.29.0" # driver 580+; older drivers: "vllm==0.11.0" "transformers>=4.56,<5"
60
+ polyserve serve Qwen/Qwen2.5-3B-Instruct --workload chat
61
+ curl localhost:8000/v1/chat/completions -H 'content-type: application/json' \
62
+ -d '{"model":"Qwen/Qwen2.5-3B-Instruct","messages":[{"role":"user","content":"hi"}]}'
63
+ ```
64
+
65
+ The first launch calibrates and caches the result per machine, model, objective and workload; later launches serve at once. To tune for your own traffic, pass a sample of it: `--workload-file prompts.jsonl`, one prompt per line. Optional: `pip install arctic-inference==0.1.1` adds suffix decoding; SGLang can live in its own environment (set `SGLANG_PYTHON` to its `python`); llama.cpp needs `llama-server` built with CUDA, on `PATH` or in `LLAMA_SERVER`.
66
+
67
+ With Docker (built on the official vLLM image):
68
+
69
+ ```bash
70
+ docker build -t polyserve .
71
+ docker run --gpus all --ipc=host -p 8000:8000 -v ~/.cache/huggingface:/root/.cache/huggingface \
72
+ -v ~/.polyserve:/root/.polyserve polyserve serve Qwen/Qwen2.5-3B-Instruct
73
+ ```
74
+
75
+ ### Supported hardware
76
+
77
+ Linux, Python 3.10–3.13. NVIDIA GPUs of compute capability 7.5 or newer run vLLM, SGLang and llama.cpp (measured on an A40, A100, H100 NVL, L4 and RTX 4090); older NVIDIA GPUs and x86 CPUs run llama.cpp. vLLM-CPU and pre-Turing GPUs have never been benchmarked ([details](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md#hardware-and-engines-measured)).
78
+
79
+ ## Learn more
80
+
81
+ - [docs/benchmarks.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md): every result, at a glance and in full, with methods, ablations, quality, limits and what is still unmeasured.
82
+ - [docs/usage.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/usage.md): how calibration works, workloads, objectives, every search option and the CLI.
83
+ - [docs/writeup.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/writeup.md): the design of the memory planner, calibration and the predictor, what the evidence does and does not support, and the roadmap.
84
+ - [benchmarks/strategies/SUMMARY.md](https://github.com/Aagam-Bothara/polyserve/blob/main/benchmarks/strategies/SUMMARY.md): every table, regenerated from the raw JSON.
85
+ - [CONTRIBUTING.md](https://github.com/Aagam-Bothara/polyserve/blob/main/CONTRIBUTING.md): development setup, tests and adding a backend.
86
+
87
+ MIT licensed.
@@ -0,0 +1,49 @@
1
+ # PolyServe
2
+
3
+ [![tests](https://github.com/Aagam-Bothara/polyserve/actions/workflows/tests.yml/badge.svg)](https://github.com/Aagam-Bothara/polyserve/actions/workflows/tests.yml)
4
+
5
+ **PolyServe is an autotuner for LLM serving: it finds the fastest configuration for your GPU and your traffic, then serves it.** Give it a model and, ideally, a sample of your prompts. It knows which settings each engine offers on each card (vLLM, SGLang and llama.cpp; weight precision, batch size, KV-cache type, speculative decoding and more), drops every configuration that will not fit in memory, measures the rest on your prompts, and serves the fastest one that meets your latency target behind an OpenAI-compatible API. What pays is what it tries and how it measures, not the order it tries things in. The winning settings changed from card to card and from one set of prompts to another, so PolyServe's pick beat a fixed rule of thumb (fp8 weights and KV cache, a short context, batch 256) by 15–97% on the three cards where both ran; but random sampling of the same settings, given the same time, did as well as PolyServe's staged order on two of those cards and better on the third. Calibration measures speed and latency; what quantization costs in answer quality is checked separately.
6
+
7
+ ## Results
8
+
9
+ Measured on rented GPUs against stock `vllm serve` and SGLang defaults:
10
+
11
+ - **The pick holds on prompts it never saw.** Llama 3.1 8B, calibrated on 300 Dolly-15k prompts and measured on 300 others: **+93%** on an A40, **+59%** on an A100, **+71%** on an RTX 4090 over stock SGLang (stock vLLM does not start there), and **+24%** on an H100 over stock vLLM with fp8 weights.
12
+ - **Your traffic decides what pays.** On one A40 with Qwen2.5-3B the pick changed with the prompts: suffix decoding on Dolly prompts (**+122%** on held-out prompts), a draft model on news-article extraction (+66.5%), an fp8 KV cache on real chat (+10%).
13
+ - **Calibration pays for itself within hours.** It took 22–56 minutes per model and workload, repaid by 0.5–2.6 hours of busy serving on held-out prompts. Measuring the busiest load first skips 29–45% of that time without changing a pick.
14
+
15
+ Every result, with methods, ablations, quality checks and limits: [docs/benchmarks.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md).
16
+
17
+ ## Quickstart
18
+
19
+ ```bash
20
+ pip install git+https://github.com/Aagam-Bothara/polyserve.git # not on PyPI yet
21
+ pip install "vllm==0.29.0" # driver 580+; older drivers: "vllm==0.11.0" "transformers>=4.56,<5"
22
+ polyserve serve Qwen/Qwen2.5-3B-Instruct --workload chat
23
+ curl localhost:8000/v1/chat/completions -H 'content-type: application/json' \
24
+ -d '{"model":"Qwen/Qwen2.5-3B-Instruct","messages":[{"role":"user","content":"hi"}]}'
25
+ ```
26
+
27
+ The first launch calibrates and caches the result per machine, model, objective and workload; later launches serve at once. To tune for your own traffic, pass a sample of it: `--workload-file prompts.jsonl`, one prompt per line. Optional: `pip install arctic-inference==0.1.1` adds suffix decoding; SGLang can live in its own environment (set `SGLANG_PYTHON` to its `python`); llama.cpp needs `llama-server` built with CUDA, on `PATH` or in `LLAMA_SERVER`.
28
+
29
+ With Docker (built on the official vLLM image):
30
+
31
+ ```bash
32
+ docker build -t polyserve .
33
+ docker run --gpus all --ipc=host -p 8000:8000 -v ~/.cache/huggingface:/root/.cache/huggingface \
34
+ -v ~/.polyserve:/root/.polyserve polyserve serve Qwen/Qwen2.5-3B-Instruct
35
+ ```
36
+
37
+ ### Supported hardware
38
+
39
+ Linux, Python 3.10–3.13. NVIDIA GPUs of compute capability 7.5 or newer run vLLM, SGLang and llama.cpp (measured on an A40, A100, H100 NVL, L4 and RTX 4090); older NVIDIA GPUs and x86 CPUs run llama.cpp. vLLM-CPU and pre-Turing GPUs have never been benchmarked ([details](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md#hardware-and-engines-measured)).
40
+
41
+ ## Learn more
42
+
43
+ - [docs/benchmarks.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/benchmarks.md): every result, at a glance and in full, with methods, ablations, quality, limits and what is still unmeasured.
44
+ - [docs/usage.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/usage.md): how calibration works, workloads, objectives, every search option and the CLI.
45
+ - [docs/writeup.md](https://github.com/Aagam-Bothara/polyserve/blob/main/docs/writeup.md): the design of the memory planner, calibration and the predictor, what the evidence does and does not support, and the roadmap.
46
+ - [benchmarks/strategies/SUMMARY.md](https://github.com/Aagam-Bothara/polyserve/blob/main/benchmarks/strategies/SUMMARY.md): every table, regenerated from the raw JSON.
47
+ - [CONTRIBUTING.md](https://github.com/Aagam-Bothara/polyserve/blob/main/CONTRIBUTING.md): development setup, tests and adding a backend.
48
+
49
+ MIT licensed.
@@ -0,0 +1,6 @@
1
+ """PolyServe: an autotuner for LLM serving.
2
+
3
+ probe -> select backends -> prepare model -> plan memory -> calibrate -> cache -> serve
4
+ """
5
+
6
+ __version__ = "0.1.0"
@@ -0,0 +1,36 @@
1
+ """Backend registry. New hardware = new class here, no core changes."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Dict
6
+
7
+ from polyserve.backends.base import Backend, BaseBackend, LaunchSpec, LlmtraceHooks, Process, free_port
8
+
9
+
10
+ def registry() -> Dict[str, BaseBackend]:
11
+ from polyserve.backends.llamacpp import LlamaCppCpuBackend, LlamaCppCudaBackend
12
+ from polyserve.backends.sglang import SglangBackend
13
+ from polyserve.backends.vllm import VllmBackend
14
+ from polyserve.backends.vllm_cpu import VllmCpuBackend
15
+
16
+ backends = [VllmBackend(), SglangBackend(), LlamaCppCudaBackend(), LlamaCppCpuBackend(), VllmCpuBackend()]
17
+ return {b.name: b for b in backends}
18
+
19
+
20
+ def get_backend(name: str) -> BaseBackend:
21
+ reg = registry()
22
+ if name not in reg:
23
+ raise KeyError(f"unknown backend {name!r}; known: {sorted(reg)}")
24
+ return reg[name]
25
+
26
+
27
+ __all__ = [
28
+ "Backend",
29
+ "BaseBackend",
30
+ "LaunchSpec",
31
+ "LlmtraceHooks",
32
+ "Process",
33
+ "free_port",
34
+ "registry",
35
+ "get_backend",
36
+ ]
@@ -0,0 +1,314 @@
1
+ """Backend interface (spec section 3) plus the subprocess wrapper every backend uses."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import logging
6
+ import os
7
+ import signal
8
+ import socket
9
+ import subprocess
10
+ import sys
11
+ import time
12
+ from dataclasses import dataclass, field
13
+ from pathlib import Path
14
+ from typing import Dict, List, Optional, Protocol, Tuple, runtime_checkable
15
+
16
+ import httpx
17
+
18
+ from polyserve.memory import MemoryModel
19
+ from polyserve.models import Config, HardwareDescriptor, ModelSpec, PreparedModel
20
+
21
+ logger = logging.getLogger(__name__)
22
+
23
+
24
+ @dataclass
25
+ class LlmtraceHooks:
26
+ """How the calibration driver should talk to (and measure) this backend."""
27
+
28
+ completions_path: str = "/v1/completions"
29
+ health_path: str = "/health"
30
+ models_path: str = "/v1/models"
31
+ stream_usage: bool = True # backend reports usage in the final stream chunk when asked
32
+ model_name: Optional[str] = None # value to put in the "model" field of requests
33
+ tokenizer_id: Optional[str] = None # HF id whose tokenizer counts tokens when usage is absent
34
+ gpu_ids: List[int] = field(default_factory=list) # NVML indices llmtrace should sample
35
+ process_memory: bool = False # sample RSS of the backend process (CPU backends)
36
+
37
+
38
+ @dataclass
39
+ class LaunchSpec:
40
+ args: List[str]
41
+ env: Dict[str, str] = field(default_factory=dict)
42
+ cwd: Optional[str] = None
43
+
44
+
45
+ class Process:
46
+ """A launched backend server: subprocess + readiness probe + logs."""
47
+
48
+ def __init__(self, spec: LaunchSpec, port: int, health_url: str, log_path: Optional[Path] = None):
49
+ self.spec = spec
50
+ self.port = port
51
+ self.health_url = health_url
52
+ self.log_path = log_path
53
+ self._proc: Optional[subprocess.Popen] = None
54
+ self._log_fh = None
55
+ self.started_at: Optional[float] = None
56
+ self.ready_at: Optional[float] = None
57
+
58
+ @property
59
+ def pid(self) -> Optional[int]:
60
+ return self._proc.pid if self._proc else None
61
+
62
+ def start(self) -> "Process":
63
+ env = dict(os.environ)
64
+ env.update(self.spec.env)
65
+ if self.log_path:
66
+ self.log_path.parent.mkdir(parents=True, exist_ok=True)
67
+ self._log_fh = open(self.log_path, "ab")
68
+ stdout = self._log_fh
69
+ else:
70
+ stdout = subprocess.DEVNULL
71
+ logger.info("launch: %s", " ".join(self.spec.args))
72
+ kwargs = {}
73
+ if sys.platform != "win32":
74
+ kwargs["start_new_session"] = True # own process group so we can kill children
75
+ self._proc = subprocess.Popen(
76
+ self.spec.args,
77
+ env=env,
78
+ cwd=self.spec.cwd,
79
+ stdout=stdout,
80
+ stderr=subprocess.STDOUT,
81
+ stdin=subprocess.DEVNULL,
82
+ **kwargs,
83
+ )
84
+ self.started_at = time.monotonic()
85
+ return self
86
+
87
+ def alive(self) -> bool:
88
+ return self._proc is not None and self._proc.poll() is None
89
+
90
+ def returncode(self) -> Optional[int]:
91
+ return self._proc.poll() if self._proc else None
92
+
93
+ def wait_ready(self, timeout: float = 600.0, poll: float = 1.0) -> bool:
94
+ deadline = time.monotonic() + timeout
95
+ with httpx.Client(timeout=5.0) as client:
96
+ while time.monotonic() < deadline:
97
+ if not self.alive():
98
+ logger.error("backend exited during startup (rc=%s); see %s", self.returncode(), self.log_path)
99
+ return False
100
+ try:
101
+ r = client.get(self.health_url)
102
+ if r.status_code < 500:
103
+ self.ready_at = time.monotonic()
104
+ return True
105
+ except httpx.HTTPError:
106
+ pass
107
+ time.sleep(poll)
108
+ logger.error("backend not ready after %.0fs; see %s", timeout, self.log_path)
109
+ return False
110
+
111
+ def stop(self, grace: float = 15.0) -> None:
112
+ if self._proc is None:
113
+ return
114
+ if self.alive():
115
+ try:
116
+ if sys.platform != "win32":
117
+ os.killpg(self._proc.pid, signal.SIGTERM) # own session: group id == server pid
118
+ else:
119
+ self._proc.terminate()
120
+ except Exception:
121
+ pass
122
+ try:
123
+ self._proc.wait(timeout=grace)
124
+ except subprocess.TimeoutExpired:
125
+ try:
126
+ if sys.platform != "win32":
127
+ os.killpg(self._proc.pid, signal.SIGKILL)
128
+ else:
129
+ self._proc.kill()
130
+ except Exception:
131
+ pass
132
+ self._proc.wait(timeout=10)
133
+ if sys.platform != "win32":
134
+ # The server can be gone while its children are not: a vLLM API server that died leaves
135
+ # VLLM::EngineCore holding the GPU, and every later trial then sees less memory. The
136
+ # children share the server's process group, so reap whatever is left of it.
137
+ try:
138
+ os.killpg(self._proc.pid, signal.SIGKILL)
139
+ except OSError:
140
+ pass
141
+ if self._log_fh:
142
+ self._log_fh.close()
143
+ self._log_fh = None
144
+
145
+ def tail_log(self, n: int = 40) -> str:
146
+ if not self.log_path or not self.log_path.exists():
147
+ return ""
148
+ try:
149
+ lines = self.log_path.read_text(errors="replace").splitlines()
150
+ return "\n".join(lines[-n:])
151
+ except OSError:
152
+ return ""
153
+
154
+
155
+ CTX_GRID = (2048, 4096, 8192, 16384, 32768, 65536, 131072)
156
+ CTX_GRID_WIDTH = 3 # sizes tried per calibration: the smallest that fits the workload and the next two
157
+
158
+
159
+ def ctx_grid(max_pos: int, min_ctx: int = 0) -> List[int]:
160
+ """Context lengths to try: the three smallest standard sizes that fit the model and hold the workload.
161
+
162
+ Returns [] if the model cannot hold `min_ctx` at all.
163
+ """
164
+ if min_ctx > max_pos:
165
+ return []
166
+ grid = [c for c in CTX_GRID if min_ctx <= c <= max_pos]
167
+ if not grid:
168
+ # Workload needs more than the largest standard size that fits: use the model max.
169
+ grid = [max_pos]
170
+ return grid[:CTX_GRID_WIDTH]
171
+
172
+
173
+ def render_extra(extra: Dict[str, object], skip: Tuple[str, ...] = ()) -> List[str]:
174
+ """Extra launch flags. `True` renders as a bare flag; `False` and None are dropped."""
175
+ out: List[str] = []
176
+ for k, v in extra.items():
177
+ if k in skip or v is None or v is False:
178
+ continue
179
+ flag = f"--{k.replace('_', '-')}"
180
+ out += [flag] if v is True else [flag, str(v)]
181
+ return out
182
+
183
+
184
+ def free_port(preferred: Optional[int] = None) -> int:
185
+ if preferred:
186
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
187
+ try:
188
+ s.bind(("127.0.0.1", preferred))
189
+ return preferred
190
+ except OSError:
191
+ pass
192
+ with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
193
+ s.bind(("127.0.0.1", 0))
194
+ return s.getsockname()[1]
195
+
196
+
197
+ @runtime_checkable
198
+ class Backend(Protocol):
199
+ name: str
200
+
201
+ def available(self, hw: HardwareDescriptor) -> bool: ...
202
+
203
+ def supports(self, hw: HardwareDescriptor, model: ModelSpec) -> bool: ...
204
+
205
+ def prepare(self, model: ModelSpec, hw: HardwareDescriptor, quants: Optional[List[str]] = None) -> PreparedModel: ...
206
+
207
+ def memory_model(self, hw: HardwareDescriptor) -> MemoryModel: ...
208
+
209
+ def estimate_memory(self, cfg: Config, model: PreparedModel, hw: HardwareDescriptor) -> int: ...
210
+
211
+ def candidate_configs(self, hw: HardwareDescriptor, model: PreparedModel, min_ctx: int = 0) -> List[Config]: ...
212
+
213
+ def launch_spec(self, cfg: Config, model: PreparedModel, port: int) -> LaunchSpec: ...
214
+
215
+ def launch(self, cfg: Config, model: PreparedModel, port: int, log_path: Optional[Path] = None) -> Process: ...
216
+
217
+ def workload_hooks(self, hw: HardwareDescriptor, model: PreparedModel) -> LlmtraceHooks: ...
218
+
219
+ def default_config(self, hw: HardwareDescriptor, model: PreparedModel, min_ctx: int = 0) -> Config: ...
220
+
221
+ def version(self, hw: HardwareDescriptor) -> Optional[str]: ...
222
+
223
+
224
+ class BaseBackend:
225
+ """Shared plumbing. Concrete backends override the abstract-ish methods."""
226
+
227
+ name: str = "base"
228
+ runtime_workspace_bytes: int = 1024 * 1024 * 1024
229
+
230
+ def available(self, hw: HardwareDescriptor) -> bool:
231
+ return hw.backend_available(self.name)
232
+
233
+ def version(self, hw: HardwareDescriptor) -> Optional[str]:
234
+ return hw.backend_version(self.name)
235
+
236
+ def memory_model(self, hw: HardwareDescriptor) -> MemoryModel:
237
+ raise NotImplementedError
238
+
239
+ def calibrated_memory(self, hw: HardwareDescriptor, kv_tokens_fn, device: str) -> MemoryModel:
240
+ """MemoryModel using this machine's fitted workspace/margin when `polyserve memory-report --apply` ran."""
241
+ from polyserve.hardware import hardware_hash
242
+ from polyserve.memcal import margin_override, workspace_override
243
+ from polyserve.memory import DEFAULT_MARGIN_FRACTION
244
+
245
+ hh = hardware_hash(hw)
246
+ ws = workspace_override(hh, self.name)
247
+ mf = margin_override(hh, self.name)
248
+ return MemoryModel(
249
+ runtime_workspace=ws if ws is not None else self.runtime_workspace_bytes,
250
+ kv_tokens_fn=kv_tokens_fn,
251
+ device=device,
252
+ margin_fraction=mf if mf is not None else DEFAULT_MARGIN_FRACTION,
253
+ calibrated=ws is not None or mf is not None,
254
+ )
255
+
256
+ def estimate_memory(self, cfg: Config, model: PreparedModel, hw: HardwareDescriptor) -> int:
257
+ from polyserve.memory import estimate
258
+
259
+ return estimate(hw, model, cfg, self.memory_model(hw)).total
260
+
261
+ def materialize(self, model: PreparedModel, quants: List[str]) -> PreparedModel:
262
+ """Download / convert weights for the quants the planner kept. Default: nothing to do."""
263
+ return model
264
+
265
+ def prefill_variants(self, cfg: Config) -> List[Config]:
266
+ """Configs differing from `cfg` only in the prefill knob. Default: this backend has none."""
267
+ return []
268
+
269
+ def disagg_launch_spec(self, cfg: Config, model: PreparedModel, port: int, role: str,
270
+ kv_transfer_config: dict, gpu_index: int, side_channel_port: int) -> LaunchSpec:
271
+ """Launch one engine of a disaggregated prefill/decode pair."""
272
+ raise NotImplementedError(f"{self.name} does not support disaggregated prefill/decode")
273
+
274
+ # ---- optional search dimensions (default: this backend offers none)
275
+
276
+ supports_tp: bool = False
277
+
278
+ def supported_quants(self, hw: HardwareDescriptor) -> List[str]:
279
+ """Every weight precision this backend could run on `hw` (what --quant filters)."""
280
+ return []
281
+
282
+ def kv_dtypes(self, hw: HardwareDescriptor) -> List[str]:
283
+ """Quantized KV-cache types worth trying on `hw`."""
284
+ return []
285
+
286
+ def batch_ladder(self) -> Tuple[int, ...]:
287
+ """Batch sizes, smallest first, that a smaller KV cache may let the search step up to."""
288
+ return ()
289
+
290
+ def prefix_variants(self, cfg: Config) -> List[Config]:
291
+ """Prefix-cache settings to try when the workload's prompts share a prefix."""
292
+ return []
293
+
294
+ def spec_variants(self, cfg: Config, model: PreparedModel) -> List[Config]:
295
+ """Speculative-decoding settings to try."""
296
+ return []
297
+
298
+ def replica_launch_spec(self, cfg: Config, model: PreparedModel, port: int, gpu_index: int) -> LaunchSpec:
299
+ """One full engine pinned to one GPU, for the replicas layout."""
300
+ spec = self.launch_spec(cfg, model, port)
301
+ spec.env = {**spec.env, "CUDA_VISIBLE_DEVICES": str(gpu_index)}
302
+ return spec
303
+
304
+ def launch_spec(self, cfg: Config, model: PreparedModel, port: int) -> LaunchSpec:
305
+ raise NotImplementedError
306
+
307
+ health_path: str = "/health"
308
+
309
+ def workload_hooks(self, hw: HardwareDescriptor, model: PreparedModel) -> LlmtraceHooks:
310
+ return LlmtraceHooks(gpu_ids=[hw.gpu.index] if hw.gpu else [])
311
+
312
+ def launch(self, cfg: Config, model: PreparedModel, port: int, log_path: Optional[Path] = None) -> Process:
313
+ spec = self.launch_spec(cfg, model, port)
314
+ return Process(spec, port, f"http://127.0.0.1:{port}{self.health_path}", log_path=log_path).start()