# 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: automl_forecasting_stage_2_tuner
description: |
  AutoML Forecasting stage 2 tuner

    Args:
        project (str):
            Required. Project to run stage 2 tuner.
        location (str):
            Cloud region for running the component (ex: us-central1).
        root_dir (str):
            The Cloud Storage location to store the output.
        worker_pool_specs_override_json (JsonArray):
            JSON worker pool specs. E.g.,
            [{"machine_spec": {"machine_type": "n1-standard-16"}},{},{},{"machine_spec": {"machine_type": "n1-standard-16"}}]
        num_selected_trials (int):
            Number of selected trials. The number of weak learners in the final model.
        deadline_hours (float):
            Number of hours the cross-validation trainer should run.
        num_parallel_trials (int):
            Number of parallel training trials.
        single_run_max_secs (int):
            Max number of seconds each training trial runs.
        metadata (TabularExampleGenMetadata):
            The forecasting example gen metadata.
        transform_output (TransformOutput):
            The transform output artifact.
        materialized_train_split (MaterializedSplit):
            The materialized train split.
        materialized_eval_split (MaterializedSplit):
            The materialized eval split.
        encryption_spec_key_name (Optional[str]):
            Customer-managed encryption key.
        tuning_result_input_path (String):
            Path to the json of hyperparameter tuning results to use when evaluating models.

    Returns:
        gcp_resources (str):
            GCP resources created by this component.
            For more details, see https://github.com/kubeflow/pipelines/blob/master/components/google-cloud/google_cloud_pipeline_components/proto/README.md.
        tuning_result_output (AutoMLTabularTuningResult):
            The trained (private) model artifact paths and their hyperparameters.
inputs:
- {name: project, type: String}
- {name: location, type: String}
- {name: root_dir, type: String}
- {name: worker_pool_specs_override_json, type: JsonArray, default: "[]"}
- {name: num_selected_trials, type: Integer}
- {name: deadline_hours, type: Float}
- {name: num_parallel_trials, type: Integer}
- {name: single_run_max_secs, type: Integer}
- {name: metadata, type: TabularExampleGenMetadata}
- {name: transform_output, type: TransformOutput}
- {name: materialized_train_split, type: MaterializedSplit}
- {name: materialized_eval_split, type: MaterializedSplit}
- {name: encryption_spec_key_name, type: String, default: ""}
- {name: tuning_result_input_path, type: AutoMLTabularTuningResult}


outputs:
- {name: gcp_resources, type: String}
- {name: tuning_result_output, type: AutoMLTabularTuningResult}

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": "automl-forecasting-stage-2-tuner-{{$.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-8"}, "container_spec": {"image_uri":"',
        'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20230322_0325','",
        "args": ["forecasting_mp_l2l_stage_2_tuner','",
        "--region=',{inputValue: location},'",
        "--transform_output_path=',{inputUri: transform_output},'",
        "--training_docker_uri=','us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20230322_0325','",
        "--component_id={{$.pipeline_task_uuid}}','",
        "--training_base_dir=',{inputValue: root_dir}, '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/train','",
        "--num_parallel_trial=',{inputValue: num_parallel_trials},'",
        "--single_run_max_secs=',{inputValue: single_run_max_secs},'",
        "--deadline_hours=',{inputValue: deadline_hours},'",
        "--num_selected_trials=', {inputValue: num_selected_trials}, '",
        "--lro_job_info=', {inputValue: root_dir}, '/{{$.pipeline_job_uuid}}/lro', '",
        "--error_file_path=', {inputValue: root_dir}, '/{{$.pipeline_job_uuid}}/{{$.pipeline_task_uuid}}/error.pb','",
        "--metadata_path=', {inputUri: metadata},'",
        "--materialized_train_split=', {inputUri: materialized_train_split},'",
        "--materialized_eval_split=', {inputUri: materialized_eval_split}, '",
        "--tuning_result_input_path=',{inputUri: tuning_result_input_path}, '",
        "--kms_key_name=', {inputValue: encryption_spec_key_name},'",
        "--gcp_resources_path=', {outputPath: gcp_resources}, '",
        "--tuning_result_output_path=', {outputUri: tuning_result_output}, '",
        "--use_json=true",
        "--log_level=ERROR",
        "--executor_input={{$.json_escape[1]}}"]}}]}}',
      ]]
