name: get_vertex_model
description: TO BE REMOVED. Gets a Vertex Model Artifact.
inputs:
- {name: model_resource_name, type: String, description: Vertex Model Resource Name.}
outputs:
- {name: model, type: google.VertexModel, description: The Vertex Model Artifact.}
implementation:
  container:
    image: python:3.8
    command:
    - sh
    - -ec
    - |
      program_path=$(mktemp)
      printf "%s" "$0" > "$program_path"
      python3 -u "$program_path" "$@"
    - |2

      import os

      def get_vertex_model(executor_input_json, model_resource_name: str):
          executor_output = {}
          executor_output['artifacts'] = {}
          for name, artifacts in executor_input_json.get('outputs',
                                                        {}).get('artifacts',
                                                                {}).items():
            artifacts_list = artifacts.get('artifacts')
            if name == 'model' and artifacts_list:
              updated_runtime_artifact = artifacts_list[0]
              updated_runtime_artifact['uri'] = model_resource_name
              updated_runtime_artifact['metadata'] = {'resourceName': model_resource_name}
              artifacts_list = {'artifacts': [updated_runtime_artifact]}

            executor_output['artifacts'][name] = artifacts_list

          # Update the output artifacts.
          os.makedirs(
              os.path.dirname(executor_input_json['outputs']['outputFile']),
              exist_ok=True)
          with open(executor_input_json['outputs']['outputFile'], 'w') as f:
            f.write(json.dumps(executor_output))

      if __name__ == '__main__':
        import argparse
        import json

        parser = argparse.ArgumentParser(description='get resource name')
        parser.add_argument('--model_resource_name', type=str)
        parser.add_argument('--executor_input', type=str)
        args, _ = parser.parse_known_args()
        executor_input_json = json.loads(args.executor_input)
        model_resource_name = args.model_resource_name
        get_vertex_model(executor_input_json, model_resource_name)
    args:
    - --executor_input
    - "{{$}}"
    - --model_resource_name
    - {inputValue: model_resource_name}
