# Buggy code
def Do(
    self,
    input_dict: Dict[Text, List[types.Artifact]],
    output_dict: Dict[Text, List[types.Artifact]],
    exec_properties: Dict[Text, Any],
) -> None:
    """Push model to target directory if blessed.

    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.  A push
          action delivers the model exports produced by Trainer to the
          destination defined in component config.
      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: A dict of execution properties, including:
        - push_destination: JSON string of pusher_pb2.PushDestination instance,
          providing instruction of destination to push model.

    Returns:
      None
    """
    self._log_startup(input_dict, output_dict, exec_properties)
    model_push = artifact_utils.get_single_instance(output_dict[PUSHED_MODEL_KEY])
    if not self.CheckBlessing(input_dict):
        model_push.set_int_custom_property("pushed", 0)
        return
    model_push_uri = model_push.uri
    model_export = artifact_utils.get_single_instance(input_dict[MODEL_KEY])
    model_export_uri = model_export.uri
    logging.info("Model pushing.")
    # Copy the model to pushing uri.
    model_path = path_utils.serving_model_path(model_export_uri)
    model_version = path_utils.get_serving_model_version(model_export_uri)
    logging.info("Model version is %s", model_version)
    io_utils.copy_dir(model_path, os.path.join(model_push_uri, model_version))
    logging.info("Model written to %s.", model_push_uri)

    # Copied to a fixed outside path, which can be listened by model server.
    #
    # If model is already successfully copied to outside before, stop copying.
    # This is because model validator might blessed same model twice (check
    # mv driver) with different blessing output, we still want Pusher to
    # handle the mv output again to keep metadata tracking, but no need to
    # copy to outside path again..
    # TODO(jyzhao): support rpc push and verification.
    push_destination = pusher_pb2.PushDestination()
    json_format.Parse(exec_properties["push_destination"], push_destination)
    serving_path = os.path.join(
        push_destination.filesystem.base_directory, model_version
    )
    if tf.io.gfile.exists(serving_path):
        logging.info(
            "Destination directory %s already exists, skipping current push.",
            serving_path,
        )
    else:
        # tf.serving won't load partial model, it will retry until fully copied.
        io_utils.copy_dir(model_path, serving_path)
        logging.info("Model written to serving path %s.", serving_path)

    model_push.set_int_custom_property("pushed", 1)
    model_push.set_string_custom_property("pushed_model", model_export_uri)
    model_push.set_int_custom_property("pushed_model_id", model_export.id)
    logging.info("Model pushed to %s.", serving_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],
) -> None:
    """Push model to target directory if blessed.

    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.  A push
          action delivers the model exports produced by Trainer to the
          destination defined in component config.
      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: A dict of execution properties, including:
        - push_destination: JSON string of pusher_pb2.PushDestination instance,
          providing instruction of destination to push model.

    Returns:
      None
    """
    self._log_startup(input_dict, output_dict, exec_properties)
    model_push = artifact_utils.get_single_instance(output_dict[PUSHED_MODEL_KEY])
    if not self.CheckBlessing(input_dict):
        self._MarkNotPushed(model_push)
        return
    model_export = artifact_utils.get_single_instance(input_dict[MODEL_KEY])
    model_path = path_utils.serving_model_path(model_export.uri)

    # Push model to the destination, which can be listened by a model server.
    #
    # If model is already successfully copied to outside before, stop copying.
    # This is because model validator might blessed same model twice (check
    # mv driver) with different blessing output, we still want Pusher to
    # handle the mv output again to keep metadata tracking, but no need to
    # copy to outside path again..
    # TODO(jyzhao): support rpc push and verification.
    push_destination = pusher_pb2.PushDestination()
    json_format.Parse(exec_properties["push_destination"], push_destination)

    destination_kind = push_destination.WhichOneof("destination")
    if destination_kind == "filesystem":
        fs_config = push_destination.filesystem
        if fs_config.versioning == _Versioning.AUTO:
            fs_config.versioning = _Versioning.UNIX_TIMESTAMP
        if fs_config.versioning == _Versioning.UNIX_TIMESTAMP:
            model_version = str(int(time.time()))
        else:
            raise NotImplementedError(
                "Invalid Versioning {}".format(fs_config.versioning)
            )
        logging.info("Model version: %s", model_version)
        serving_path = os.path.join(fs_config.base_directory, model_version)

        if tf.io.gfile.exists(serving_path):
            logging.info(
                "Destination directory %s already exists, skipping current push.",
                serving_path,
            )
        else:
            # tf.serving won't load partial model, it will retry until fully copied.
            io_utils.copy_dir(model_path, serving_path)
            logging.info("Model written to serving path %s.", serving_path)
    else:
        raise NotImplementedError(
            "Invalid push destination {}".format(destination_kind)
        )

    # Copy the model to pushing uri for archiving.
    io_utils.copy_dir(model_path, model_push.uri)
    self._MarkPushed(
        model_push, pushed_destination=serving_path, pushed_version=model_version
    )
    logging.info("Model pushed to %s.", model_push.uri)