name: feature_extractor_error_analysis
description: |
  Extracts feature embeddings using the Starburst V5 model, and stores them into a private GCS bucket

  Args:
      project (str):
          Required. GCP Project ID.
      location (Optional[str]):
          GCP Region.
          If not set, defaulted to `us-central1`.
      root_dir (str):
          Required. The GCS directory for keeping staging files.
          A random subdirectory will be created under the directory to keep job info for resuming
          the job in case of failure.
      test_dataset (Optional[google.VertexDataset])
          Required. A google.VertexDataset artifact of the test dataset.
      training_dataset (Optional[google.VertexDataset])
          Required. A google.VertexDataset artifact of the training dataset.
      preprocessed_test_dataset_storage_source (str):
          Required. Google Cloud Storage URI to preprocessed test dataset for Vision Analysis pipelines.
      preprocessed_training_dataset_storage_source (str):
          Required. Google Cloud Storage URI to preprocessed training dataset for Vision Error Analysis pipelines.
      output_dir (str):
          Required. The Google Cloud Storage location where the outputs will be saved to.
      dataflow_service_account (Optional[str]):
          Optional. Service account to run the dataflow job.
          If not set, dataflow will use the default woker service account.

          For more details, see https://cloud.google.com/dataflow/docs/concepts/security-and-permissions#default_worker_service_account
      dataflow_disk_size (Optional[int]):
          Optional. The disk size (in GB) of the machine executing the evaluation run.
          If not set, defaulted to `50`.
      dataflow_machine_type (Optional[str]):
          Optional. The machine type executing the evaluation run.
          If not set, defaulted to `n1-standard-4`.
      dataflow_workers_num (Optional[int]):
          Optional. The number of workers executing the evaluation run.
          If not set, defaulted to `10`.
      dataflow_max_workers_num (Optional[int]):
          Optional. The max number of workers executing the evaluation run.
          If not set, defaulted to `25`.
      dataflow_subnetwork (Optional[str]):
          Dataflow's fully qualified subnetwork name, when empty the default subnetwork will be
          used. More details:
          https://cloud.google.com/dataflow/docs/guides/specifying-networks#example_network_and_subnetwork_specifications
      dataflow_use_public_ips (Optional[bool]):
          Specifies whether Dataflow workers use public IP addresses.
      encryption_spec_key_name (Optional[str]):
          Customer-managed encryption key for the Dataflow job. If this is set, then all resources
          created by the Dataflow job will be encrypted with the provided encryption key.
  Returns:
      embedding_output_file (str):
          Google Cloud Storage URI to computed feature embeddings of the image datasets.
inputs:
  - { name: project, type: String }
  - { name: location, type: String, default: "us-central1" }
  - { name: root_dir, type: String }
  - { name: test_dataset, type: google.VertexDataset, optional: True }
  - { name: training_dataset, type: google.VertexDataset, optional: True }
  - { name: preprocessed_test_dataset_storage_source, type: String }
  - { name: preprocessed_training_dataset_storage_source, type: String }
  - { name: output_dir, type: String }
  - { name: dataflow_service_account, type: String, default: "" }
  - { name: dataflow_disk_size, type: Integer, default: 50 }
  - { name: dataflow_machine_type, type: String, default: "n1-standard-8" }
  - { name: dataflow_workers_num, type: Integer, default: 1 }
  - { name: dataflow_max_workers_num, type: Integer, default: 5 }
  - { name: dataflow_subnetwork, type: String, default: "" }
  - { name: dataflow_use_public_ips, type: Boolean, default: "true" }
  - { name: encryption_spec_key_name, type: String, default: "" }
outputs:
  - { name: embedding_output_file, type: String }
  - { name: gcp_resources, type: String }
implementation:
  container:
    image: gcr.io/cloud-aiplatform-private/starburst/v5/cmle:20230203_1221_RC00
    args:
      - --project_id
      - { inputValue: project }
      - --location
      - { inputValue: location }
      - --root_dir
      - { inputValue: root_dir }
      - --test_dataset_resource_name
      - "{{$.inputs.artifacts['test_dataset'].metadata['resourceName']}}"
      - --test_dataset_storage_source
      - { inputValue: preprocessed_test_dataset_storage_source }
      - --training_dataset_resource_name
      - "{{$.inputs.artifacts['training_dataset'].metadata['resourceName']}}"
      - --training_dataset_storage_source
      - { inputValue: preprocessed_training_dataset_storage_source }
      - --output_dir
      - { inputValue: output_dir }
      - --dataflow_job_prefix
      - "feature-extraction-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}"
      - --dataflow_service_account
      - { inputValue: dataflow_service_account }
      - --dataflow_disk_size
      - { inputValue: dataflow_disk_size }
      - --dataflow_machine_type
      - { inputValue: dataflow_machine_type }
      - --dataflow_workers_num
      - { inputValue: dataflow_workers_num }
      - --dataflow_max_workers_num
      - { inputValue: dataflow_max_workers_num }
      - --dataflow_subnetwork
      - { inputValue: dataflow_subnetwork }
      - --dataflow_use_public_ips
      - { inputValue: dataflow_use_public_ips }
      - --kms_key_name
      - { inputValue: encryption_spec_key_name }
      - --embedding_output_file
      - { outputPath: embedding_output_file }
      - --gcp_resources
      - { outputPath: gcp_resources }
      - --executor_input
      - "{{$}}"
