UBN tests

Summary:
A bunch of tests are flaky or failing. This seems to be holding up an FBPkg release (which fired a UBN), and then is also just creating a lot of tasks for flaky or failing tests.

This diff was originally used to push to py sci stack v2, and update all the broken/tests and code, however I've since split it into that v2 upgrade (next diff in stack) and this smaller diff which just resolves already-broken tests.

Failing tests have two kinds of causes: test code is broken, or test is super heavy weight and timing out always. Flaky tests seem to just be heavyweight, and their success seems to depend on whether they get enough resources to finish in 10 min or not.

First category I just fix the test or correct the code (usually some API changes in py 3.8 ).

Second category is a mixed bag: where possible I rewrote the test to mock out the expensive calls (I'm looking at you TsFeatures, which in some tests was called nearly 10 times). If mocking out the expensive part seemed infeasible, or just sufficiently daunting, I labelled the test as `heavyweight` and `long_running`, which uses more resources but hopefully stops timing out. This only ended up happening for the flaky tests, which were on the line of not timing out, so hopefully this fix should do the trick.

#buildall

Reviewed By: irumata, igorsugak

Differential Revision: D58582196

fbshipit-source-id: 04d7fec197747237a24f953c274333972a02e820

kats/compat/compat.py

@@ -13,9 +13,11 @@
     # Python < 3.8
     import importlib_metadata as metadata
 
+import logging
 from typing import Any, Callable, Type, Union
 
 import packaging
+import statsmodels
 
 from packaging import version as pv
 
@@ -60,6 +62,11 @@ def __init__(self, version: V) -> None:
             self.version: VERSION_TYPE = version
 
     def _parse(self, version: str) -> Union[pv.Version, pv.LegacyVersion]:
+        if version == "statsmodels":
+            return pv.Version(
+                statsmodels.__version__
+            )  # TODO: importlib.metadata.version is spuriously returning 0.0.0 as statsmodels version, breaking compat
+
         try:
             version = metadata.version(version)
         except metadata.PackageNotFoundError:

kats/compat/statsmodels.py

@@ -83,7 +83,7 @@ def __init__(
                 )
             if initialization_method != "missing":
                 logging.warning(
-                    "ExponentialSmoothing parameter 'initialization_method' not supported by statsmodels"
+                    f"ExponentialSmoothing parameter 'initialization_method' not supported by statsmodels, detected {version.version}"
                 )
             if initial_level is not None:
                 logging.warning(
@@ -102,21 +102,21 @@ def __init__(
                     "ExponentialSmoothing parameter 'seasonal_periods' not supported by statsmodels"
                 )
             super().__init__(
-                endog,
-                trend,
-                damped_trend,
-                seasonal,
-                seasonal_periods,
-                dates,
-                freq,
-                missing,
+                endog=endog,
+                trend=trend,
+                damped_trend=damped_trend,
+                seasonal=seasonal,
+                seasonal_periods=seasonal_periods,
+                dates=dates,
+                freq=freq,
+                missing=missing,
             )
         else:
             super().__init__(
-                endog,
-                trend,
-                damped_trend,
-                seasonal,
+                endog=endog,
+                trend=trend,
+                damped_trend=damped_trend,
+                seasonal=seasonal,
                 seasonal_periods=seasonal_periods,
                 initialization_method=initialization_method,
                 initial_level=initial_level,

kats/models/globalmodel/utils.py

@@ -11,6 +11,7 @@
 
 import numpy as np
 import torch
+from kats.compat.compat import Version
 from kats.consts import TimeSeriesData
 from kats.tsfeatures.tsfeatures import TsFeatures
 from numba import jit
@@ -21,6 +22,8 @@
 
 import pandas as pd
 
+PANDAS_VERSION = Version("pandas")
+
 """
 A module for utility functions of global models, including:
     1) Helper functions for preprocessing and calculating error metrics.
@@ -999,14 +1002,22 @@ def _get_last_date_feature(
         ans = np.zeros(n * m)
         pdt = pd.to_datetime(time[:, -1])
         indices = []
-        # compute day of week indices
-        indices.append(pdt.dayofweek.values + offset)
-        # compute bi-week indices
-        indices.append((pdt.weekofyear.values - 1) // 2 + 7 + offset)
+        if PANDAS_VERSION < "2.0.3":
+            # compute day of week indices
+            indices.append(pdt.dayofweek.values + offset)
+
+            # compute bi-week indices
+            indices.append((pdt.weekofyear.values - 1) // 2 + 7 + offset)
+        else:
+            # compute day of week indices
+            indices.append(pdt.isocalendar().day.values + offset)
+            # compute bi-week indices
+            indices.append((pdt.isocalendar().week.values - 1) // 2 + 7 + offset)
+
         # compute day of month indices
         indices.append(pdt.day.values + 6 + 27 + offset)
         indices = np.concatenate(indices)
-        ans[indices] = 1.0
+        ans[indices.astype(int)] = 1.0
         return torch.tensor(ans.reshape(n, -1), dtype=torch.get_default_dtype())
 
     @staticmethod

kats/tests/models/test_globalmodel.py

@@ -10,6 +10,7 @@
 from functools import partial
 from typing import Dict, List, Union
 from unittest import mock, TestCase
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 import pandas as pd
@@ -38,10 +39,46 @@
     S2Cell,
     split,
 )
+from kats.tsfeatures.tsfeatures import TsFeatures
 
 from parameterized.parameterized import parameterized
 
 
+def _MOCK_GET_TSFEATURES(
+    x: np.ndarray,
+    time: np.ndarray,
+) -> torch.Tensor:
+    """
+    Mocks the private method of GMFeature so that a reduced set of features is computed and the test stops timing out
+    """
+    features = []
+
+    for i in range(len(x)):
+        features.append(
+            np.log(
+                np.abs(
+                    list(
+                        # pyre-fixme[16]: `List` has no attribute `values`.
+                        TsFeatures(selected_features=["length", "mean"])
+                        .transform(
+                            TimeSeriesData(
+                                pd.DataFrame(
+                                    {"time": time[i], "value": x[i]}, copy=False
+                                ).dropna()  # NaNs would mess-up tsfeatures
+                            )
+                        )
+                        .values()
+                    )
+                )
+            )
+        )
+    features = torch.tensor(features)
+    # filter out NaN and inf
+    features[torch.isnan(features)] = 0.0
+    features[torch.isinf(features)] = 0.0
+    return features
+
+
 def get_ts(
     n: int, start_time: str, seed: int = 560, freq: str = "D", has_nans: bool = True
 ) -> TimeSeriesData:
@@ -428,10 +465,13 @@ def test_other(self) -> None:
 
 
 class GMFeatureTest(TestCase):
-    def test_gmfeature(self) -> None:
+    @patch("kats.models.globalmodel.utils.GMFeature._get_tsfeatures")
+    def test_gmfeature(self, _get_ts_features_mock: MagicMock) -> None:
         x = np.row_stack([np.abs(ts.value.values[:10]) for ts in TSs])
         time = np.row_stack([ts.time.values[:10] for ts in TSs])
 
+        _get_ts_features_mock.return_value = _MOCK_GET_TSFEATURES(x, time)
+
         gmfs = [
             GMFeature(feature_type="tsfeatures"),
             GMFeature(feature_type="last_date"),