# Copyright 2021 The Kubeflow Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

name: dataproc_create_spark_r_batch
description: |
  Create a Dataproc SparkR batch workload and wait for it to finish.

    Args:
        project (str):
            Required: Project to run the Dataproc batch workload.
        location (Optional[str]):
            Location of the Dataproc batch workload. If not set, default to `us-central1`.
        batch_id (Optional[str]):
            The ID to use for the batch, which will become the final component of the batch's
            resource name. If none is specified, a default name will be generated by the component.

            This value must be 4-63 characters. Valid characters are /[a-z][0-9]-/.
        labels (Optional[dict]):
            The labels to associate with this batch. Label keys must contain 1 to 63 characters,
            and must conform to RFC 1035. Label values may be empty, but, if present, must contain
            1 to 63 characters, and must conform to RFC 1035. No more than 32 labels can be
            associated with a batch.

            An object containing a list of "key": value pairs.
            Example: { "name": "wrench", "mass": "1.3kg", "count": "3" }.
        container_image (Optional[str]):
            Optional custom container image for the job runtime environment. If not specified, a
            default container image will be used.
        runtime_config_version (Optional[str]):
            Version of the batch runtime.
        runtime_config_properties (Optional[dict]):
            Runtime configuration for a workload.
        service_account (Optional[str]):
            Service account that used to execute workload.
        network_tags (Optional[Sequence]):
            Tags used for network traffic control.
        kms_key (Optional[str]):
            The Cloud KMS key to use for encryption.
        network_uri (Optional[str]):
            Network URI to connect workload to.
        subnetwork_uri (Optional[str]):
            Subnetwork URI to connect workload to.
        metastore_service (Optional[str]):
            Resource name of an existing Dataproc Metastore service.
        spark_history_dataproc_cluster (Optional[str]):
            The Spark History Server configuration for the workload.
        main_r_file_uri (str):
            Required. The HCFS URI of the main R file to use as the driver. Must be a .R or .r file.
        file_uris (Optional[Sequence]):
            HCFS URIs of files to be placed in the working directory of each executor.
        archive_uris (Optional[Sequence]):
            HCFS URIs of archives to be extracted into the working directory of each executor.
        args (Optional[Sequence]):
            The arguments to pass to the driver.

    Returns:
        gcp_resources (str):
            Serialized gcp_resources proto tracking the Dataproc batch workload.
            For more details, see https://github.com/kubeflow/pipelines/blob/master/components/google-cloud/google_cloud_pipeline_components/proto/README.md.
inputs:
- {name: project, type: String}
- {name: location, type: String, optional: true, default: 'us-central1'}
- {name: batch_id, type: String, optional: true, default: ''}
- {name: labels, type: JsonObject, optional: true, default: '{}'}
- {name: container_image, type: String, optional: true, default: ''}
- {name: runtime_config_version, type: String, optional: true, default: ''}
- {name: runtime_config_properties, type: JsonObject, optional: true, default: '{}'}
- {name: service_account, type: String, optional: true, default: ''}
- {name: network_tags, type: JsonArray, optional: true, default: '[]'}
- {name: kms_key, type: String, optional: true, default: ''}
- {name: network_uri, type: String, optional: true, default: ''}
- {name: subnetwork_uri, type: String, optional: true, default: ''}
- {name: metastore_service, type: String, optional: true, default: ''}
- {name: spark_history_dataproc_cluster, type: String, optional: true, default: ''}
- {name: main_r_file_uri, type: String}
- {name: file_uris, type: JsonArray, optional: true, default: '[]'}
- {name: archive_uris, type: JsonArray, optional: true, default: '[]'}
- {name: args, type: JsonArray, optional: true, default: '[]'}
outputs:
- {name: gcp_resources, type: String}
implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:2.0.0b1
    command: [python3, -u, -m, google_cloud_pipeline_components.container.v1.dataproc.create_spark_r_batch.launcher]
    args: [
      --type, DataprocSparkRBatch,
      --payload,
      concat: [
          '{',
            '"labels": ', {inputValue: labels},
            ', "runtime_config": {',
              '"version": "', {inputValue: runtime_config_version}, '"',
              ', "container_image": "', {inputValue: container_image}, '"',
              ', "properties": ', {inputValue: runtime_config_properties},
            '}',
            ', "environment_config": {',
              '"execution_config": {',
                '"service_account": "', {inputValue: service_account}, '"',
                ', "network_tags": ', {inputValue: network_tags},
                ', "kms_key": "', {inputValue: kms_key}, '"',
                ', "network_uri": "', {inputValue: network_uri}, '"',
                ', "subnetwork_uri": "', {inputValue: subnetwork_uri}, '"',
              '}',
              ', "peripherals_config": {',
                '"metastore_service": "', {inputValue: metastore_service}, '"',
                ', "spark_history_server_config": { ',
                  '"dataproc_cluster": "', {inputValue: spark_history_dataproc_cluster}, '"',
                '}',
              '}',
            '}',
            ', "spark_r_batch": {',
              '"main_r_file_uri": "', {inputValue: main_r_file_uri}, '"',
              ', "file_uris": ', {inputValue: file_uris},
              ', "archive_uris": ', {inputValue: archive_uris},
              ', "args": ', {inputValue: args},
            '}',
          '}'
      ],
      --project, {inputValue: project},
      --location, {inputValue: location},
      --batch_id, {inputValue: batch_id},
      --gcp_resources, {outputPath: gcp_resources}
    ]
