code: def run(argv=None):
      """Helper for obtaining prediction summary."""
      parser = argparse.ArgumentParser()
      parser.add_argument(
          "--prediction_path",
          required=True,
          help=(
              "The GCS folder that contains BatchPrediction results, containing "
              "prediction.results-NNNNN-of-NNNNN files in the json format. "
              "Output will be also stored in this folder, as a file"
              "'prediction.summary.json'."
          ),
      )
      parser.add_argument(
          "--metric_fn_encoded",
          required=True,
          help=(
              "An encoded function that calculates and returns a tuple of "
              "metric(s) for a given instance (as a dictionary). It should be "
              "encoded via base64.b64encode(dill.dumps(fn, recurse=True))."
          ),
      )
      parser.add_argument(
          "--metric_keys",
          required=True,
          help=(
              "A comma-separated keys of the aggregated metric(s) in the summary "
              "output. The order and the size of the keys must match to the "
              "output of metric_fn. The summary will have an additional key, "
              "'count', to represent the total number of instances, so this flag "
              "shouldn't include 'count'."
          ),
      )
      known_args, pipeline_args = parser.parse_known_args(argv)
  
      metric_fn = dill.loads(
          base64.b64decode(known_args.metric_fn_encoded)
      )
      if not callable(metric_fn):
          raise ValueError(
              "--metric_fn_encoded must be an encoded callable."
          )
      metric_keys = known_args.metric_keys.split(",")
  
      with beam.Pipeline(
          options=beam.pipeline.PipelineOptions(pipeline_args)
      ) as pipe:
          prediction_result_pattern = os.path.join(
              known_args.prediction_path, "prediction.results-*-of-*"
          )
          prediction_summary_path = os.path.join(
              known_args.prediction_path, "prediction.summary.json"
          )
          # This is apache-beam ptransform's convention
          _ = (
              pipe
              | "ReadPredictionResult"
              >> beam.io.ReadFromText(
                  prediction_result_pattern, coder=JsonCoder()
              )
              | "Summary" >> MakeSummary(metric_fn, metric_keys)
              | "Write"
              >> beam.io.WriteToText(
                  prediction_summary_path,
                  shard_name_template="",  # without trailing -NNNNN-of-NNNNN.
                  coder=JsonCoder(),
              )
          )
  
file_path: airflow/airflow/providers/google/cloud/utils/mlengine_prediction_summary.py
imports: from __future__ import annotations
  import argparse
  import base64
  import json
  import logging
  import os
  import apache_beam
  import dill
  from apache_beam.coders.coders import Coder
label_desc: Use image node_id 2 for calling a global function with example usage: run(argv) without return types