Added SageMaker batch transform support.

src/gluonts/model/forecast.py

@@ -610,9 +610,12 @@ class OutputType(str, Enum):
 
 
 class Config(pydantic.BaseModel):
- num_eval_samples: int = pydantic.Schema(..., alias="num_samples")
- output_types: Set[OutputType]
- quantiles: List[str] # FIXME: validate list elements
+ num_eval_samples: int = pydantic.Schema(100, alias="num_samples")
+ output_types: Set[OutputType] = {"quantiles", "mean"}
+ # FIXME: validate list elements
+ quantiles: List[str] = ["0.1", "0.5", "0.9"]
 
  class Config:
  allow_population_by_alias = True
+ # store additional fields
+ extra = "allow"

src/gluonts/shell/__main__.py

@@ -113,18 +113,17 @@ def serve_command(
  logging.info("Run 'serve' command")
 
  if not force_static and forecaster is not None:
- gunicorn_app = serve.make_gunicorn_app(
- env=None,
- forecaster_type=forecaster_type_by_name(forecaster),
- settings=Settings(),
+ forecaster_type: Optional[Forecaster] = forecaster_type_by_name(
+ forecaster
  )
  else:
- gunicorn_app = serve.make_gunicorn_app(
- env=ServeEnv(Path(data_path)),
- forecaster_type=None,
- settings=Settings(),
- )
+ forecaster_type = None
 
+ gunicorn_app = serve.make_gunicorn_app(
+ env=ServeEnv(Path(data_path)),
+ forecaster_type=forecaster_type,
+ settings=Settings(),
+ )
  gunicorn_app.run()
 
 

src/gluonts/shell/sagemaker/__init__.py

@@ -11,13 +11,15 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
+import os
 import json
 from pathlib import Path
 from typing import Dict, Optional
 
 from pydantic import BaseModel
 
 from gluonts.dataset.common import FileDataset, MetaData
+from gluonts.model.forecast import Config as ForecastConfig
 from gluonts.support.util import map_dct_values
 from .params import decode_sagemaker_parameters
 from .path import ServePaths, TrainPaths
@@ -44,9 +46,20 @@ def __init__(self, path: Path = Path("/opt/ml")) -> None:
 
 
 class ServeEnv:
+ path: ServePaths
+ batch_config: Optional[ForecastConfig]
+
  def __init__(self, path: Path = Path("/opt/ml")) -> None:
  self.path = ServePaths(path)
 
+ batch_transform = os.environ.get("SAGEMAKER_BATCH", "false") == "true"
+ if batch_transform:
+ self.batch_config = ForecastConfig.parse_raw(
+ os.environ["INFERENCE_CONFIG"]
+ )
+ else:
+ self.batch_config = None
+
 
 def _load_inputdataconfig(
  inputdataconfig: Path