[kotselidis]

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, or inactive again if the download fails. The download code is the existing one, moved into a method.
• LOAD MODEL runs the download in the pool controller's background load task, before the inference pools start. 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.
• 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 one batch.
  • 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 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 leaves it inactive;
  • forecasts before the download are rejected, and other models are not affected.
• 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 before loading 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 (new)
• 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