Merge branch 'main' into new_model_integrations

services/boilerplate/dirutil.py

@@ -0,0 +1,35 @@
+"""Utilities for directory operations."""
+
+import os
+from pathlib import Path
+
+
+def resolve_model_path(search_path: str, model_id: str) -> Path:
+    """Find the first path under search_path for model_id. All entries in
+    search_path must be:
+    * an existing directory
+    * must be readable by the current process
+
+    Args:
+        search_path (str): A unix-like ":" separated list of directories such a "dir1:dir2"
+        model_id (str): a model_id (which is really just a subdirectory under dir1 or dir2)
+
+    Returns:
+        Path: the first matching path, None if no path is fount.
+    """
+
+    _amodeldir_found = next(
+        (
+            adir
+            for adir in (Path(p) for p in search_path.split(":"))
+            if adir.exists()
+            and adir.is_dir()
+            and os.access(adir, os.R_OK)
+            and (adir / model_id).exists()
+            and os.access(adir / model_id, os.R_OK)
+        ),
+        None,
+    )
+    if not _amodeldir_found:
+        return None
+    return _amodeldir_found / model_id

services/boilerplate/inference_payloads.py

@@ -101,12 +101,10 @@ class ForecastingMetadataInput(BaseMetadataInput):
 
 
 class BaseParameters(BaseModel):
-    model_config = ConfigDict(extra="forbid", protected_namespaces=())
-
+    model_config = ConfigDict(extra="allow", protected_namespaces=())
 
-class ForecastingParameters(BaseModel):
-    model_config = ConfigDict(extra="forbid", protected_namespaces=())
 
+class ForecastingParameters(BaseParameters):
     prediction_length: Optional[int] = Field(
         description="The prediction length for the forecast."
         " The service will return this many periods beyond the last"

services/finetuning/Makefile

@@ -2,10 +2,13 @@ CONTAINER_BUILDER ?= docker
 
 # copies boilerplate code to suitable locations
 boilerplate:
+	rm tsfmfinetuning/.gitignore || true
+	echo "# THIS FILE IS AUTOMATICALLY GENERATED, YOUR CHANGES WILL BE OVERWRITTEN" > tsfmfinetuning/.gitignore
 	for f in ../boilerplate/*.py; do \
 		echo $$f; \
 		cat ../boilerplate/warning.txt > tsfmfinetuning/$$(basename $$f); \
 		cat $$f>>tsfmfinetuning/$$(basename $$f); \
+		echo $$(basename $$f) >> tsfmfinetuning/.gitignore; \
 		done 
 
 image:

services/finetuning/tsfmfinetuning/.gitignore

@@ -1,3 +1,6 @@
-inference_payloads.py
-hfutil.py
+# THIS FILE IS AUTOMATICALLY GENERATED, YOUR CHANGES WILL BE OVERWRITTEN
+dataframe_checks.py
+dirutil.py
 errors.py
+hfutil.py
+inference_payloads.py

services/inference/.gitignore

@@ -1,3 +1 @@
-# These version placeholders will be replaced later during substitution.
-__version__ = "0.0.0"
-__version_tuple__ = (0, 0, 0)
+prometheus_metrics

services/inference/Makefile

@@ -2,10 +2,13 @@ CONTAINER_BUILDER ?= docker
 
 # copies boilerplate code to suitable locations
 boilerplate:
+	rm tsfminference/.gitignore || true
+	echo "# THIS FILE IS AUTOMATICALLY GENERATED, YOUR CHANGES WILL BE OVERWRITTEN" > tsfminference/.gitignore
 	for f in ../boilerplate/*.py; do \
 		echo $$f; \
 		cat ../boilerplate/warning.txt > tsfminference/$$(basename $$f); \
 		cat $$f>>tsfminference/$$(basename $$f); \
+		echo $$(basename $$f) >> tsfminference/.gitignore; \
 		done 
 
 create_prometheus_metrics_dir:
@@ -16,7 +19,7 @@ create_prometheus_metrics_dir:
 start_service_local: create_prometheus_metrics_dir boilerplate
 	PROMETHEUS_MULTIPROC_DIR=./prometheus_metrics \
 	TSFM_PYTHON_LOGGING_LEVEL="ERROR" \
-	TSFM_MODEL_DIR=./mytest-tsfm \
+	TSFM_MODEL_DIR=./foobaz:./mytest-tsfm \
 	TSFM_ALLOW_LOAD_FROM_HF_HUB=1 \
 	python -m gunicorn \
 	-w 1 \

services/inference/tests/test_inference_lib.py

@@ -2,7 +2,9 @@
 #
 
 import copy
+import json
 import os
+import tempfile
 from datetime import timedelta
 from pathlib import Path
 
@@ -112,6 +114,14 @@ def test_forecast_with_good_data(ts_data_base: pd.DataFrame, forecasting_input_b
         return
     df = copy.deepcopy(data)
     input.data = df.to_dict(orient="list")
+
+    # useful for generating sample payload files
+    if int(os.environ.get("TSFM_TESTS_DO_VERBOSE_DUMPS", "0")) == 1:
+        with open(f"{tempfile.gettempdir()}/{model_id}.payload.json", "w") as out:
+            foo = copy.deepcopy(df)
+            foo["date"] = foo["date"].apply(lambda x: x.isoformat())
+            json.dump(foo.to_dict(orient="list"), out)
+
     runtime: InferenceRuntime = InferenceRuntime(config=config)
     po: PredictOutput = runtime.forecast(input=input)
     results = pd.DataFrame.from_dict(po.results[0])

services/inference/tsfminference/.gitignore

@@ -1,4 +1,6 @@
-inference_payloads.py
+# THIS FILE IS AUTOMATICALLY GENERATED, YOUR CHANGES WILL BE OVERWRITTEN
+dataframe_checks.py
+dirutil.py
 errors.py
 hfutil.py
-dataframe_checks.py
+inference_payloads.py

services/inference/tsfminference/__init__.py

@@ -38,7 +38,20 @@
 )
 
 # use TSFM_MODEL_DIR preferentially. If not set, use HF_HOME or the system tempdir if that's not set.
-TSFM_MODEL_DIR: Path = Path(os.environ.get("TSFM_MODEL_DIR", os.environ.get("HF_HOME", tempfile.gettempdir())))
-
-if not TSFM_MODEL_DIR.exists():
-    raise Exception(f"TSFM_MODEL_DIR {TSFM_MODEL_DIR} does not exist.")
+TSFM_MODEL_DIR: str = os.environ.get("TSFM_MODEL_DIR", os.environ.get("HF_HOME", tempfile.gettempdir()))
+
+# basic checks
+# make sure at least one of them is a valid directory
+# make sure it's readable as well
+_amodeldir_found = next(
+    (
+        adir
+        for adir in (Path(p) for p in TSFM_MODEL_DIR.split(":"))
+        if adir.exists() and adir.is_dir() and os.access(adir, os.R_OK)
+    ),
+    None,
+)
+if not _amodeldir_found and not TSFM_ALLOW_LOAD_FROM_HF_HUB:
+    raise Exception(
+        f"None of the values given in TSFM_MODEL_DIR {TSFM_MODEL_DIR} are an existing and readable directory."
+    )

services/inference/tsfminference/inference.py

@@ -17,6 +17,7 @@
 from . import TSFM_ALLOW_LOAD_FROM_HF_HUB, TSFM_MODEL_DIR
 from .constants import API_VERSION
 from .dataframe_checks import check
+from .dirutil import resolve_model_path
 from .errors import error_message
 from .inference_payloads import ForecastingInferenceInput, ForecastingMetadataInput, PredictOutput
 from .service_handler import ForecastingServiceHandler
@@ -47,8 +48,8 @@ def add_routes(self, app):
         app.include_router(self.router)
 
     def _modelspec(self, model_id: str):
-        model_path = TSFM_MODEL_DIR / model_id
-        if not model_path.exists():
+        model_path = resolve_model_path(TSFM_MODEL_DIR, model_id)
+        if not model_path:
             raise HTTPException(status_code=404, detail=f"model {model_id} not found.")
         handler, e = ForecastingServiceHandler.load(model_id=model_id, model_path=model_path)
         if handler.handler_config:
@@ -84,16 +85,16 @@ def forecast(self, input: ForecastingInferenceInput):
         return answer
 
     def _forecast_common(self, input_payload: ForecastingInferenceInput) -> PredictOutput:
-        model_path = TSFM_MODEL_DIR / input_payload.model_id
+        model_path = resolve_model_path(TSFM_MODEL_DIR, input_payload.model_id)
 
-        if not model_path.is_dir():
+        if not model_path:
             LOGGER.info(f"Could not find model at path: {model_path}")
             if TSFM_ALLOW_LOAD_FROM_HF_HUB:
                 model_path = input_payload.model_id
                 LOGGER.info(f"Using HuggingFace Hub: {model_path}")
             else:
                 return None, RuntimeError(
-                    f"Could not load model {input_payload.model_id} from {TSFM_MODEL_DIR.as_posix()}. If trying to load directly from the HuggingFace Hub please ensure that `TSFM_ALLOW_LOAD_FROM_HF_HUB=1`"
+                    f"Could not load model {input_payload.model_id} from {TSFM_MODEL_DIR}. If trying to load directly from the HuggingFace Hub please ensure that `TSFM_ALLOW_LOAD_FROM_HF_HUB=1`"
                 )
 
         handler, e = ForecastingServiceHandler.load(model_id=input_payload.model_id, model_path=model_path)

services/tests/test_dirutil.py

@@ -0,0 +1,17 @@
+import os
+import tempfile
+from pathlib import Path
+
+from tsfminference.dirutil import resolve_model_path
+
+
+def test_resolve_model_path():
+    with tempfile.TemporaryDirectory() as dir1:
+        with tempfile.TemporaryDirectory() as dir2:
+            dirpath = f"{dir1}:{dir2}"
+            os.mkdir(Path(dir1) / "amodel")
+            assert resolve_model_path(dirpath, "amodel") == Path(dir1) / "amodel"
+            assert resolve_model_path(dirpath, "foobar") is None
+            assert resolve_model_path("fzbatt:zap", "amodel") is None
+            os.mkdir(Path(dir2) / "anewmodel")
+            assert resolve_model_path(dirpath, "anewmodel") == Path(dir2) / "anewmodel"

tests/toolkit/test_get_model.py

@@ -66,3 +66,24 @@ def test_get_model():
             model = get_model(model_path=mp, context_length=cl, prediction_length=fl)
             assert model.config.prediction_length == fl
             assert model.config.context_length == cl
+
+    mp = "ibm/ttm-research-r2"
+    for cl in range(1, 2000, 500):
+        for fl in range(1, 900, 90):
+            model = get_model(model_path=mp, context_length=cl, prediction_length=fl)
+            if model.config.prediction_filter_length is not None:
+                assert model.config.prediction_filter_length == fl
+
+    mp = "ibm-granite/granite-timeseries-ttm-r2"
+    for cl in range(1, 2000, 500):
+        for fl in range(1, 900, 90):
+            model = get_model(model_path=mp, context_length=cl, prediction_length=fl)
+            if model.config.prediction_filter_length is not None:
+                assert model.config.prediction_filter_length == fl
+
+    mp = "ibm-granite/granite-timeseries-ttm-r1"
+    for cl in range(512, 2000, 500):
+        for fl in range(1, 720, 90):
+            model = get_model(model_path=mp, context_length=cl, prediction_length=fl)
+            if model.config.prediction_filter_length is not None:
+                assert model.config.prediction_filter_length == fl

tsfm_public/models/tinytimemixer/modeling_tinytimemixer.py

@@ -1802,6 +1802,15 @@ def forward(
         Returns:
 
         """
+        if past_values.dim() != 3:
+            raise ValueError(
+                "`past_values` must have 3 dimensions of shape `(batch_size, sequence_length, num_input_channels)`."
+            )
+        if past_values.shape[1] > self.config.context_length:
+            past_values = past_values[:, -self.config.context_length :, :]
+        elif past_values.shape[1] < self.config.context_length:
+            raise ValueError("Context length in `past_values` is shorter that TTM context_length.")
+
         if self.loss == "mse":
             loss = nn.MSELoss(reduction="mean")
         elif self.loss == "mae":

tsfm_public/toolkit/get_model.py

@@ -48,6 +48,7 @@ def get_model(
     context_length: int = None,
     prediction_length: int = None,
     freq_prefix_tuning: bool = None,
+    force_return: bool = True,
     **kwargs,
 ):
     """
@@ -93,28 +94,51 @@ def get_model(
             with open(os.path.join(config_dir, "ttm.yaml"), "r") as file:
                 model_revisions = yaml.safe_load(file)
 
-            if prediction_length <= 96:
-                selected_prediction_length = 96
-            elif prediction_length <= 192:
-                selected_prediction_length = 192
-            elif prediction_length <= 336:
-                selected_prediction_length = 336
-            elif prediction_length <= 720:
-                selected_prediction_length = 720
+            max_supported_horizon = SUPPORTED_LENGTHS[model_path_type]["FL"][-1]
+            if prediction_length > max_supported_horizon:
+                if force_return:
+                    selected_prediction_length = max_supported_horizon
+                    LOGGER.warning(
+                        f"The requested forecast horizon is greater than the maximum supported horizon ({max_supported_horizon}). Returning TTM model with horizon {max_supported_horizon} since `force_return=True`."
+                    )
+                else:
+                    raise ValueError(f"Currently supported maximum prediction_length = {max_supported_horizon}")
             else:
-                raise ValueError("Currently supported maximum prediction_length = 720")
+                for h in SUPPORTED_LENGTHS[model_path_type]["FL"]:
+                    if prediction_length <= h:
+                        selected_prediction_length = h
+                        break
 
-            LOGGER.info(f"Selected prediction_length = {selected_prediction_length}")
+            LOGGER.info(f"Selected TTM `prediction_length` = {selected_prediction_length}")
 
-            if selected_prediction_length != prediction_length:
+            if selected_prediction_length > prediction_length:
                 prediction_filter_length = prediction_length
                 LOGGER.warning(
-                    f"Requested `prediction_length` ({prediction_length}) is not exactly equal to any of the available TTM prediction lengths.\n\
-                    Hence, TTM will forecast using the `prediction_filter_length` argument to provide the requested prediction length.\n\
-                    Supported context lengths (CL) and forecast/prediction lengths (FL) for Model Card: {model_path} are\n\
-                    {SUPPORTED_LENGTHS[model_path_type]}"
+                    f"Requested `prediction_length` ({prediction_length}) is not exactly equal to any of the available TTM prediction lengths. Hence, TTM will forecast using the `prediction_filter_length` argument to provide the requested prediction length. Supported context lengths (CL) and forecast/prediction lengths (FL) for Model Card: {model_path} are {SUPPORTED_LENGTHS[model_path_type]}"
                 )
 
+            # Choose closest context length
+            available_context_lens = sorted(SUPPORTED_LENGTHS[model_path_type]["CL"], reverse=True)
+            selected_context_length = None
+            for cl in available_context_lens:
+                if cl <= context_length:
+                    selected_context_length = cl
+                    if cl < context_length:
+                        LOGGER.warning(
+                            f"Selecting TTM context length ({selected_context_length}) < Requested context length ({context_length} since exact match was not found.)"
+                        )
+                    break
+            if selected_context_length is None:
+                if force_return:
+                    selected_context_length = available_context_lens[-1]
+                    LOGGER.warning(
+                        f"Requested context length is too short. Requested = {context_length}. Available lengths for model_type = {model_path_type} are: {available_context_lens}. Returning the shortest context length model possible since `force_return=True`. Data needs to be handled properly, and it can affect the performance!"
+                    )
+                else:
+                    raise ValueError(
+                        f"Requested context length is too short. Requested = {context_length}. Available lengths for model_type = {model_path_type} are: {available_context_lens}. To return the shortest context length model possible, set `force_return=True`."
+                    )
+
             if freq_prefix_tuning is None:
                 # Default model preference (freq / nofreq)
                 if model_path_type == 1 or model_path_type == 2:  # for granite use nofreq models
@@ -124,8 +148,8 @@ def get_model(
                 else:
                     freq_prefix = None
             else:
-                raise Exception(
-                    "In current implementation, set freq_prefix_tuning to None for automatic model selection accordingly.."
+                raise ValueError(
+                    "In the current implementation, set `freq_prefix_tuning` to None for automatic model selection accordingly."
                 )
                 if freq_prefix_tuning:
                     freq_prefix = "freq"
@@ -135,11 +159,11 @@ def get_model(
             try:
                 if model_path_type == 1 or model_path_type == 2:
                     ttm_model_revision = model_revisions["ibm-granite-models"][
-                        f"r{model_path_type}-{context_length}-{selected_prediction_length}-{freq_prefix}"
+                        f"r{model_path_type}-{selected_context_length}-{selected_prediction_length}-{freq_prefix}"
                     ]["revision"]
                 elif model_path_type == 3:
                     ttm_model_revision = model_revisions["research-use-models"][
-                        f"r2-{context_length}-{selected_prediction_length}-{freq_prefix}"
+                        f"r2-{selected_context_length}-{selected_prediction_length}-{freq_prefix}"
                     ]["revision"]
                 else:
                     raise Exception(
@@ -149,9 +173,10 @@ def get_model(
                 raise ValueError(
                     f"Model not found, possibly because of wrong context_length. Supported context lengths (CL) and forecast/prediction lengths (FL) for Model Card: {model_path} are {SUPPORTED_LENGTHS[model_path_type]}"
                 )
+        else:
+            prediction_filter_length = prediction_length
 
         # Load model
-
         model = TinyTimeMixerForPrediction.from_pretrained(
             model_path,
             revision=ttm_model_revision,