chronos handler to support chronos-bolt-* models

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/finetuning/tsfmfinetuning/finetuning.py

@@ -7,6 +7,7 @@
 from typing import Any, Dict, Tuple, Union
 
 import pandas as pd
+import torch
 from fastapi import APIRouter, HTTPException
 from starlette import status
 from transformers import EarlyStoppingCallback, Trainer, TrainingArguments, set_seed
@@ -215,7 +216,7 @@ def _finetuning_common(
             metric_for_best_model="eval_loss",  # Metric to monitor for early stopping
             greater_is_better=False,  # For loss
             label_names=["future_values"],
-            use_cpu=True,  # only needed for testing on Mac :(
+            use_cpu=not torch.cuda.is_available(),
         )
 
         callbacks = []

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/pyproject.toml

@@ -49,7 +49,7 @@ urllib3 = { version = ">=1.26.19,<2" } # see https://github.com/urllib3/urllib3/
 aiohttp = { version = ">=3.10.11" }
 
 # ***********Chronos*********
-chronos = { git = "https://github.com/amazon-science/chronos-forecasting.git" }
+chronos-forecasting = { git = "https://github.com/amazon-science/chronos-forecasting.git" }
 
 [[tool.poetry.source]]
 name = "pytorch"

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
 
@@ -13,6 +15,7 @@
 from fastapi import HTTPException
 from pytest import FixtureRequest
 from tsfminference import TSFM_CONFIG_FILE, TSFM_MODEL_DIR
+from tsfminference.dirutil import resolve_model_path
 from tsfminference.inference import InferenceRuntime
 from tsfminference.inference_payloads import (
     ForecastingInferenceInput,
@@ -35,7 +38,7 @@
 
 
 def min_context_length(model_id):
-    model_path: Path = TSFM_MODEL_DIR / model_id
+    model_path: Path = resolve_model_path(TSFM_MODEL_DIR, model_id)
     assert model_path.exists(), f"{model_path} does not exist!"
     handler, e = ForecastingServiceHandler.load(model_id=model_id, model_path=model_path)
     return handler.handler_config.minimum_context_length
@@ -112,6 +115,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/tests/test_inference_service.py

@@ -21,7 +21,8 @@
     "ttm-1536-96-r2": {"context_length": 1536, "prediction_length": 96},
     "ibm/test-patchtst": {"context_length": 512, "prediction_length": 96},
     "ibm/test-patchtsmixer": {"context_length": 512, "prediction_length": 96},
-    "chronos-t5-tiny": {"context_length": 512, "prediction_length": 96},
+    "chronos-t5-tiny": {"context_length": 512, "prediction_length": 16},
+    "chronos-bolt-tiny": {"context_length": 512, "prediction_length": 16},
 }
 
 
@@ -369,24 +370,28 @@ def test_zero_shot_forecast_inference(ts_data):
     assert counts["output_data_points"] == (prediction_length // 4) * len(params["target_columns"][1:])
 
 
-@pytest.mark.parametrize("ts_data", ["chronos-t5-tiny"], indirect=True)
+@pytest.mark.parametrize("ts_data", ["chronos-t5-tiny", "chronos-bolt-tiny"], indirect=True)
 def test_zero_shot_forecast_inference_chronos(ts_data):
     test_data, params = ts_data
 
     prediction_length = params["prediction_length"]
     model_id = params["model_id"]
     model_id_path: str = model_id
-
     id_columns = params["id_columns"]
+    num_samples = 10
 
     # test single
     test_data_ = test_data[test_data[id_columns[0]] == "a"].copy()
 
+    parameters = {
+        "prediction_length": params["prediction_length"],
+    }
+    if model_id == "chronos-t5-tiny":
+        parameters["num_samples"] = num_samples
+
     msg = {
         "model_id": model_id_path,
-        "parameters": {
-            "prediction_length": params["prediction_length"],
-        },
+        "parameters": parameters,
         "schema": {
             "timestamp_column": params["timestamp_column"],
             "id_columns": params["id_columns"],
@@ -400,6 +405,7 @@ def test_zero_shot_forecast_inference_chronos(ts_data):
     assert len(df_out) == 1
     assert df_out[0].shape[0] == prediction_length
 
+    # test with future data. should throw error.
     test_data_ = test_data[test_data[id_columns[0]] == "a"].copy()
     future_data = extend_time_series(
         select_by_index(test_data_, id_columns=params["id_columns"], start_index=-1),
@@ -414,9 +420,7 @@ def test_zero_shot_forecast_inference_chronos(ts_data):
 
     msg = {
         "model_id": model_id,
-        "parameters": {
-            # "prediction_length": params["prediction_length"],
-        },
+        "parameters": parameters,
         "schema": {
             "timestamp_column": params["timestamp_column"],
             "id_columns": params["id_columns"],
@@ -430,6 +434,74 @@ def test_zero_shot_forecast_inference_chronos(ts_data):
     out, _ = get_inference_response(msg)
     assert "Chronos does not support or require future exogenous." in out.text
 
+    # test multi-time series
+    num_ids = test_data[id_columns[0]].nunique()
+    test_data_ = test_data.copy()
+
+    msg = {
+        "model_id": model_id_path,
+        "parameters": parameters,
+        "schema": {
+            "timestamp_column": params["timestamp_column"],
+            "id_columns": params["id_columns"],
+            "target_columns": params["target_columns"],
+        },
+        "data": encode_data(test_data_, params["timestamp_column"]),
+        "future_data": {},
+    }
+
+    df_out, _ = get_inference_response(msg)
+
+    assert len(df_out) == 1
+    assert df_out[0].shape[0] == prediction_length * num_ids
+
+    # test multi-time series multi-id
+    multi_df = []
+    for grp in ["A", "B"]:
+        td = test_data.copy()
+        td["id2"] = grp
+        multi_df.append(td)
+    test_data_ = pd.concat(multi_df, ignore_index=True)
+    new_id_columns = id_columns + ["id2"]
+
+    num_ids = test_data_[new_id_columns[0]].nunique() * test_data_[new_id_columns[1]].nunique()
+
+    msg = {
+        "model_id": model_id_path,
+        "parameters": parameters,
+        "schema": {
+            "timestamp_column": params["timestamp_column"],
+            "id_columns": new_id_columns,
+            "target_columns": params["target_columns"],
+        },
+        "data": encode_data(test_data_, params["timestamp_column"]),
+        "future_data": {},
+    }
+
+    df_out, _ = get_inference_response(msg)
+    assert len(df_out) == 1
+    assert df_out[0].shape[0] == prediction_length * num_ids
+
+    # single series, less columns, no id
+    test_data_ = test_data[test_data[id_columns[0]] == "a"].copy()
+
+    msg = {
+        "model_id": model_id_path,
+        "parameters": parameters,
+        "schema": {
+            "timestamp_column": params["timestamp_column"],
+            "id_columns": [],
+            "target_columns": ["HULL"],
+        },
+        "data": encode_data(test_data_, params["timestamp_column"]),
+        "future_data": {},
+    }
+
+    df_out, counts = get_inference_response(msg)
+    assert len(df_out) == 1
+    assert df_out[0].shape[0] == prediction_length
+    assert df_out[0].shape[1] == 2
+
 
 @pytest.mark.parametrize("ts_data", ["ttm-r2-etth-finetuned-control"], indirect=True)
 def test_future_data_forecast_inference(ts_data):

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."
+    )