address nits

databricks · Oct 7, 2024 · ac27c97 · ac27c97
1 parent 0060de6
commit ac27c97
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/runtime/databricks/automl_runtime/forecast/model.py b/runtime/databricks/automl_runtime/forecast/model.py
@@ -57,7 +57,7 @@ def mlflow_forecast_log_model(forecast_model: ForecastModel,
     :param forecast_model: Forecast model wrapper
     :param sample_input: sample input Dataframes for model inference
     """
-    # TODO: we should not be logging without a signature since it cannot be registered to UC then
+    # TODO: [ML-46185] we should not be logging without a signature since it cannot be registered to UC then
     try:
         signature = forecast_model.infer_signature(sample_input)
     except Exception: # noqa

diff --git a/runtime/tests/automl_runtime/forecast/deepar/model_test.py b/runtime/tests/automl_runtime/forecast/deepar/model_test.py
@@ -25,7 +25,9 @@
 from gluonts.torch.model.predictor import PyTorchPredictor
 
 from databricks.automl_runtime.forecast.deepar.model import (
-    DeepARModel, mlflow_deepar_log_model, DEEPAR_ADDITIONAL_PIP_DEPS
+    DeepARModel, 
+    mlflow_deepar_log_model, 
+    DEEPAR_ADDITIONAL_PIP_DEPS
 )
 
 
@@ -106,7 +108,7 @@ def test_model_save_and_load_single_series(self):
         run_id = run.info.run_id
 
         # check if all additional dependencies are logged
-        self._check_requirements(run_id)
+        self.check_requirements(run_id)
 
         # load the model and predict
         loaded_model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")

diff --git a/runtime/tests/automl_runtime/forecast/pmdarima/model_test.py b/runtime/tests/automl_runtime/forecast/pmdarima/model_test.py
@@ -26,8 +26,11 @@
 from pmdarima.arima import ARIMA
 
 from databricks.automl_runtime.forecast.pmdarima.model import (
-    ArimaModel, MultiSeriesArimaModel, AbstractArimaModel, \
-    mlflow_arima_log_model, ARIMA_ADDITIONAL_PIP_DEPS
+    ArimaModel, 
+    MultiSeriesArimaModel, 
+    AbstractArimaModel,
+    mlflow_arima_log_model, 
+    ARIMA_ADDITIONAL_PIP_DEPS,
 )