xpark.dataset.TimeSeriesForecast#

class xpark.dataset.TimeSeriesForecast(_local_model: str | None = 'amazon/chronos-2', /, *, base_url: str | None = None, api_key: str = 'NOT_SET', max_retries: int = 3, max_qps: int | None = None, max_concurrency: int | None = None, prediction_length: int, quantile_levels: tuple[float, ...] | list[float] = (0.1, 0.5, 0.9), context_length: int | None = None, cross_learning: bool = False, limit_prediction_length: bool = True, freq: str | None = None, **kwargs: Any)#

Time-series forecasting processor backed by Chronos-2.

Three back-ends are supported and picked based on the current worker’s environment (matches the dispatch used by TextEmbedding):

  • GPU local pipeline (when CUDA_VISIBLE_DEVICES is set, i.e. num_workers={"GPU": ..}): a Chronos2Pipeline is loaded with device_map="auto" and lands on the visible GPU.

  • Remote HTTP (when XPARK_IO_WORKER is set, i.e. num_workers={"IO": ..}): calls are POSTed to a tcray-hosted Chronos-2 /infer endpoint. base_url is required; api_key (or the CHRONOS_API_KEY env var) must be supplied.

  • CPU local pipeline (fallback, i.e. num_workers={"CPU": ..}): the same Chronos2Pipeline is loaded with device_map="auto" and lands on CPU.

The GPU and CPU branches share TimeSeriesForecastLocal; the device_map="auto" argument to Chronos2Pipeline.from_pretrained handles device placement automatically.

The processor treats one row as one complete time series. The target column is expected to be a list<float> containing the historical observations of a single series (already sorted by timestamp and equally spaced). The second positional argument is the timestamps column (list<timestamp>) carrying each row’s historical timestamps; the processor uses it to derive the future timestamps returned alongside the quantile predictions. Any number of additional feature columns can be passed as further positional arguments; each must also be a list<float> of the same length as target and is forwarded to Chronos-2 as a past-only covariate. Future covariates are intentionally not exposed here.

Parameters:

  • _local_model – Chronos-2 model id used by the local pipeline back-end. Available models: [‘amazon/chronos-2’]. Required for the GPU/CPU branches; ignored on IO workers.

  • base_url – Full /infer URL of the remote Chronos-2 endpoint (e.g. https://.../v2/models/m-chronos-2-.../infer). Required for the remote HTTP back-end (IO worker); ignored on GPU/CPU workers.

  • api_key – Bearer token for the remote endpoint. Falls back to the CHRONOS_API_KEY environment variable when unset.

  • max_retries – Number of additional retries on transient failures (HTTP 408/425/429/5xx and connection errors) for the remote back-end. Backs off exponentially with jitter, capped at 8s.

  • **kwargs –

    Transport-level keyword arguments forwarded to every remote request. Only applies to the remote HTTP back-end. The accepted keys are a closed set – extra_headers, timeout, params; any other key raises TypeError at construction time with the allowed set spelled out.

    • extra_headers: per-request headers appended on top of the defaults (bearer auth). Use this for upstream traffic-tagging headers; the name matches the OpenAI-SDK-based operators so the same keyword works across all remote operators.

    • timeout: per-request timeout in seconds; also seeds the underlying httpx client’s default. Defaults to 120s when omitted.

    • params: query-string parameters for every request.

  • prediction_length – Number of future steps to predict for each series. Must be positive. The remote HTTP backend additionally enforces prediction_length <= 1024 (Chronos-2 cloud endpoint hard cap); the local backend has no such ceiling and instead honours limit_prediction_length – pass limit_prediction_length=False to go beyond the model’s recommended horizon locally.

  • quantile_levels –

    Quantile levels to return. Must not contain duplicates (otherwise the resulting field names would collide and the constructor raises ValueError). The output column is a struct containing a leading "timestamp" field of type list<timestamp[ns]> (the future timestamps inferred from the timestamps input column) followed by one list<float32> field per requested quantile level; every inner list has length prediction_length.

    Quantile field naming rule. Each quantile field name is produced by Python’s built-in format(q, ".10g"). For typical usage where q has at most 3 decimal digits (between 0.001 and 0.999), this is exactly the literal you would write yourself, so just wrap the number in quotes when indexing the output, e.g. row["forecast"]["0.1"], row["forecast"]["0.5"], row["forecast"]["0.975"].

  • context_length – Maximum history length used during inference. Defaults to the model’s built-in context length (typically 2048). Histories longer than context_length are left-truncated by Chronos-2.

  • cross_learning – Whether to enable cross-learning mode in Chronos-2. Recommended only when individual histories are very short.

  • limit_prediction_length – Local-only flag. If True (default), raise when prediction_length is greater than the model’s default prediction length. Ignored by the remote HTTP backend, which is always capped at 1024 by the cloud endpoint itself.

  • freq – Optional frequency string (e.g. "h", "D", "W") used to generate future timestamps. When None (default) the frequency is inferred per row; pass an explicit freq to skip inference (faster, and required for series shorter than 3 points).

Examples

from xpark.dataset import TimeSeriesForecast, from_items
from xpark.dataset.expressions import col

# One row per series; ``target`` is the historical observation list,
# ``ts`` is the corresponding timestamp list (equally spaced).
ds = from_items([
    {
        "id": "a",
        "target": [1.0, 2.0, 3.0, 4.0, 5.0],
        "ts": pd.date_range("2026-01-01", periods=5, freq="D").tolist(),
    },
])
ds = ds.with_column(
    "forecast",
    TimeSeriesForecast(prediction_length=3, quantile_levels=[0.1, 0.5, 0.9])
    .options(num_workers={"GPU": 1})
    .with_column(col("target"), col("ts")),
)
# row["forecast"] -> {"timestamp": [2026-01-06, 2026-01-07, 2026-01-08],
#                     "0.1": [...], "0.5": [...], "0.9": [...]}
print(ds.take_all())

Using the remote HTTP back-end (each xpark batch_size becomes one request; tune batch_size on the caller side to trade off payload size vs. throughput):

ds = ds.with_column(
    "forecast",
    TimeSeriesForecast(
        prediction_length=24,
        base_url="https://tcray-aerogate.../v2/models/m-chronos-2-.../infer",
        api_key=os.environ["CHRONOS_API_KEY"],
    )
    .options(num_workers={"IO": 4}, batch_size=32)
    .with_column(col("target"), col("ts")),
)

With past-only covariates (each covariate is a separate list<float> column with the same length as target); covariates follow the timestamps column positionally:

ds = ds.with_column(
    "forecast",
    TimeSeriesForecast(prediction_length=7)
    .with_column(col("target"), col("ts"), col("humidity"), col("wind_speed")),
)

Methods

__call__(target, timestamps, *features)

Call self as a function.

options(**kwargs)

with_column(target, timestamps, *features)

__call__(target: pa.ChunkedArray, timestamps: pa.ChunkedArray, *features: pa.ChunkedArray) → pa.Array#

Call self as a function.

options(**kwargs: Unpack[ExprUDFOptions]) → Self#
with_column(target: pa.ChunkedArray, timestamps: pa.ChunkedArray, *features: pa.ChunkedArray) → pa.Array#