Batch transform debug (awslabs#404)

* Added debug to batch-transform.

src/gluonts/shell/serve/app.py

@@ -121,18 +121,35 @@ def invocations() -> Response:
 
 
 def batch_inference_invocations(predictor_factory, configuration) -> Flask:
+ DEBUG = configuration.dict().get("DEBUG")
  predictor = predictor_factory({"configuration": configuration.dict()})
 
  def invocations() -> Response:
  request_data = request.data.decode("utf8").strip()
  instances = list(map(json.loads, request_data.splitlines()))
+ predictions = []
 
- predictions = handle_predictions(predictor, instances, configuration)
+ # we have to take this as the initial start-time since the first
+ # forecast is produced before the loop in predictor.predict
+ start = time.time()
 
- lines = list(map(json.dumps, map(jsonify_floats, predictions)))
+ forecast_iter = predictor.predict(
+ ListDataset(instances, predictor.freq),
+ num_eval_samples=configuration.num_eval_samples,
+ )
+
+ for forecast in forecast_iter:
+ end = time.time()
+ prediction = forecast.as_json_dict(configuration)
 
- # force line break at the end
- lines.append("")
+ if DEBUG:
+ prediction["debug"] = {"timing": end - start}
+
+ predictions.append(prediction)
+
+ start = time.time()
+
+ lines = list(map(json.dumps, map(jsonify_floats, predictions)))
  return Response("\n".join(lines), mimetype="application/jsonlines")
 
  return invocations

src/gluonts/testutil/shell.py

@@ -93,7 +93,8 @@ def batch_invocations(
  url=self.url("/invocations"), data="\n".join(instances)
  )
 
- assert response.status_code == 200
+ if response.status_code != 200:
+ raise RuntimeError(response.content.decode("utf-8"))
 
  predictions = list(map(json.loads, response.text.splitlines()))
  assert len(predictions) == len(instances)