# Buggy code
def Do(
    self,
    input_dict: Dict[Text, List[types.Artifact]],
    output_dict: Dict[Text, List[types.Artifact]],
    exec_properties: Dict[Text, Any],
):
    """Overrides the tfx_pusher_executor.

    Args:
      input_dict: Input dict from input key to a list of artifacts, including:
        - model_export: exported model from trainer.
        - model_blessing: model blessing path from model_validator.
      output_dict: Output dict from key to a list of artifacts, including:
        - model_push: A list of 'ModelPushPath' artifact of size one. It will
          include the model in this push execution if the model was pushed.
      exec_properties: Mostly a passthrough input dict for
        tfx.components.Pusher.executor.  custom_config.ai_platform_serving_args
        is consumed by this class.  For the full set of parameters supported by
        Google Cloud AI Platform, refer to
        https://cloud.google.com/ml-engine/docs/tensorflow/deploying-models#creating_a_model_version.

    Returns:
      None
    Raises:
      ValueError:
        If ai_platform_serving_args is not in exec_properties.custom_config.
        If Serving model path does not start with gs://.
      RuntimeError: if the Google Cloud AI Platform training job failed.
    """
    self._log_startup(input_dict, output_dict, exec_properties)
    model_push = artifact_utils.get_single_instance(
        output_dict[tfx_pusher_executor.PUSHED_MODEL_KEY]
    )
    if not self.CheckBlessing(input_dict):
        model_push.set_int_custom_property("pushed", 0)
        return

    model_export = artifact_utils.get_single_instance(
        input_dict[tfx_pusher_executor.MODEL_KEY]
    )
    model_export_uri = model_export.uri

    exec_properties_copy = exec_properties.copy()
    custom_config = exec_properties_copy.pop("custom_config", {})
    ai_platform_serving_args = custom_config[SERVING_ARGS_KEY]
    if not ai_platform_serving_args:
        raise ValueError("'ai_platform_serving_args' is missing in 'custom_config'")
    # Deploy the model.
    model_path = path_utils.serving_model_path(model_export_uri)
    # Note: we do not have a logical model version right now. This
    # model_version is a timestamp mapped to trainer's exporter.
    model_version = os.path.basename(model_path)
    executor_class_path = "%s.%s" % (self.__class__.__module__, self.__class__.__name__)
    runner.deploy_model_for_aip_prediction(
        model_path,
        model_version,
        ai_platform_serving_args,
        executor_class_path,
    )

    model_push.set_int_custom_property("pushed", 1)
    model_push.set_string_custom_property("pushed_model", model_path)


# Traceback (KeyError)
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-38-f67dc69712f4> in <module>
----> 1 context.run(pusher)

~/XXX/tfx/tfx/tfx/orchestration/experimental/interactive/interactive_context.py in run_if_ipython(*args, **kwargs)
64       # __IPYTHON__ variable is set by IPython, see
65       # https://ipython.org/ipython-doc/rel-0.10.2/html/interactive/reference.html#embedding-ipython.
---> 66       return fn(*args, **kwargs)
67     else:
68       absl.logging.warning(

~/XXX/tfx/tfx/tfx/orchestration/experimental/interactive/interactive_context.py in run(self, component, enable_cache, beam_pipeline_args)
166         component, pipeline_info, driver_args, metadata_connection,
167         beam_pipeline_args, additional_pipeline_args)
--> 168     execution_id = launcher.launch().execution_id
169
170     return execution_result.ExecutionResult(

~/XXX/tfx/tfx/tfx/orchestration/launcher/base_component_launcher.py in launch(self)
203                          execution_decision.input_dict,
204                          execution_decision.output_dict,
--> 205                          execution_decision.exec_properties)
206
207     absl.logging.info('Running publisher for %s',

~/XXX/tfx/tfx/tfx/orchestration/launcher/in_process_component_launcher.py in _run_executor(self, execution_id, input_dict, output_dict, exec_properties)
65         executor_context)  # type: ignore
66
---> 67     executor.Do(input_dict, output_dict, exec_properties)

~/XXX/tfx/tfx/tfx/extensions/google_cloud_big_query_ml/pusher/executor.py in Do(self, input_dict, output_dict, exec_properties)
59     """
60     self._log_startup(input_dict, output_dict, exec_properties)
---> 61     model_push = artifact_utils.get_single_instance(output_dict['model_push'])
62     if not self.CheckBlessing(input_dict):
63       model_push.set_int_custom_property('pushed', 0)

KeyError: 'model_push'

# Fixed code
def Do(
    self,
    input_dict: Dict[Text, List[types.Artifact]],
    output_dict: Dict[Text, List[types.Artifact]],
    exec_properties: Dict[Text, Any],
):
    """Overrides the tfx_pusher_executor.

    Args:
      input_dict: Input dict from input key to a list of artifacts, including:
        - model_export: exported model from trainer.
        - model_blessing: model blessing path from model_validator.
      output_dict: Output dict from key to a list of artifacts, including:
        - model_push: A list of 'ModelPushPath' artifact of size one. It will
          include the model in this push execution if the model was pushed.
      exec_properties: Mostly a passthrough input dict for
        tfx.components.Pusher.executor.  custom_config.ai_platform_serving_args
        is consumed by this class.  For the full set of parameters supported by
        Google Cloud AI Platform, refer to
        https://cloud.google.com/ml-engine/docs/tensorflow/deploying-models#creating_a_model_version.

    Raises:
      ValueError:
        If ai_platform_serving_args is not in exec_properties.custom_config.
        If Serving model path does not start with gs://.
      RuntimeError: if the Google Cloud AI Platform training job failed.
    """
    self._log_startup(input_dict, output_dict, exec_properties)
    model_push = artifact_utils.get_single_instance(
        output_dict[tfx_pusher_executor.PUSHED_MODEL_KEY]
    )
    if not self.CheckBlessing(input_dict):
        self._MarkNotPushed(model_push)
        return

    model_export = artifact_utils.get_single_instance(
        input_dict[tfx_pusher_executor.MODEL_KEY]
    )

    exec_properties_copy = exec_properties.copy()
    custom_config = exec_properties_copy.pop(_CUSTOM_CONFIG_KEY, {})
    ai_platform_serving_args = custom_config.get(SERVING_ARGS_KEY)
    if not ai_platform_serving_args:
        raise ValueError("'ai_platform_serving_args' is missing in 'custom_config'")
    # Deploy the model.
    io_utils.copy_dir(
        src=path_utils.serving_model_path(model_export.uri), dst=model_push.uri
    )
    model_path = model_push.uri
    # TODO(jjong): Introduce Versioning.
    # Note that we're adding "v" prefix as Cloud AI Prediction only allows the
    # version name that starts with letters, and contains letters, digits,
    # underscore only.
    model_version = "v{}".format(int(time.time()))
    executor_class_path = "%s.%s" % (self.__class__.__module__, self.__class__.__name__)
    runner.deploy_model_for_aip_prediction(
        model_path,
        model_version,
        ai_platform_serving_args,
        executor_class_path,
    )

    self._MarkPushed(
        model_push,
        pushed_destination=_CAIP_MODEL_VERSION_PATH_FORMAT.format(
            project_id=ai_platform_serving_args["project_id"],
            model=ai_platform_serving_args["model_name"],
            version=model_version,
        ),
        pushed_version=model_version,
    )