Skip to content

Commit 446c5bc

Browse files
author
Li Zhang
committed
Add vLLM plugin for Chronos-2 inference
Signed-off-by: Li Zhang <lzhanga@amazon.com>
1 parent f951d9a commit 446c5bc

20 files changed

Lines changed: 3025 additions & 0 deletions

‎pyproject.toml‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,10 @@ test = [
5252
"fev>=0.6.1",
5353
"pandas[pyarrow]>=2.0,<2.4",
5454
]
55+
vllm = [
56+
"vllm>=0.13.0",
57+
"pydantic>=2.0",
58+
]
5559
typecheck = ["mypy~=1.9"]
5660
dev = [
5761
"gluonts[pro]~=0.16",
@@ -65,6 +69,12 @@ dev = [
6569
"pandas>=2.0,<2.4",
6670
]
6771

72+
[project.entry-points."vllm.general_plugins"]
73+
chronos2 = "chronos.chronos2.vllm:register_chronos2_model"
74+
75+
[project.entry-points."vllm.io_processor_plugins"]
76+
chronos2 = "chronos.chronos2.vllm:get_chronos2_io_processor"
77+
6878
[project.urls]
6979
Homepage = "https://github.com/amazon-science/chronos-forecasting"
7080
Issues = "https://github.com/amazon-science/chronos-forecasting/issues"

‎src/chronos/chronos2/vllm/README.md‎

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
# Chronos-2 vLLM Plugin
2+
3+
A [vLLM](https://github.com/vllm-project/vllm) plugin that adds support for [Chronos-2](https://github.com/amazon-science/chronos-forecasting) time series forecasting via the `/pooling` API endpoint.
4+
5+
## Overview
6+
7+
Chronos-2 is an encoder-only time series foundation model for zero-shot forecasting. This plugin integrates it with vLLM using the **IOProcessor** plugin interface, so forecast requests are served through vLLM's standard pooling endpoint.
8+
9+
### Features
10+
11+
- **Zero-shot forecasting** — no fine-tuning required
12+
- **Quantile predictions** — probabilistic forecasts with customizable quantile levels
13+
- **Cross-series learning** — information sharing across time series in a batch
14+
- **Covariates support** — past and future covariates (numeric and categorical)
15+
- **Batch forecasting** — process multiple time series in a single request
16+
17+
## Installation
18+
19+
Requires Python 3.10+ and vLLM 0.13.0+.
20+
21+
```bash
22+
pip install chronos-forecasting[vllm]
23+
```
24+
25+
## Quick Start
26+
27+
### 1. Start the Server
28+
29+
```bash
30+
vllm serve amazon/chronos-2 \
31+
--io-processor-plugin chronos2 \
32+
--runner pooling \
33+
--enforce-eager \
34+
--no-enable-prefix-caching \
35+
--skip-tokenizer-init \
36+
--enable-mm-embeds \
37+
--dtype float32 \
38+
--max-model-len 8192
39+
```
40+
41+
### 2. Send a Forecast Request
42+
43+
```bash
44+
curl -X POST http://localhost:8000/pooling \
45+
-H "Content-Type: application/json" \
46+
-d '{
47+
"model": "amazon/chronos-2",
48+
"task": "plugin",
49+
"data": {
50+
"inputs": [
51+
{
52+
"target": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
53+
"item_id": "series_1"
54+
}
55+
],
56+
"parameters": {
57+
"prediction_length": 5,
58+
"quantile_levels": [0.1, 0.5, 0.9]
59+
}
60+
}
61+
}'
62+
```
63+
64+
### 3. Parse the Response
65+
66+
```json
67+
{
68+
"request_id": null,
69+
"created_at": 1739397600,
70+
"data": {
71+
"predictions": [
72+
{
73+
"mean": [11.0, 12.1, 13.0, 14.2, 15.1],
74+
"0.1": [9.5, 10.3, 11.0, 11.8, 12.5],
75+
"0.5": [11.0, 12.0, 13.0, 14.0, 15.0],
76+
"0.9": [12.5, 13.8, 15.0, 16.3, 17.5],
77+
"item_id": "series_1"
78+
}
79+
]
80+
}
81+
}
82+
```
83+
84+
## API Reference
85+
86+
### Request Format
87+
88+
| Field | Type | Required | Description |
89+
|---|---|---|---|
90+
| `model` | `str` | ✅ | Model name (e.g., `"amazon/chronos-2"`) |
91+
| `task` | `str` | ✅ | Must be `"plugin"` |
92+
| `data.inputs` | `list` | ✅ | List of time series inputs (1–1024) |
93+
| `data.parameters` | `dict` | | Forecast parameters |
94+
95+
#### Time Series Input (`data.inputs[*]`)
96+
97+
| Field | Type | Required | Description |
98+
|---|---|---|---|
99+
| `target` | `list[float]` or `list[list[float]]` | ✅ | Historical values (min 5 observations). 1-D for univariate, 2-D for multivariate. |
100+
| `item_id` | `str` | | Identifier echoed in response |
101+
| `start` | `str` | | ISO 8601 timestamp |
102+
| `past_covariates` | `dict[str, list]` | | Past covariate arrays (must match target length) |
103+
| `future_covariates` | `dict[str, list]` | | Future covariate arrays (must match `prediction_length`) |
104+
105+
#### Parameters (`data.parameters`)
106+
107+
| Field | Type | Default | Description |
108+
|---|---|---|---|
109+
| `prediction_length` | `int` | `1` | Forecast horizon (1–1024) |
110+
| `quantile_levels` | `list[float]` | `[0.1, 0.5, 0.9]` | Quantile levels in (0, 1) |
111+
| `freq` | `str` | `null` | Pandas frequency string (e.g., `"D"`, `"H"`) |
112+
| `batch_size` | `int` | `256` | Inference batch size |
113+
| `cross_learning` | `bool` | `false` | Enable cross-series learning |
114+
115+
### Response Format
116+
117+
Each prediction in `data.predictions` contains:
118+
119+
| Field | Type | Description |
120+
|---|---|---|
121+
| `mean` | `list[float]` | Point forecast (mean/median) |
122+
| `"0.1"`, `"0.5"`, etc. | `list[float]` | Named quantile columns matching `quantile_levels` |
123+
| `item_id` | `str` | Echoed from input (if provided) |
124+
125+
## Architecture
126+
127+
The vLLM model wrapper (`Chronos2ForForecasting`) is a thin adapter that delegates all computation to the existing `chronos.chronos2.model.Chronos2Model`. No model architecture is duplicated.
128+
129+
### Module Structure
130+
131+
```
132+
src/chronos/chronos2/vllm/
133+
├── __init__.py # Plugin entry point & registration
134+
├── model.py # Chronos2ForForecasting (thin vLLM wrapper)
135+
├── multimodal.py # MM pipeline for "timeseries" modality
136+
├── io_processor.py # Chronos2IOProcessor (request/response handling)
137+
├── protocol/
138+
│ ├── __init__.py
139+
│ ├── forecast.py # Pydantic models (TimeSeriesInput, ForecastParameters, etc.)
140+
│ ├── validation.py # Input validation logic
141+
│ └── data_prep.py # Tensor preparation from validated inputs
142+
└── utils/
143+
├── __init__.py
144+
├── helpers.py # Utility functions
145+
└── quantiles.py # Quantile selection & interpolation
146+
```
147+
148+
### Key Classes
149+
150+
| Class | File | Purpose |
151+
|---|---|---|
152+
| `Chronos2ForForecasting` | `model.py` | Thin vLLM wrapper — delegates to `chronos.chronos2.model.Chronos2Model` |
153+
| `Chronos2IOProcessor` | `io_processor.py` | Request parsing, validation, pre/post processing |
154+
| `ForecastParameters` | `protocol/forecast.py` | Pydantic validation for forecast parameters |
155+
| `TimeSeriesInput` | `protocol/forecast.py` | Pydantic validation for time series inputs |
156+
| `ForecastPrediction` | `protocol/forecast.py` | Pydantic model for forecast output |
157+
158+
## Troubleshooting
159+
160+
### Server Flags
161+
162+
The following flags are required for Chronos-2:
163+
164+
```bash
165+
--io-processor-plugin chronos2 # Enable the forecast IOProcessor
166+
--enforce-eager # Chronos-2 doesn't support CUDA graphs
167+
--no-enable-prefix-caching # Not applicable for time series
168+
--skip-tokenizer-init # Chronos-2 doesn't use a text tokenizer
169+
```
170+
171+
### Plugin Not Loading
172+
173+
1. Verify installation: `pip list | grep chronos-forecasting`
174+
2. Check entry points:
175+
```python
176+
from importlib.metadata import entry_points
177+
print(list(entry_points(group='vllm.general_plugins')))
178+
```
179+
3. Enable debug logging: `VLLM_LOGGING_LEVEL=DEBUG vllm serve ...`

‎src/chronos/chronos2/vllm/__init__.py‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
"""vLLM Chronos-2 Time Series Forecasting Model Plugin.
2+
3+
This plugin registers the Chronos-2 model with vLLM's ModelRegistry,
4+
allowing it to be used with vLLM's inference engine for time series forecasting.
5+
"""
6+
7+
8+
def register_chronos2_model() -> None:
9+
"""Register Chronos-2 models with vLLM's ModelRegistry.
10+
11+
This function is called automatically when the plugin is loaded
12+
through vLLM's plugin discovery mechanism.
13+
"""
14+
try:
15+
from vllm.logger import init_logger
16+
from vllm.model_executor.models.registry import ModelRegistry
17+
18+
logger = init_logger(__name__)
19+
20+
ModelRegistry.register_model(
21+
"Chronos2Model",
22+
"chronos.chronos2.vllm.model:Chronos2ForForecasting",
23+
)
24+
25+
logger.info("Successfully registered Chronos-2 model with vLLM")
26+
27+
except Exception as e:
28+
from vllm.logger import init_logger
29+
30+
logger = init_logger(__name__)
31+
logger.error(f"Failed to register Chronos-2 model: {e}")
32+
raise
33+
34+
35+
def get_chronos2_io_processor():
36+
"""
37+
Factory function for IOProcessor plugin registration.
38+
39+
This function is called by vLLM when --io-processor-plugin is set to 'chronos2'.
40+
Returns the fully qualified name (module.Class format) as a string.
41+
"""
42+
return "chronos.chronos2.vllm.io_processor.Chronos2IOProcessor"
43+
44+
45+
__all__ = [
46+
"register_chronos2_model",
47+
"get_chronos2_io_processor",
48+
]

0 commit comments

Comments
 (0)