Repository: awslabs/gluonts

Issue:
Add forward-fields support to shell for batch-transform
When doing batch transform, one might want to store additional fields such as `item_id` alongside the prediction to be able to simply match input and output values.

Currently, we just ignore all additional values.


Discussion:
This might be obsolete, since batch transform might be able to do it itself:

https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_DataProcessing.html
If I have this json dataset format:

`{"item_id": 1000, "store_id": 20, "start": "2011-01-29", "feat_static_cat": [0, 0, 0, 0], "feat_dynamic_real": [[30, 30, 30, 30, 30, 30 ..], [4, 4, 4, 4, 4, 4, 4, ...], [4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4....], [2, 2, 2, 2, 2, 2, 2, 2, 2, 2...], [0, 0, 0, 1, 1, 1...], [2.240234375, 2.240234375, 2.240234375, 2.240234375 ...]], "target": [5, 5, 1, 1, 0, 3, 3, 4, 9, 4, 7, 1, 2, 4, 14 ...]}`

Is there someway to access `store_id` on forecast inference result ?

`[gluonts.model.forecast.SampleForecast(freq="D", info=None, item_id="1000", samples=numpy.array()`

Resolving patch:
diff --git a/src/gluonts/dataset/common.py b/src/gluonts/dataset/common.py
--- a/src/gluonts/dataset/common.py
+++ b/src/gluonts/dataset/common.py
@@ -505,6 +505,8 @@ def serialize_field(field):
             field = field.astype(np.object_)
             field[nan_ix] = "NaN"
             return field.tolist()
+        if isinstance(field, (int, float)):
+            return field
         return str(field)
 
     return {k: serialize_field(v) for k, v in data.items() if v is not None}
diff --git a/src/gluonts/shell/serve/__init__.py b/src/gluonts/shell/serve/__init__.py
--- a/src/gluonts/shell/serve/__init__.py
+++ b/src/gluonts/shell/serve/__init__.py
@@ -14,7 +14,7 @@
 import logging
 import multiprocessing
 from ipaddress import IPv4Address
-from typing import Optional, Type, Union
+from typing import List, Optional, Type, Union
 
 from flask import Flask
 from gunicorn.app.base import BaseApplication
@@ -55,6 +55,7 @@ class Config:
     gluonts_batch_fallback_predictor: str = (
         "gluonts.model.trivial.mean.MeanPredictor"
     )
+    gluonts_forward_fields: List[str] = []
 
     sagemaker_batch: bool = False
     sagemaker_batch_strategy: str = "SINGLE_RECORD"
diff --git a/src/gluonts/shell/serve/app.py b/src/gluonts/shell/serve/app.py
--- a/src/gluonts/shell/serve/app.py
+++ b/src/gluonts/shell/serve/app.py
@@ -245,6 +245,10 @@ def invocations() -> Response:
 
         log_scored(when=end_time)
 
+        for forward_field in settings.gluonts_forward_fields:
+            for input_item, prediction in zip(dataset, predictions):
+                prediction[forward_field] = input_item.get(forward_field)
+
         lines = list(map(json.dumps, map(jsonify_floats, predictions)))
         return Response("\n".join(lines), mimetype="application/jsonlines")
 


Test patch:
diff --git a/test/shell/test_shell.py b/test/shell/test_shell.py
--- a/test/shell/test_shell.py
+++ b/test/shell/test_shell.py
@@ -100,6 +100,7 @@ def batch_transform(monkeypatch, train_env):
     }
 
     monkeypatch.setenv("INFERENCE_CONFIG", json.dumps(inference_config))
+    monkeypatch.setenv("GLUONTS_FORWARD_FIELDS", json.dumps(["foo"]))
     return inference_config
 
 
@@ -233,11 +234,14 @@ def test_dynamic_batch_shell(
     assert execution_parameters["MaxPayloadInMB"] == 6
 
     for entry in train_env.datasets["train"]:
+        entry["foo"] = 42
         forecast = dynamic_server.batch_invocations([entry])[0]
 
         for output_type in batch_transform["output_types"]:
             assert output_type in forecast
 
+        assert forecast["foo"] == 42
+
         act_mean = np.array(forecast["mean"])
         act_samples = np.array(forecast["samples"])