[AINode] Add Timestral as a built-in forecasting model, downloaded on load by kotselidis · Pull Request #19272 · apache/iotdb

kotselidis · 2026-11-13T21:06:15Z

Description

This PR adds Timestral (Mistral AI) to AINode's built-in forecasting models, next to timer_xl, sundial, chronos2, moirai2, toto and moment. Timestral is a patch-based transformer for multivariate probabilistic forecasting. It forecasts 21 quantiles (0.01 to 0.99) from a context of up to 32,768 steps, and it can use past-only and future-known covariates. The model is registered as timestral, with the weights of mistralai/Timestral-209M. The 209M model is the largest of the four published sizes (5M, 28M, 123M, 209M) and the one Timestral's own documentation uses.

Unlike the other built-in models, AINode does not download Timestral's weights at startup. They are downloaded when a user loads the model with LOAD MODEL. The reason is the weights' licence, below.

Weights licence

Timestral's code is under the Apache License 2.0. Its weights are not. Every published size is under a "Modified MIT" licence (licence text), whose second condition reads:

You are not authorized to exercise any rights under this license if the global consolidated monthly revenue of your company (or that of your employer) exceeds $20 million (or its equivalent in another currency) for the preceding month.

The weights are not part of the AINode distribution in either case. However, AINode downloads the weights of every built-in model when it starts. As an automatic built-in, Timestral would put weights that some users are not licensed to use on every AINode installation, without anyone asking for them. With this PR, the weights reach a machine only when someone runs LOAD MODEL timestral, which is a deliberate act by that user. The same question applies to moirai2, whose weights (Salesforce/moirai-2.0-R-small) are under CC BY-NC 4.0 and are downloaded at startup today. The new download_on_start flag could be applied to it as well. That is left to the maintainers, and this PR does not change moirai2.

The maintainers may prefer a different policy, such as not shipping models with restricted weights at all. The Timestral integration does not depend on the on-request download, so it can follow whichever policy is chosen.

Usage

SHOW MODELS
-- timestral | timestral | builtin | inactive   (weights not downloaded)

SELECT * FROM FORECAST(model_id => 'timestral', targets => (SELECT time, s FROM db.t) ORDER BY time, output_length => 24)
-- 1504: Model [timestral] is not downloaded at start. Use 'LOAD MODEL timestral TO DEVICES ...' to download and load it; 'SHOW MODELS' shows it as active once it is downloaded.

LOAD MODEL timestral TO DEVICES 'cpu'
-- downloads config.json and model.safetensors from mistralai/Timestral-209M, then loads the model

SHOW MODELS
-- timestral | timestral | builtin | active

SELECT * FROM FORECAST(model_id => 'timestral', targets => (SELECT time, s FROM db.t) ORDER BY time, output_length => 24)
-- 24 forecast rows

Once downloaded, the weights stay in AINode's model directory, and the model is active from the next start on.

Example: the input is 1024 points of 50 + 10·sin(2πt/24) + 0.01·t. The request is a TForecastReq with outputLength = 24, sent to AINode's inference manager on CPU with the Timestral-209M checkpoint. The response is 24 values (the median forecast). The first six:

step	forecast	true value
1	51.507	51.580
2	50.514	50.591
3	50.168	50.260
4	50.564	50.611
5	51.589	51.620
6	53.205	53.219

The mean absolute error is 0.031 over 24 steps and 0.049 over 96 steps.

On-request download of built-in models

ModelInfo gets a download_on_start flag, which is True by default. Every existing built-in model keeps it, so their behaviour does not change.
ModelStorage registers a built-in model with download_on_start=False as inactive at startup, unless its config.json and model.safetensors are already in its directory. ModelStorage.download_builtin_model downloads the two files and marks the model active. If the download fails, or the downloaded config.json cannot be read, the model goes back to inactive and the load stops. The download code is the existing one, moved into a method.
LOAD MODEL starts the download in its own thread, and the pool controller starts the inference pools when the download ends. The LOAD MODEL statement returns immediately, as it does today. SHOW MODELS shows activating during the download, and SHOW LOADED MODELS shows the model once it is loaded.
A download does not hold up the pool controller. Loads and unloads of other models run while it is in progress. Loads and unloads of the downloading model wait, and run after it is loaded, in the order they were requested. Stopping AINode does not wait for a download either.
FORECAST and CALL INFERENCE on such a model before the download return error 1504 with the message shown above. Without this check, they would fail while loading the missing weights.

Model code

iotdb/ainode/core/model/timestral/ contains the inference code from mistralai/timestral 1.0.0 (Apache License 2.0), in the same way as chronos2/ contains the code from chronos-forecasting. The model, forecaster, layers, masks, scaler, covariate encoding and configuration are unchanged except for these points:

Imports use the iotdb.ainode.core.model.timestral package, and files carry the ASF header followed by Mistral AI's copyright notice.
The Hugging Face download helper (loading.py) and TimestralForecaster.from_pretrained are left out. AINode downloads and loads the checkpoint itself.

AINode integration

modeling_timestral.TimestralForPrediction is the model, loaded through PyTorchModelHubMixin (hub_mixin_cls in ModelInfo), a loading path AINode already supports. config.json is read into Timestral's own TimestralConfig, which rejects unknown keys. The weights load with strict=True. The rotary-embedding frequencies are recomputed whenever the model is moved to a device, as Timestral's own loader does after moving the model.
pipeline_timestral.TimestralPipeline runs Timestral's TimestralForecaster with its default settings:
- Each input is one series, and its targets are forecast jointly.
- A covariate that appears only in past_covariates becomes a past covariate. A covariate that also appears in future_covariates becomes a known covariate, built from its past and future values.
- Inputs of the same shape without covariates run in batches of up to 64 targets. Timestral forecasts a batch as one set of rows, one per target, and its variate-attention mask grows with the square of the number of rows. With the bound, the memory of a batch of requests grows linearly with the number of requests.
- Forecasts are averaged with the reflected forecast of the negated series, and are clipped at 0 for series whose context is non-negative.
- Horizons longer than one decode block (1024 steps) are forecast block by block, so any output_length that AINode accepts (up to 2880) works.
- The pipeline returns the 0.5 quantile.

Dependencies

There are no new dependencies. Timestral needs only PyTorch, NumPy, safetensors and huggingface_hub, which AINode already pins. pyproject.toml, ainode.spec and ainode.xml are unchanged. ainode.spec already collects every module under iotdb.ainode.core.

Licences

The code is under the Apache License 2.0 (licence). This PR adds it to LICENSE, NOTICE and AINode's NOTICE-binary. See the note above about the licence of the weights.

Tests

tests/test_timestral.py (unittest, CPU only, no downloads) uses a two-layer Timestral with random weights. The tests:
- save it as a checkpoint and check that it reloads with identical weights;
- check that a checkpoint with a missing weight, or a config with an unknown key, is rejected;
- check that the rotary-embedding frequencies are rebuilt when the model is cast;
- load it through AINode's built-in loader and the TimestralPipeline, and check the output shapes for 1-D and 2-D targets over more than one decode block, and that the result is the forecaster's median;
- check that inputs of the same shape are split into batches by the number of targets, and that each forecast equals the forecast of its own series;
- check how covariates are assembled and that covariates change the forecast.
tests/test_builtin_model_download.py (unittest, no network: hf_hub_download is mocked) checks that:
- Timestral is inactive and not downloaded at startup, while the other built-ins are;
- download_builtin_model fetches it and marks it active, and a second call downloads nothing;
- it is active at the next startup once downloaded;
- a failed download, or a downloaded config.json that cannot be read, leaves it inactive;
- forecasts before the download are rejected, and other models are not affected.
tests/test_builtin_model_load.py (unittest, no downloads) checks that during a download the pool controller still runs other models' loads and unloads. It also checks that the downloading model's unload runs after its load, that a failed download skips only the load, and that the controller stops without waiting for the download.
The AINode CI job runs these three test files, next to test_dependency_security.py and test_sundial_flow_loss.py.
Integration tests: AINodeTestUtils.ON_REQUEST_BUILTIN_MODEL_MAP lists timestral. The SHOW MODELS test accepts it in either state, and the new onRequestBuiltInModelLoadInferenceAndForecastTest loads it on the CPU, which downloads it, then runs CALL INFERENCE and FORECAST.
A run of AINode's RPC handler on CPU confirms the whole flow. SHOW MODELS first reports inactive, and FORECAST returns 1504. LOAD MODEL timestral TO DEVICES 'cpu' downloads the files and starts one request pool. FORECAST then returns the 24 values above through that pool.
With the published checkpoints on CPU, the forecasts are bit-identical to the unmodified Timestral 1.0.0 code for a multivariate series, a 1100-step forecast spanning two decode blocks, and a series with past and known covariates.

This PR has:

been self-reviewed.
added or updated version, license, or notice information
added comments explaining the "why" and the intent of the code wherever would not be obvious
for an unfamiliar reader.
added unit tests or modified existing tests to cover new code paths, ensuring the threshold
for code coverage.
added integration tests.
added documentation for new or modified features or behaviors.

Key changed/added classes (or packages if there are too many classes) in this PR

iotdb/ainode/core/model/timestral/ (new): the Timestral model and forecaster, TimestralForPrediction, TimestralPipeline
iotdb/ainode/core/model/model_info.py: ModelInfo.download_on_start; timestral entry in BUILTIN_HF_TRANSFORMERS_MODEL_MAP
iotdb/ainode/core/model/model_storage.py: download_builtin_model; built-ins with download_on_start=False are not downloaded at start
iotdb/ainode/core/inference/pool_controller.py: LOAD MODEL downloads such a model in its own thread, then loads it
iotdb/ainode/core/rpc/handler.py: forecasts on a model that is not downloaded yet are rejected
iotdb-core/ainode/tests/test_timestral.py, tests/test_builtin_model_download.py, tests/test_builtin_model_load.py (new)
.github/workflows/cluster-it-1c1d1a.yml: runs the new test files
integration-test/.../ainode/utils/AINodeTestUtils.java, .../it/AINodeSharedClusterIT.java: ON_REQUEST_BUILTIN_MODEL_MAP and its tests
LICENSE, NOTICE, iotdb-core/ainode/src/assembly/resources/NOTICE-binary

Copilot

🟡 Changes recommended

Malformed metadata can still trigger inference-pool creation, combined-series attention can exhaust memory, and slow downloads block management operations and delay shutdown.

3 open issues

What changed in this PR

Adds Timestral-209M as an AINode built-in forecasting model. Its restricted-license weights are downloaded explicitly through LOAD MODEL, rather than automatically at startup.

Changes:

Adds Timestral inference code and a median-forecast pipeline supporting multivariate inputs, covariates, and multi-block horizons.
Introduces on-request downloads, model-state handling, and RPC checks for unavailable weights.
Adds unit and integration coverage and updates third-party license notices.

File	Description
`NOTICE`	Adds Timestral attribution.
`LICENSE`	Records the vendored code’s Apache license.
`iotdb-core/ainode/tests/test_timestral.py`	Tests checkpoint loading, forecasts, and covariates.
`iotdb-core/ainode/tests/test_builtin_model_download.py`	Tests download states, retries, and pre-download rejection.
`iotdb-core/ainode/src/assembly/resources/NOTICE-binary`	Adds attribution to AINode binary notices.
`iotdb-core/ainode/iotdb/ainode/core/rpc/handler.py`	Rejects inference before on-request weights are available.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/scaler.py`	Implements normalization and forecast rescaling.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/pipeline_timestral.py`	Adapts AINode inputs and returns median forecasts.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/modeling_timestral.py`	Integrates strict checkpoint loading and device transitions.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/model.py`	Implements the Timestral network and quantile prediction.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/masks.py`	Defines attention masks and temporal positions.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/layers.py`	Adds transformer layers and rotary embeddings.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/forecaster.py`	Handles block decoding, covariates, and forecast adjustments.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/covariate_encoding.py`	Encodes categorical covariates numerically.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/config.py`	Defines and validates checkpoint configuration.
`iotdb-core/ainode/iotdb/ainode/core/model/timestral/__init__.py`	Adds the package initializer and attribution.
`iotdb-core/ainode/iotdb/ainode/core/model/model_storage.py`	Manages deferred downloads and stored-model states.
`iotdb-core/ainode/iotdb/ainode/core/model/model_info.py`	Registers Timestral and adds the startup-download flag.
`iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py`	Exposes built-in model downloads to callers.
`iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py`	Downloads deferred weights before loading inference pools.
`integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java`	Defines expected on-request built-in models.
`integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeSharedClusterIT.java`	Covers model listing and load/inference/forecast flow.

_{Review effort: Balanced}

CRZbulabula · 2026-11-14T12:10:28Z

+                succeeded = self._download_builtin_model_files(model_id, model_dir)
+            finally:
+                self._set_model_download_result(model_id, succeeded)
+            return succeeded


Reproduced with a malformed local config.json: download_builtin_model("timestral") returns True even though _set_model_download_result() leaves the model INACTIVE. The controller consequently proceeds to create a pool whose checkpoint cannot load. Please propagate configuration-processing failure into the returned result and add a regression test for malformed JSON.

CRZbulabula · 2026-11-14T12:10:26Z

+        Forecast the series, batching those of the same shape without covariates.
+
+        Returns
+        -------
+        list of torch.Tensor
+            One tensor per series, of shape (target_count, output_length, quantile_count).
+        """
+        output_length = infer_kwargs.get("output_length", 96)
+        outputs: list[torch.Tensor | None] = [None] * len(inputs)
+        batches: dict[tuple, list[int]] = {}
+        for i, item in enumerate(inputs):
+            if item["past_covariates"] is None and item["known_covariates"] is None:
+                batches.setdefault(tuple(item["context"].shape), []).append(i)
+            else:
+                batches[("series", i)] = [i]


I checked the actual mask builder with meta tensors: one 100-target series with 1,056 patch slots produces a dense boolean mask of shape [1056, 1, 100, 100] (10,560,000 bytes), while ten such series produce [1056, 1, 1000, 1000] (1,056,000,000 bytes). That is 100 times the mask memory for ten times the input; the normal request scheduler can select ten requests at once. Please preserve a separate batch dimension for independent series or process them separately/in bounded batches, and add a regression covering multiple concurrent requests.

CRZbulabula · 2026-11-14T12:10:23Z

+        model_info = self._model_manager.get_model_info(model_id)
+        if not model_info.download_on_start:
+            # Built-in models that are not downloaded at start are downloaded when loaded
+            if not self._model_manager.download_builtin_model(model_id):


Confirmed with an event-controlled download: a task queued behind LOAD MODEL timestral does not execute until the download completes. This is the controller's only management worker, and stop() joins it before closing existing pools, so a slow download also delays unrelated unloads and shutdown. Please move the transfer off this queue and make shutdown cancellation bounded, while preserving ordering for the affected model.

CRZbulabula

Reviewed e3475fa. I independently reproduced the three issues in the existing review threads and added the evidence there: synchronous downloads block the shared controller, failed configuration processing still reports download success, and batching independent series causes quadratic mask-memory growth across requests. These should be addressed before merging.

All ten new Python tests pass locally (six model tests and four download tests), using Python 3.11 / PyTorch 2.11.0 on CPU. This does not validate the pinned PyTorch 2.14.1 or GPU behavior. The current AINode CI job is green, but its explicit Python test commands only select test_dependency_security.py and test_sundial_flow_loss.py; please include both new regression suites in CI as well.

I replied separately in #19270 about the optional download design and the remaining weights-license question.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

Copilot

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

Add Timestral to AINode as a built-in model downloaded on load

e3475fa

HTHou requested a review from CRZbulabula November 14, 2026 00:09

SpriCoder requested a balanced review from Copilot November 14, 2026 08:59

Copilot started reviewing on behalf of SpriCoder November 14, 2026 08:59

View session

Copilot AI reviewed Nov 14, 2026

View reviewed changes

CRZbulabula reviewed Nov 14, 2026

View reviewed changes

Fail LOAD MODEL when a downloaded config.json cannot be read

fb8f3b4

Copilot AI balanced review requested due to automatic review settings November 14, 2026 17:22

Copilot AI reviewed Nov 14, 2026

Download on-request model weights without blocking the pool controller

f6c1bec

Copilot AI balanced review requested due to automatic review settings November 14, 2026 17:30

Copilot AI reviewed Nov 14, 2026

Bound the number of target rows in a Timestral batch

8e171ce

Copilot AI balanced review requested due to automatic review settings November 14, 2026 17:35

Copilot AI reviewed Nov 14, 2026

Run the Timestral and on-request download tests in the AINode CI job

aac7c49

Copilot AI balanced review requested due to automatic review settings November 14, 2026 17:38

Copilot AI reviewed Nov 14, 2026

Navigation Menu

Repository navigation

[AINode] Add Timestral as a built-in forecasting model, downloaded on load - #19272

[AINode] Add Timestral as a built-in forecasting model, downloaded on load#19272
kotselidis wants to merge 5 commits into
apache:masterfrom
kotselidis:ainode-timestral

kotselidis commented Nov 13, 2026 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

CRZbulabula Nov 14, 2026

Uh oh!

CRZbulabula Nov 14, 2026

Uh oh!

CRZbulabula Nov 14, 2026

Uh oh!

CRZbulabula left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

kotselidis commented Nov 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Weights licence

Usage

On-request download of built-in models

Model code

AINode integration

Dependencies

Licences

Tests

Key changed/added classes (or packages if there are too many classes) in this PR

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

🟡 Changes recommended

Uh oh!

CRZbulabula Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

CRZbulabula Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

CRZbulabula Nov 14, 2026

Choose a reason for hiding this comment

Uh oh!

CRZbulabula left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kotselidis commented Nov 13, 2026 •

edited

Loading