# Copyright 2023 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: prophet_trainer
description: |
  Launch a Prophet custom training job using Vertex CustomJob API.

    Args:
        project (str):
            The GCP project that runs the pipeline components.
        location (str):
            The GCP region for Vertex AI.
        root_dir (str):
            The Cloud Storage location to store the output.
        time_column (str):
            Name of the column that identifies time order in the time series.
        time_series_identifier_column (str):
            Name of the column that identifies the time series.
        target_column (str):
            Name of the column that the model is to predict values for.
        forecast_horizon (int):
            The number of time periods into the future for which forecasts will
            be created. Future periods start after the latest timestamp for each
            time series.
        optimization_objective (str):
            Optimization objective for tuning. Supported metrics come from
            Prophet's performance_metrics function. These are mse, rmse, mae,
            mape, mdape, smape, and coverage.
        data_granularity_unit (str):
            String representing the units of time for the time column.
        predefined_split_column (str):
            The predefined_split column name.
           A string that represents a list of comma separated CSV filenames.
        source_bigquery_uri (str):
            The BigQuery table path of format
            bq (str)://bq_project.bq_dataset.bq_table
        window_column (str):
            Name of the column that should be used to filter input rows.  The
            column should contain either booleans or string booleans; if the
            value of the row is True, generate a sliding window from that row.
        max_num_trials (Optional[int]):
            Maximum number of tuning trials to perform per time series. There
            are up to 100 possible combinations to explore for each time series.
            Recommended values to try are 3, 6, and 24.
        encryption_spec_key_name (Optional[str]):
            Customer-managed encryption key.
        dataflow_machine_type (Optional[str]):
            The dataflow machine type used for training.
        dataflow_max_num_workers (Optional[int]):
            The max number of Dataflow workers used for training.
        dataflow_disk_size_gb (Optional[int]):
            Dataflow worker's disk size in GB during training.
        dataflow_service_account (Optional[str]):
            Custom service account to run dataflow jobs.
        dataflow_subnetwork (Optional[str]):
            Dataflow's fully qualified subnetwork name, when empty the default
            subnetwork will be used.
        dataflow_use_public_ips (Optional[bool]):
            Specifies whether Dataflow workers use public IP addresses.

    Returns:
        gcp_resources (str):
            Serialized gcp_resources proto tracking the custom training job.
        unmanaged_container_model (google.UnmanagedContainerModel):
            The UnmanagedContainerModel artifact.
inputs:
  - { name: project, type: String }
  - { name: location, type: String }
  - { name: root_dir, type: String }
  - { name: target_column, type: String}
  - { name: time_column, type: String}
  - { name: time_series_identifier_column, type: String}
  - { name: forecast_horizon, type: Integer}
  - { name: window_column, type: String}
  - { name: optimization_objective, type: String, default: "rmse"}
  - { name: data_granularity_unit, type: String}
  - { name: predefined_split_column, type: String}
  - { name: max_num_trials, type: Integer, default: 6}
  - { name: encryption_spec_key_name, type: String, default: ""}
  - { name: source_bigquery_uri, type: String}
  - { name: dataflow_max_num_workers, type: Integer, default: 10}
  - { name: dataflow_machine_type, type: String, default: "n1-standard-1"}
  - { name: dataflow_disk_size_gb, type: Integer, default: 40}
  - { name: dataflow_service_account, type: String, default: ""}
  - { name: dataflow_subnetwork, type: String, default: ""}
  - { name: dataflow_use_public_ips, type: Boolean, default: "true"}

outputs:
  - { name: gcp_resources, type: String }
  - { name: unmanaged_container_model, type: google.UnmanagedContainerModel }
  - { name: evaluated_examples_directory, type: system.Artifact }

implementation:
  container:
    image: gcr.io/ml-pipeline/google-cloud-pipeline-components:1.0.32
    command: [python3, -u, -m, google_cloud_pipeline_components.container.v1.custom_job.launcher]
    args: [
      --type, CustomJob,
      --project, { inputValue: project },
      --location, { inputValue: location },
      --gcp_resources, { outputPath: gcp_resources },
      --payload,
      concat: [
          '{"display_name": "prophet-trainer-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}", ',
          '"encryption_spec": {"kms_key_name":"', { inputValue: encryption_spec_key_name }, '"}, ',
          '"job_spec": {"worker_pool_specs": [{"replica_count":"1", ',
          '"machine_spec": {"machine_type": "n1-standard-4"}, ',
          # TODO(b/260611238): Replace with prophet-training tag.
          '"container_spec": {"image_uri":"us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/training:20230322_0325", ',
          '"args": ["prophet_trainer", "',
          '--job_name={{$.pipeline_job_name}}", "',
          '--dataflow_worker_container_image=us-docker.pkg.dev/vertex-ai/automl-tabular/dataflow-worker:20230322_0325", "',
          '--prediction_container_image=us-docker.pkg.dev/vertex-ai/automl-tabular/fte-prediction-server:20230322_0325", "',
          '--artifacts_dir=', {inputValue: root_dir}, '/{{$.pipeline_job_uuid}}/model/", "',
          '--evaluated_examples_dir=', {inputValue: root_dir}, '/{{$.pipeline_job_uuid}}/eval/", "',
          '--region=', { inputValue: location}, '", "',
          '--source_bigquery_uri=', { inputValue: source_bigquery_uri}, '", "',
          '--target_column=', { inputValue: target_column}, '", "',
          '--time_column=', { inputValue: time_column}, '", "',
          '--time_series_identifier_column=', { inputValue: time_series_identifier_column}, '", "',
          '--forecast_horizon=', { inputValue: forecast_horizon}, '", "',
          '--window_column=', { inputValue: window_column}, '", "',
          '--optimization_objective=', { inputValue: optimization_objective}, '", "',
          '--data_granularity_unit=', { inputValue: data_granularity_unit}, '", "',
          '--predefined_split_column=', { inputValue: predefined_split_column}, '", "',
          '--max_num_trials=', { inputValue: max_num_trials}, '", "',
          '--dataflow_project=', { inputValue: project}, '", "',
          '--dataflow_max_num_workers=', { inputValue: dataflow_max_num_workers}, '", "',
          '--dataflow_machine_type=', { inputValue: dataflow_machine_type}, '", "',
          '--dataflow_disk_size_gb=', { inputValue: dataflow_disk_size_gb}, '", "',
          '--dataflow_service_account=', { inputValue: dataflow_service_account}, '", "',
          '--dataflow_subnetwork=', { inputValue: dataflow_subnetwork}, '", "',
          '--dataflow_use_public_ips=', { inputValue: dataflow_use_public_ips}, '", "',
          '--gcp_resources_path=', {outputPath: gcp_resources}, '", "',
          '--executor_input={{$.json_escape[1]}}"]}}]}}'
      ]
    ]
