Backports for v0.14.0rc2 (#3032)

* Fix WaveNet inputs (#3022)

* Fix version range for lightning (#3023)

* Allow VCS install using GLUONTS_FALLBACK_VERSION (#3028)

* Make `mypy` checks opt-out instead of opt-in, fix several type issues (#3027)

Co-authored-by: Oleksandr Shchur <shchuro@amazon.com>

* No API docs for nursery. (#3030)

* Fix: #3030. (#3031)

* Add support for Pydantic v1 and v2. (#3026)

---------

Co-authored-by: Oleksandr Shchur <shchuro@amazon.com>
Co-authored-by: ddelange <14880945+ddelange@users.noreply.github.com>
Co-authored-by: Jasper <schjaspe@amazon.de>

.github/workflows/style_type_checks.yml

@@ -22,6 +22,7 @@ jobs:
  pip install click black mypy
  pip install types-python-dateutil
  pip install types-waitress
+ pip install types-PyYAML
  - name: Style and type checks
  run: |
  just black

docs/Makefile

@@ -8,21 +8,26 @@ SPHINXBUILD ?= sphinx-build
 SOURCEDIR = .
 BUILDDIR = _build
 
-APIDOC = sphinx-apidoc
-APIDOC_OPTS = --implicit-namespaces --separate --module-first
-APIDOC_ROOT = gluonts
-
 # Put it first so that "make" without argument is like "make help".
 help:
  @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
 .PHONY: help Makefile
 
 apidoc:
- @rm -Rf api/$(APIDOC_ROOT)
- @$(APIDOC) $(APIDOC_OPTS) -o api/$(APIDOC_ROOT) ../src/$(APIDOC_ROOT) setup* */bin/* test docs *pycache*
- @rm -Rf api/$(APIDOC_ROOT)/modules.rst
- @sed -i"" -e "s/$(APIDOC_ROOT) package/API Docs/" api/$(APIDOC_ROOT)/$(APIDOC_ROOT).rst
+ @rm -Rf api/gluonts
+
+ @sphinx-apidoc \
+ --implicit-namespaces \
+ --separate \
+ --module-first \
+ -o api/gluonts \
+ ../src/gluonts \
+ ../src/gluonts/nursery/* \
+ ../src/gluonts/pydantic.py
+
+ @rm -Rf api/gluonts/modules.rst
+ @sed -i"" -e "s/gluonts package/API Docs/" api/gluonts/gluonts.rst
 
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).

requirements/requirements-pytorch.txt

@@ -1,9 +1,6 @@
 torch>=1.9,<3
-lightning>=1.8,<2.2
+lightning>=2.0,<2.2
 # Capping `lightning` does not cap `pytorch_lightning`, so we cap manually
-pytorch_lightning>=1.8,<2.2
-# Need to pin protobuf (for now)
-# See: https://github.com/PyTorchLightning/pytorch-lightning/issues/13159
-protobuf~=3.19.0
+pytorch_lightning>=2.0,<2.2
 scipy~=1.10; python_version > "3.7.0"
-scipy~=1.7.3; python_version <= "3.7.0"
+scipy~=1.7.3; python_version <= "3.7.0"

requirements/requirements.txt

@@ -1,6 +1,6 @@
 numpy~=1.16
 pandas>=1.0,<3
-pydantic~=1.7
+pydantic>=1.7,<3
 tqdm~=4.23
 toolz~=0.10
 

setup.py

@@ -57,18 +57,24 @@ def run(self):
  # otherwise a module-not-found error is thrown
  import mypy.api
 
- folders = [
- str(p.parent.resolve()) for p in ROOT.glob("src/**/.typesafe")
+ excluded_folders = [
+ str(p.parent.relative_to(ROOT)) for p in ROOT.glob("src/**/.typeunsafe")
  ]
 
- print(
- "The following folders contain a `.typesafe` marker file "
- "and will be type-checked with `mypy`:"
- )
- for folder in folders:
+ if len(excluded_folders) > 0:
+ print(
+ "The following folders contain a `.typeunsafe` marker file "
+ "and will *not* be type-checked with `mypy`:"
+ )
+ for folder in excluded_folders:
  print(f" {folder}")
 
- std_out, std_err, exit_code = mypy.api.run(folders)
+ args = [str(ROOT / "src")]
+ for folder in excluded_folders:
+ args.append("--exclude")
+ args.append(folder)
+
+ std_out, std_err, exit_code = mypy.api.run(args)
 
  print(std_out, file=sys.stdout)
  print(std_err, file=sys.stderr)
@@ -78,7 +84,7 @@ def run(self):
  f"""
  Mypy command
 
- mypy {" ".join(folders)}
+ mypy {" ".join(args)}
 
  returned a non-zero exit code. Fix the type errors listed above
  and then run

src/gluonts/core/component.py

@@ -19,10 +19,15 @@
 from typing import Any, Type, TypeVar
 
 import numpy as np
-from pydantic import BaseConfig, BaseModel, ValidationError, create_model
 
 from gluonts.core import fqname_for
 from gluonts.exceptions import GluonTSHyperparametersError
+from gluonts.pydantic import (
+ BaseConfig,
+ BaseModel,
+ ValidationError,
+ create_model,
+)
 
 
 logger = logging.getLogger(__name__)
@@ -252,7 +257,7 @@ def validated(base_model=None):
  >>> c = ComplexNumber(y=None)
  Traceback (most recent call last):
  ...
- pydantic.error_wrappers.ValidationError: 1 validation error for
+ pydantic.v1.error_wrappers.ValidationError: 1 validation error for
  ComplexNumberModel
  y
  none is not an allowed value (type=type_error.none.not_allowed)
@@ -262,7 +267,7 @@ def validated(base_model=None):
  accessed through the ``Model`` attribute of the decorated initializer.
 
  >>> ComplexNumber.__init__.Model
- <class 'pydantic.main.ComplexNumberModel'>
+ <class 'pydantic.v1.main.ComplexNumberModel'>
 
  The Pydantic model is synthesized automatically from on the parameter
  names and types of the decorated initializer. In the ``ComplexNumber``

src/gluonts/core/serde/_base.py

@@ -20,9 +20,8 @@
 
 from toolz.dicttoolz import valmap
 
-from pydantic import BaseModel
-
 from gluonts.core import fqname_for
+from gluonts.pydantic import BaseModel
 
 bad_type_msg = textwrap.dedent(
  """

src/gluonts/core/serde/_dataclass.py

@@ -24,10 +24,8 @@
  TypeVar,
 )
 
-import pydantic
-import pydantic.dataclasses
-
 from gluonts.itertools import select
+from gluonts.pydantic import pydantic, dataclass as pydantic_dataclass
 
 T = TypeVar("T")
 
@@ -152,7 +150,7 @@ class Config:
  arbitrary_types_allowed = True
 
  # make `cls` a dataclass
- pydantic.dataclasses.dataclass(
+ pydantic_dataclass(
  init=init,
  repr=repr,
  eq=eq,

src/gluonts/core/settings.py

@@ -76,8 +76,7 @@ def fn(debug):
 from operator import attrgetter
 from typing import Any
 
-import pydantic
-from pydantic.utils import deep_update
+from gluonts.pydantic import pydantic, deep_update
 
 
 class ListElement:

src/gluonts/dataset/common.py

@@ -22,19 +22,17 @@
 import pandas as pd
 from pandas.tseries.frequencies import to_offset
 
-import pydantic
-
 from gluonts import json
 from gluonts.itertools import Cached, Map
 from gluonts.dataset.field_names import FieldName
 from gluonts.dataset.schema import Translator
 from gluonts.exceptions import GluonTSDataError
+from gluonts.pydantic import pydantic
 
 
 from . import Dataset, DatasetCollection, DataEntry, DataBatch # noqa
 from . import jsonl, DatasetWriter
 
-
 arrow: Optional[ModuleType]
 
 try:

src/gluonts/dataset/loader.py

@@ -15,7 +15,6 @@
 from typing import Callable, Iterable, Optional
 
 import numpy as np
-from pydantic import BaseModel
 
 from gluonts.dataset import DataBatch, Dataset
 from gluonts.itertools import (
@@ -25,6 +24,7 @@
  batcher,
  rows_to_columns,
 )
+from gluonts.pydantic import BaseModel
 from gluonts.transform import (
  AdhocTransform,
  Identity,

src/gluonts/ev/aggregations.py

@@ -48,6 +48,8 @@ class Sum(Aggregation):
  partial_result: Optional[Union[List[np.ndarray], np.ndarray]] = None
 
  def step(self, values: np.ndarray) -> None:
+ assert self.axis is None or isinstance(self.axis, tuple)
+
  summed_values = np.ma.sum(values, axis=self.axis)
 
  if self.axis is None or 0 in self.axis:
@@ -57,9 +59,12 @@ def step(self, values: np.ndarray) -> None:
  else:
  if self.partial_result is None:
  self.partial_result = []
+ assert isinstance(self.partial_result, list)
  self.partial_result.append(summed_values)
 
  def get(self) -> np.ndarray:
+ assert self.axis is None or isinstance(self.axis, tuple)
+
  if self.axis is None or 0 in self.axis:
  return np.ma.copy(self.partial_result)
 
@@ -85,6 +90,8 @@ class Mean(Aggregation):
  n: Optional[Union[int, np.ndarray]] = None
 
  def step(self, values: np.ndarray) -> None:
+ assert self.axis is None or isinstance(self.axis, tuple)
+
  if self.axis is None or 0 in self.axis:
  summed_values = np.ma.sum(values, axis=self.axis)
  if self.partial_result is None:
@@ -101,10 +108,14 @@ def step(self, values: np.ndarray) -> None:
  self.partial_result = []
 
  mean_values = np.ma.mean(values, axis=self.axis)
+ assert isinstance(self.partial_result, list)
  self.partial_result.append(mean_values)
 
  def get(self) -> np.ndarray:
+ assert self.axis is None or isinstance(self.axis, tuple)
+
  if self.axis is None or 0 in self.axis:
+ assert isinstance(self.partial_result, np.ndarray)
  return self.partial_result / self.n
 
  return np.ma.concatenate(self.partial_result)

src/gluonts/ev/metrics.py

@@ -139,6 +139,9 @@ def __call__(self, axis: Optional[int] = None) -> Metric:
 
 
 class BaseMetricDefinition:
+ def __call__(self, axis):
+ raise NotImplementedError()
+
  def __add__(self, other) -> MetricDefinitionCollection:
  if isinstance(other, MetricDefinitionCollection):
  return other + self
@@ -154,7 +157,7 @@ def add(self, *others):
 
 @dataclass
 class MetricDefinitionCollection(BaseMetricDefinition):
- metrics: List[MetricDefinition]
+ metrics: List[BaseMetricDefinition]
 
  def __call__(self, axis: Optional[int] = None) -> MetricCollection:
  return MetricCollection([metric(axis=axis) for metric in self.metrics])

src/gluonts/ev/stats.py

@@ -18,6 +18,7 @@
 
 def num_masked_target_values(data: Dict[str, np.ndarray]) -> np.ndarray:
  if np.ma.isMaskedArray(data["label"]):
+ assert isinstance(data["label"], np.ma.MaskedArray)
  return data["label"].mask.astype(float)
  else:
  return np.zeros(data["label"].shape)

src/gluonts/exceptions.py

@@ -13,7 +13,7 @@
 
 from typing import Any
 
-from pydantic.error_wrappers import ValidationError, display_errors
+from gluonts.pydantic import ValidationError, display_errors
 
 
 class GluonTSException(Exception):

src/gluonts/ext/naive_2/_predictor.py

@@ -21,7 +21,7 @@
 from gluonts.model.predictor import RepresentablePredictor
 
 
-def seasonality_test(past_ts_data: np.array, season_length: int) -> bool:
+def seasonality_test(past_ts_data: np.ndarray, season_length: int) -> bool:
  """
  Test the time series for seasonal patterns by performing a 90% auto-
  correlation test.

src/gluonts/ext/r_forecast/_hierarchical_predictor.py

@@ -32,7 +32,7 @@
  "erm",
 ]
 
-HIERARCHICAL_SAMPLE_FORECAST_METHODS = [] # TODO: Add `depbu_mint`.
+HIERARCHICAL_SAMPLE_FORECAST_METHODS: List[str] = [] # TODO: Add `depbu_mint`.
 
 SUPPORTED_HIERARCHICAL_METHODS = (
  HIERARCHICAL_POINT_FORECAST_METHODS + HIERARCHICAL_SAMPLE_FORECAST_METHODS
@@ -200,6 +200,8 @@ def _get_r_forecast(self, data: Dict) -> Dict:
  else:
  hier_ts = self._hts_pkg.gts(y_bottom_ts, groups=nodes)
 
+ assert isinstance(self.params["num_samples"], int)
+
  forecast = self._r_method(hier_ts, r_params)
 
  all_forecasts = list(forecast)
@@ -244,7 +246,7 @@ def _forecast_dict_to_obj(
  forecast_dict: Dict,
  forecast_start_date: pd.Timestamp,
  item_id: Optional[str],
- info: Dict,
+ info: Optional[Dict],
  ) -> SampleForecast:
  samples = np.array(forecast_dict["samples"])
 

src/gluonts/ext/r_forecast/_predictor.py

@@ -91,7 +91,7 @@ def __init__(
  self,
  freq: str,
  prediction_length: int,
- period: int = None,
+ period: Optional[int] = None,
  trunc_length: Optional[int] = None,
  save_info: bool = False,
  r_file_prefix: str = "",
@@ -215,7 +215,7 @@ def _forecast_dict_to_obj(
  forecast_dict: Dict,
  forecast_start_date: pd.Timestamp,
  item_id: Optional[str],
- info: Dict,
+ info: Optional[Dict],
  ) -> Forecast:
  """
  Returns object of type `gluonts.model.Forecast`.

src/gluonts/ext/r_forecast/_univariate_predictor.py

@@ -11,7 +11,7 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
-from typing import Dict, Optional
+from typing import Dict, Optional, Any
 
 import numpy as np
 import pandas as pd
@@ -79,7 +79,7 @@ def __init__(
  freq: str,
  prediction_length: int,
  method_name: str = "ets",
- period: int = None,
+ period: Optional[int] = None,
  trunc_length: Optional[int] = None,
  save_info: bool = False,
  params: Dict = dict(),
@@ -100,7 +100,7 @@ def __init__(
  self.method_name = method_name
  self._r_method = self._robjects.r[method_name]
 
- self.params = {
+ self.params: Dict[str, Any] = {
  "prediction_length": self.prediction_length,
  "frequency": self.period,
  }
@@ -185,7 +185,7 @@ def _forecast_dict_to_obj(
  forecast_dict: Dict,
  forecast_start_date: pd.Timestamp,
  item_id: Optional[str],
- info: Dict,
+ info: Optional[Dict],
  ) -> QuantileForecast:
  stats_dict = {"mean": forecast_dict["mean"]}