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

.github/workflows/cluster-it-1c1d1a.yml

@@ -81,6 +81,9 @@ jobs:
           PY
           "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_dependency_security.py'
           "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_sundial_flow_loss.py'
+          "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_timestral.py'
+          "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_builtin_model_download.py'
+          "$AINODE_VENV_PY" -m unittest discover -v -s tests -p 'test_builtin_model_load.py'
       - name: Upload Artifact
         if: failure()
         uses: actions/upload-artifact@v6

LICENSE

@@ -426,3 +426,14 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 
 --------------------------------------------------------------------------------
+
+The following files include code modified from Timestral project.
+
+./iotdb-core/ainode/iotdb/ainode/core/model/timestral/*
+
+The Timestral is open source software licensed under the Apache License 2.0
+Project page: https://github.com/mistralai/timestral
+License: https://github.com/mistralai/timestral/blob/main/LICENSE
+Copyright 2026 Mistral AI
+
+--------------------------------------------------------------------------------

NOTICE

@@ -34,6 +34,15 @@ This product includes source code derived from the chronos-forecasting project:
 
 ============================================================================
 
+This product includes source code derived from the Timestral project:
+
+  Timestral
+  Copyright 2026 Mistral AI
+  Licensed under the Apache License, Version 2.0
+  https://github.com/mistralai/timestral
+
+============================================================================
+
 Apache Commons Collections
 Copyright 2001-2019 The Apache Software Foundation
 

integration-test/src/test/java/org/apache/iotdb/ainode/it/AINodeSharedClusterIT.java

@@ -52,6 +52,7 @@
 
 import static org.apache.iotdb.ainode.utils.AINodeTestUtils.BUILTIN_LTSM_MAP;
 import static org.apache.iotdb.ainode.utils.AINodeTestUtils.BUILTIN_MODEL_MAP;
+import static org.apache.iotdb.ainode.utils.AINodeTestUtils.ON_REQUEST_BUILTIN_MODEL_MAP;
 import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkHeader;
 import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkModelNotOnSpecifiedDevice;
 import static org.apache.iotdb.ainode.utils.AINodeTestUtils.checkModelOnSpecifiedDevice;
@@ -236,13 +237,20 @@ private void showBuiltInModelTest(Statement statement) throws SQLException {
       ResultSetMetaData resultSetMetaData = resultSet.getMetaData();
       checkHeader(resultSetMetaData, "ModelId,ModelType,Category,State");
       while (resultSet.next()) {
-        built_in_model_count++;
         FakeModelInfo modelInfo =
             new FakeModelInfo(
                 resultSet.getString(1),
                 resultSet.getString(2),
                 resultSet.getString(3),
                 resultSet.getString(4));
+        FakeModelInfo onRequestModelInfo = ON_REQUEST_BUILTIN_MODEL_MAP.get(modelInfo.getModelId());
+        if (onRequestModelInfo != null) {
+          // Inactive until another test loads it
+          assertEquals(onRequestModelInfo.getModelType(), modelInfo.getModelType());
+          assertEquals(onRequestModelInfo.getCategory(), modelInfo.getCategory());
+          continue;
+        }
+        built_in_model_count++;
         assertTrue(BUILTIN_MODEL_MAP.containsKey(modelInfo.getModelId()));
         assertEquals(BUILTIN_MODEL_MAP.get(modelInfo.getModelId()), modelInfo);
       }
@@ -459,6 +467,21 @@ public void largeTimeSeriesModelLoadInferenceAndForecastTest()
     }
   }
 
+  @Test
+  public void onRequestBuiltInModelLoadInferenceAndForecastTest()
+      throws SQLException, InterruptedException {
+    try (Connection treeConnection = EnvFactory.getEnv().getConnection(BaseEnv.TREE_SQL_DIALECT);
+        Statement treeStatement = treeConnection.createStatement();
+        Connection tableConnection = EnvFactory.getEnv().getConnection(BaseEnv.TABLE_SQL_DIALECT);
+        Statement tableStatement = tableConnection.createStatement()) {
+      for (FakeModelInfo modelInfo : ON_REQUEST_BUILTIN_MODEL_MAP.values()) {
+        // LOAD MODEL downloads the weights before loading the model
+        loadRunAndUnloadModelOnDevices(
+            treeStatement, tableStatement, modelInfo.getModelId(), "cpu");
+      }
+    }
+  }
+
   private void loadRunAndUnloadModelOnDevices(
       Statement treeStatement, Statement tableStatement, String modelId, String devices)
       throws SQLException, InterruptedException {

integration-test/src/test/java/org/apache/iotdb/ainode/utils/AINodeTestUtils.java

@@ -65,6 +65,16 @@ public class AINodeTestUtils {
                   "moment", new FakeModelInfo("moment", "moment", "builtin", "active")))
           .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 
+  /**
+   * Built-in models whose weights are downloaded only when they are loaded. They are inactive until
+   * then.
+   */
+  public static final Map<String, FakeModelInfo> ON_REQUEST_BUILTIN_MODEL_MAP =
+      Stream.of(
+              new AbstractMap.SimpleEntry<>(
+                  "timestral", new FakeModelInfo("timestral", "timestral", "builtin", "inactive")))
+          .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+
   public static final Map<String, FakeModelInfo> BUILTIN_MODEL_MAP;
 
   static {

iotdb-core/ainode/iotdb/ainode/core/inference/pool_controller.py

@@ -42,6 +42,7 @@
 )
 from iotdb.ainode.core.log import Logger
 from iotdb.ainode.core.manager.model_manager import ModelManager
+from iotdb.ainode.core.model.model_constants import ModelStates
 from iotdb.ainode.core.util.atmoic_int import AtomicInt
 from iotdb.ainode.core.util.batch_executor import BatchExecutor
 from iotdb.ainode.core.util.decorator import synchronized
@@ -66,6 +67,9 @@ def __init__(self, result_queue: mp.Queue):
 
         # for pool instances control
         self._task_queue = queue.Queue()
+        # Models whose weights are downloading for a LOAD MODEL, with the load and unload
+        # tasks for them that arrived during the download. Only the worker thread uses it.
+        self._tasks_waiting_for_download: Dict[str, list] = {}
         self._pool_control_worker_thread = threading.Thread(
             target=self._worker_loop, daemon=True
         )
@@ -192,15 +196,39 @@ def _worker_loop(self):
                     "PoolController received task None, the worker loop is existed."
                 )
                 break
-            task_fn, args, kwargs = task
             try:
-                task_fn(*args, **kwargs)
+                self._run_task(*task)
             except Exception as e:
                 logger.error(f"Error executing task: {e}")
             finally:
                 self._task_queue.task_done()
 
+    def _run_task(self, task_fn, args, kwargs):
+        model_id = args[0]
+        if (
+            task_fn != self._finish_download_task
+            and model_id in self._tasks_waiting_for_download
+        ):
+            # Keep the order of a model's tasks: run them after its download ends
+            self._tasks_waiting_for_download[model_id].append((task_fn, args, kwargs))
+            return
+        task_fn(*args, **kwargs)
+
     def _load_one_model_task(self, model_id: str, device_id_list: list[torch.device]):
+        model_info = self._model_manager.get_model_info(model_id)
+        if not model_info.download_on_start and model_info.state != ModelStates.ACTIVE:
+            # Built-in models that are not downloaded at start are downloaded when loaded.
+            # The download runs in its own thread, so that tasks for other models and
+            # stop() do not wait for it.
+            self._tasks_waiting_for_download[model_id] = []
+            threading.Thread(
+                target=self._download_model,
+                args=(model_id, device_id_list),
+                name=f"{ThreadName.INFERENCE_POOL_CONTROLLER.value}-download-{model_id}",
+                daemon=True,
+            ).start()
+            return
+
         def _load_one_model_on_device_task(device_id: torch.device):
             if not self.has_pool_on_device(device_id):
                 self._expand_pools_on_device(model_id, device_id, 1)
@@ -216,6 +244,37 @@ def _load_one_model_on_device_task(device_id: torch.device):
             load_model_futures, return_when=concurrent.futures.ALL_COMPLETED
         )
 
+    def _download_model(self, model_id: str, device_id_list: list[torch.device]):
+        try:
+            succeeded = self._model_manager.download_builtin_model(model_id)
+        except Exception as e:
+            logger.error(f"[Inference] Error downloading model {model_id}: {e}")
+            succeeded = False
+        if not self._stop_event.is_set():
+            self._task_queue.put(
+                (
+                    self._finish_download_task,
+                    (model_id, device_id_list, succeeded),
+                    {},
+                )
+            )
+
+    def _finish_download_task(
+        self, model_id: str, device_id_list: list[torch.device], succeeded: bool
+    ):
+        waiting_tasks = self._tasks_waiting_for_download.pop(model_id)
+        if succeeded:
+            self._load_one_model_task(model_id, device_id_list)
+        else:
+            logger.error(
+                f"[Inference] Cannot load model {model_id}: downloading it failed."
+            )
+        for task in waiting_tasks:
+            try:
+                self._run_task(*task)
+            except Exception as e:
+                logger.error(f"Error executing task: {e}")
+
     def _unload_one_model_task(self, model_id: str, device_id_list: list[torch.device]):
         def _unload_one_model_on_device_task(device_id: torch.device):
             if self.has_request_pools(model_id, device_id):

iotdb-core/ainode/iotdb/ainode/core/manager/model_manager.py

@@ -87,6 +87,9 @@ def get_model_info(
     ) -> Optional[ModelInfo]:
         return self._model_storage.get_model_info(model_id, category)
 
+    def download_builtin_model(self, model_id: str) -> bool:
+        return self._model_storage.download_builtin_model(model_id)
+
     def _refresh(self):
         """Refresh the model list (re-scan the file system)"""
         self._model_storage.discover_all_models()

iotdb-core/ainode/iotdb/ainode/core/model/model_info.py

@@ -32,6 +32,7 @@ def __init__(
         repo_id: str = "",
         auto_map: Optional[Dict] = None,
         hub_mixin_cls: Optional[str] = None,
+        download_on_start: bool = True,
     ):
         self.model_id = model_id
         self.model_type = model_type
@@ -41,6 +42,8 @@ def __init__(
         self.repo_id = repo_id
         self.auto_map = auto_map
         self.hub_mixin_cls = hub_mixin_cls
+        # False for a built-in model whose weights are downloaded only when it is loaded
+        self.download_on_start = download_on_start
 
     def __repr__(self):
         return (
@@ -178,4 +181,16 @@ def __repr__(self):
             "AutoModelForCausalLM": "modeling_moment.MomentForPrediction",
         },
     ),
+    "timestral": ModelInfo(
+        model_id="timestral",
+        category=ModelCategory.BUILTIN,
+        state=ModelStates.INACTIVE,
+        model_type="timestral",
+        pipeline_cls="pipeline_timestral.TimestralPipeline",
+        repo_id="mistralai/Timestral-209M",
+        hub_mixin_cls="modeling_timestral.TimestralForPrediction",
+        # The weights are not under an open-source licence, so AINode downloads them only
+        # when the model is loaded (LOAD MODEL), not at start
+        download_on_start=False,
+    ),
 }