name: evaluation_data_sampler
description: |
  Randomly downsamples an input dataset to a specified size for computing Vertex XAI feature
  attributions for AutoML Tables and custom models. Creates a Dataflow job with Apache Beam to
  downsample the dataset.

  Args:
      project (str):
          Required. Project to retrieve dataset from.
      location (Optional[str]):
          Location to retrieve dataset from.
          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.
      gcs_source_uris (Sequence[str]):
          Required. Google Cloud Storage URI(-s) to your instances to run
          data sampler on. They must match `instances_format`.
          May contain wildcards. For more information on wildcards, see
          https://cloud.google.com/storage/docs/gsutil/addlhelp/WildcardNames.
      instances_format (Optional[str]):
          The format in which instances are given, must be one of
          the model's supported input storage formats.
          If not set, default to "jsonl".
      sample_size (Optional[int]):
          Sample size of the randomly sampled dataset. 10k by default.
      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_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:
      gcs_output_directory (JsonArray):
          JsonArray of the downsampled dataset GCS output.
      gcp_resources (str):
          Serialized gcp_resources proto tracking the data sampler.
inputs:
  - { name: project, type: String }
  - { name: location, type: String, default: "us-central1" }
  - { name: root_dir, type: String }
  - { name: gcs_source_uris, type: JsonArray, default: "[]" }
  - { name: instances_format, type: String, default: "jsonl"}
  - { name: sample_size, type: Integer, default: 10000 }
  - { name: dataflow_service_account, type: String, default: "" }
  - { 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: gcs_output_directory, type: JsonArray }
  - { name: gcp_resources, type: String }
implementation:
  container:
    image: gcr.io/ml-pipeline/model-evaluation:v0.3
    command:
      - python
      - /main.py
    args:
      - --task
      - "data_sampler"
      - --display_name
      - "data-sampler-run"
      - --project_id
      - { inputValue: project }
      - --location
      - { inputValue: location }
      - --root_dir
      - "{{$.inputs.parameters['root_dir']}}/{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}"
      - --gcs_source_uris
      - { inputValue: gcs_source_uris }
      - --instances_format
      - { inputValue: instances_format}
      - --sample_size
      - { inputValue: sample_size }
      - --dataflow_service_account
      - { inputValue: dataflow_service_account }
      - --dataflow_subnetwork
      - { inputValue: dataflow_subnetwork }
      - --dataflow_use_public_ips
      - { inputValue: dataflow_use_public_ips }
      - --kms_key_name
      - { inputValue: encryption_spec_key_name }
      - --gcs_output_directory
      - { outputPath: gcs_output_directory }
      - --gcp_resources
      - { outputPath: gcp_resources }
      - --executor_input
      - "{{$}}"
