[AINode] Add TiRex-2 as a built-in forecasting model

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 TiRex-2 project.
+
+./iotdb-core/ainode/iotdb/ainode/core/model/tirex2/*
+
+The TiRex-2 is open source software licensed under the Apache License 2.0
+Project page: https://github.com/NX-AI/tirex-2
+License: https://github.com/NX-AI/tirex-2/blob/main/LICENSE
+Copyright 2026 NXAI GmbH
+
+--------------------------------------------------------------------------------

NOTICE

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

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

@@ -397,6 +397,7 @@ public static void forecastTableFunctionErrorTest(
     errorTest(
         statement, invalidOutputLengthSQLWithZero, "701: OUTPUT_LENGTH should be greater than 0");
 
+    int maxOutputLength = AINodeTestUtils.getMaxOutputLength(modelInfo.getModelId());
     String invalidOutputLengthSQLWithOutOfRange =
         String.format(
             FORECAST_TABLE_FUNCTION_SQL_TEMPLATE,
@@ -405,13 +406,15 @@ public static void forecastTableFunctionErrorTest(
             5760,
             2880,
             5760,
-            2881,
+            maxOutputLength + 1,
             1,
             "time");
     errorTest(
         statement,
         invalidOutputLengthSQLWithOutOfRange,
-        "1599: Error occurred while executing forecast:[Attribute output_length expect value between 1 and 2880, got 2881 instead.]");
+        String.format(
+            "1599: Error occurred while executing forecast:[Attribute output_length expect value between 1 and %d, got %d instead.]",
+            maxOutputLength, maxOutputLength + 1));
 
     String invalidOutputIntervalSQL =
         String.format(

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

@@ -62,7 +62,9 @@ public class AINodeTestUtils {
               new AbstractMap.SimpleEntry<>(
                   "toto", new FakeModelInfo("toto", "toto", "builtin", "active")),
               new AbstractMap.SimpleEntry<>(
-                  "moment", new FakeModelInfo("moment", "moment", "builtin", "active")))
+                  "moment", new FakeModelInfo("moment", "moment", "builtin", "active")),
+              new AbstractMap.SimpleEntry<>(
+                  "tirex2", new FakeModelInfo("tirex2", "tirex2", "builtin", "active")))
           .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 
   public static final Map<String, FakeModelInfo> BUILTIN_MODEL_MAP;
@@ -95,6 +97,17 @@ public class AINodeTestUtils {
     BUILTIN_MODEL_MAP = Collections.unmodifiableMap(tmp);
   }
 
+  /** AINode's default limit on output_length (ain_inference_max_output_length). */
+  public static final int DEFAULT_MAX_OUTPUT_LENGTH = 2880;
+
+  /** Built-in models whose own forecast horizon is below AINode's limit. */
+  private static final Map<String, Integer> MODEL_MAX_OUTPUT_LENGTH =
+      Collections.singletonMap("tirex2", 320);
+
+  public static int getMaxOutputLength(String modelId) {
+    return MODEL_MAX_OUTPUT_LENGTH.getOrDefault(modelId, DEFAULT_MAX_OUTPUT_LENGTH);
+  }
+
   private static final Logger LOGGER = LoggerFactory.getLogger(AINodeTestUtils.class);
 
   public static void checkHeader(ResultSetMetaData resultSetMetaData, String title)

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

@@ -217,6 +217,21 @@ def _do_inference_and_construct_resp(
             resp_list.append(resp)
         return resp_list
 
+    def _max_output_length(self, model_id: str) -> int:
+        """AINode's output_length limit, lowered to the model's own limit if it has one.
+
+        The check runs before a request reaches an inference pool, where a failing
+        request would stop the pool.
+        """
+        max_output_length = (
+            AINodeDescriptor().get_config().get_ain_inference_max_output_length()
+        )
+        if self._model_manager.is_model_registered(model_id):
+            model_limit = self._model_manager.get_model_info(model_id).max_output_length
+            if model_limit is not None:
+                max_output_length = min(max_output_length, model_limit)
+        return max_output_length
+
     def _run_forecast(
         self,
         req,
@@ -232,17 +247,10 @@ def _run_forecast(
 
             inference_attrs = extract_attrs(req)
             output_length = int(inference_attrs.get("output_length", 96))
-            if (
-                output_length
-                > AINodeDescriptor().get_config().get_ain_inference_max_output_length()
-            ):
+            max_output_length = self._max_output_length(model_id)
+            if not 1 <= output_length <= max_output_length:
                 raise NumericalRangeException(
-                    "output_length",
-                    output_length,
-                    1,
-                    AINodeDescriptor()
-                    .get_config()
-                    .get_ain_inference_max_output_length(),
+                    "output_length", output_length, 1, max_output_length
                 )
 
             model_inputs: list[dict[str, torch.Tensor | dict[str, torch.Tensor]]] = [

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,
+        max_output_length: Optional[int] = None,
     ):
         self.model_id = model_id
         self.model_type = model_type
@@ -41,6 +42,9 @@ def __init__(
         self.repo_id = repo_id
         self.auto_map = auto_map
         self.hub_mixin_cls = hub_mixin_cls
+        # The longest forecast the model can produce, if it is shorter than AINode's own limit.
+        # Requests are checked against it before they reach an inference pool.
+        self.max_output_length = max_output_length
 
     def __repr__(self):
         return (
@@ -178,4 +182,14 @@ def __repr__(self):
             "AutoModelForCausalLM": "modeling_moment.MomentForPrediction",
         },
     ),
+    "tirex2": ModelInfo(
+        model_id="tirex2",
+        category=ModelCategory.BUILTIN,
+        state=ModelStates.INACTIVE,
+        model_type="tirex2",
+        pipeline_cls="pipeline_tirex2.TiRex2Pipeline",
+        repo_id="NX-AI/TiRex-2",
+        hub_mixin_cls="modeling_tirex2.TiRex2ForPrediction",
+        max_output_length=320,
+    ),
 }

iotdb-core/ainode/iotdb/ainode/core/model/tirex2/__init__.py

@@ -0,0 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# Copyright (c) NXAI GmbH.
+# SPDX-License-Identifier: Apache-2.0

iotdb-core/ainode/iotdb/ainode/core/model/tirex2/component/__init__.py

@@ -0,0 +1,19 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+#
+# Copyright (c) NXAI GmbH.
+# SPDX-License-Identifier: Apache-2.0