TSFM_MODE_DIR can be multiple dirs under byom

ibm-granite · Dec 10, 2024 · a6f8175 · a6f8175
1 parent 573de88
commit a6f8175
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/services/inference/tests/test_inference_lib.py b/services/inference/tests/test_inference_lib.py
@@ -15,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,
@@ -37,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