# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.


# Define variable `mlflow_tracking_uri`
mlflow_tracking_uri: &mlflow_tracking_uri "http://127.0.0.1:5000" 

# Define the workflow
workflow:
  name: "MLflow"

# Define the tasks within the workflow
tasks:
  - name: train_xgboost_native
    task_type: MLFlowProjectsCustom 
    repository: https://github.com/mlflow/mlflow#examples/xgboost/xgboost_native
    mlflow_tracking_uri: *mlflow_tracking_uri
    parameters: -P learning_rate=0.2 -P colsample_bytree=0.8 -P subsample=0.9
    experiment_name: xgboost

  - name: train_automl
    task_type: MLFlowProjectsAutoML 
    mlflow_tracking_uri: *mlflow_tracking_uri
    parameters: time_budget=30;estimator_list=['lgbm']
    experiment_name: automl_iris
    model_name: iris_A
    automl_tool: flaml
    data_path: /data/examples/iris

  - name: deploy_docker
    task_type: MLflowModels 
    deps: [train_automl]
    model_uri: models:/iris_A/Production
    mlflow_tracking_uri: *mlflow_tracking_uri
    deploy_mode: DOCKER
    port: 7002

  - name: train_basic_algorithm
    task_type: MLFlowProjectsBasicAlgorithm 
    mlflow_tracking_uri: *mlflow_tracking_uri
    parameters: n_estimators=200;learning_rate=0.2
    experiment_name: basic_algorithm_iris
    model_name: iris_B
    algorithm: lightgbm
    data_path: /data/examples/iris
    search_params: max_depth=[5, 10];n_estimators=[100, 200]

  - name: deploy_mlflow
    deps: [train_basic_algorithm]
    task_type: MLflowModels
    model_uri: models:/iris_B/Production
    mlflow_tracking_uri: *mlflow_tracking_uri
    deploy_mode: MLFLOW
    port: 7001