# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3.8 - AzureML
#     language: python
#     name: python38-azureml
# ---

# ## Azure ML Pipeline - Autoencoder Training & Registration
# This notebook demonstrates creation & execution of an Azure ML pipeline designed to load data from an attached ADLS Gen 2 datastore, split the data into test/train sets, train an autoencoder using TensorFlow & Keras, then run an A/B test to check whether the new model performs better than the current 'best' registered model. If the new model is found to outperform the current best performer, it is registered to the AML workspace. This notebook was built as part of a larger MLOps solution where this pipeline is deployed and triggered via Azure DevOps and the best performing model is published to a real-time endpoint

# Import required packages
from azureml.core import Workspace, Experiment, Datastore, Environment, Dataset
from azureml.core.compute import ComputeTarget, AmlCompute, DataFactoryCompute
from azureml.core.compute_target import ComputeTargetException
from azureml.core.runconfig import RunConfiguration
from azureml.core.conda_dependencies import CondaDependencies
from azureml.core.runconfig import DEFAULT_CPU_IMAGE
from azureml.pipeline.core import Pipeline, PipelineParameter, PipelineData
from azureml.pipeline.steps import PythonScriptStep
from azureml.pipeline.core import PipelineParameter, PipelineData
from azureml.data.output_dataset_config import OutputTabularDatasetConfig, OutputDatasetConfig, OutputFileDatasetConfig
from azureml.data.datapath import DataPath
from azureml.data.data_reference import DataReference
from azureml.data.sql_data_reference import SqlDataReference
from azureml.pipeline.steps import DataTransferStep
import logging

# ### Connect to Azure ML Workspace, Provision Compute Resources, and get References to Datastores
# Connect to workspace using config associated config file. Get a reference to you pre-existing AML compute cluster or provision a new cluster to facilitate processing. Finally, get references to your default blob datastore.

# +
# Connect to AML Workspace
ws = Workspace.from_config()

#Select AML Compute Cluster
cpu_cluster_name = 'cpucluster'

# Verify that cluster does not exist already
try:
    cpu_cluster = ComputeTarget(workspace=ws, name=cpu_cluster_name)
    print('Found an existing cluster, using it instead.')
except ComputeTargetException:
    compute_config = AmlCompute.provisioning_configuration(vm_size='STANDARD_D3_V2',
                                                           min_nodes=0,
                                                           max_nodes=1)
    cpu_cluster = ComputeTarget.create(ws, cpu_cluster_name, compute_config)
    cpu_cluster.wait_for_completion(show_output=True)
    
#Get default datastore
default_ds = ws.get_default_datastore()
# -

# ### Create Run Configuration
# The RunConfiguration defines the environment used across all python steps. You can optionally add additional conda or pip packages to be added to your environment. [More details here.](https://docs.microsoft.com/en-us/python/api/azureml-core/azureml.core.conda_dependencies.condadependencies?view=azure-ml-py)
# ~~~
# run_config.environment.python.conda_dependencies = CondaDependencies.create(conda_packages=['requests'])
# run_config.environment.python.conda_dependencies.add_pip_package('azureml-opendatasets')
# ~~~
# Here, we also register the environment to the AML workspace so that it can be used for future retraining and inferencing operations.

# +
run_config = RunConfiguration()
run_config.docker.use_docker = True
run_config.environment = Environment(name='tf_keras_autoencoder_env')
run_config.environment.docker.base_image = DEFAULT_CPU_IMAGE
run_config.environment.python.conda_dependencies = CondaDependencies.create()
run_config.environment.python.conda_dependencies.set_pip_requirements([
    'requests==2.26.0',
    'pandas==0.25.3',
    'numpy==1.19.2',
    'scikit-learn==0.22.2.post1',
    'joblib==0.14.1',
    'h5py==3.1.0',
    'tensorflow==2.6.0',
    'keras==2.6.0',
    'azureml-defaults==1.33.0',
    'matplotlib'
])
run_config.environment.python.conda_dependencies.set_python_version('3.8.10')

#Register environment for reuse 
run_config.environment.register(ws)
# -

# ### Define Output Datasets
# Below we define the configuration for datasets that will be passed between steps in our pipeline. Note, in all cases we specify the datastore that should hold the datasets and whether they should be registered following step completion or not. This can optionally be disabled by removing the register_on_complete() call. upload_file_dataset is intended to hold the data within an uploaded CSV file and processed_dataset will contain our uploaded data post-processing.

autoencoder_raw_data = OutputFileDatasetConfig(name='Autoencoder_Raw_Data', destination=(default_ds, 'autoencoder_raw_data/{run-id}')).read_delimited_files().register_on_complete(name='Autoencoder_Raw_Data')
autoencoder_training_data = OutputFileDatasetConfig(name='Autoencoder_Training_Data', destination=(default_ds, 'autoencoder_training_data/{run-id}')).read_delimited_files().register_on_complete(name='Autoencoder_Training_Data')
autoencoder_testing_data = OutputFileDatasetConfig(name='Autoencoder_Testing_Data', destination=(default_ds, 'autoencoder_testing_data/{run-id}')).read_delimited_files().register_on_complete(name='Autoencoder_Testing_Data')

# ### Define Pipeline Data
# Pipeline data represents intermediate files/data in an Azure Machine Learning pipeline that needs to be shuttled between steps. In our case we are fitting a MinMax scaler in an early step that needs to be consumed in a later step. Additionall, we will register the trained scaler and model artifacts in the final step if our model performs better than the current model. For both cases, we use `PipelineData` objects to facilitate this motion of files between steps. [More details can be found here.](https://docs.microsoft.com/en-us/python/api/azureml-pipeline-core/azureml.pipeline.core.pipelinedata?view=azure-ml-py)

split_to_train_pipeline_data = PipelineData(name='Autoencoder_SplitScale_Outputs', datastore=default_ds)
train_to_evaluate_pipeline_data = PipelineData(name='Autoencoder_Training_Outputs', datastore=default_ds)

# ### Define Pipeline Parameters
# PipelineParameter objects serve as variable inputs to an Azure ML pipeline and can be specified at runtime. Below we specify the number of epochs (`num_epochs`) and batch size (`batch_size`) as variable parameters which can be passed into the pipeline at runtime.

num_epochs = PipelineParameter(name='num_epochs', default_value=10)
batch_size = PipelineParameter(name='batch_size', default_value=10)

# ### Define Pipeline Steps
# The pipeline below consists of four distinct steps all of which execute an associated python script located in the `./pipeline_script_steps` dir. First, we call `get_data.py` and retrieve data from the registered ADLS Gen 2 datastore and register this dataset as `Autoencoder_Raw_Data`. From here we run `split_and_scale.py` which normalizes the raw data using a `MinMaxScaler` and then splits into test and train datasets - both of which are subsequently registered. Then, we pass the test and training datasets into a step that runs `train_model.py` which trains the autoencoder and computes as Mean Squared Error (MSE) metric. Finally, the final step executes `evaluate_and_register.py` which compares the MSE of our model to the MSE of the currently registered autoencoder. If the current model performs better, or no model has been registered to-date, the model is registered in the workspace.

# +
#Get raw data from registered ADLS Gen2 datastore
#Register tabular dataset after retrieval
get_data_step = PythonScriptStep(
    name='Get Data from ADLS Gen2',
    script_name='get_data.py',
    arguments =['--autoencoder_raw_data', autoencoder_raw_data],
    outputs=[autoencoder_raw_data],
    compute_target=cpu_cluster,
    source_directory='./pipeline_step_scripts',
    allow_reuse=False,
    runconfig=run_config
)

#Normalize the raw data using a MinMaxScaler
#and then split into test and train datasets
split_scale_step = PythonScriptStep(
    name='Split and Scale Raw Data',
    script_name='split_and_scale.py',
    arguments =['--autoencoder_training_data', autoencoder_training_data,
                '--autoencoder_testing_data', autoencoder_testing_data,
                '--split_to_train_pipeline_data', split_to_train_pipeline_data],
    inputs=[autoencoder_raw_data.as_input(name='Autoencoder_Raw_Data')],
    outputs=[autoencoder_training_data, autoencoder_testing_data, split_to_train_pipeline_data],
    compute_target=cpu_cluster,
    source_directory='./pipeline_step_scripts',
    allow_reuse=False,
    runconfig=run_config
)

#Train autoencoder using raw data as an input
#Raw data will be preprocessed and registered as train/test datasets
#Scaler and train autoencoder will be saved out
train_model_step = PythonScriptStep(
    name='Train TF/Keras Autoencoder',
    script_name='train_model.py',
    arguments =[
                '--train_to_evaluate_pipeline_data', train_to_evaluate_pipeline_data,
                '--split_to_train_pipeline_data', split_to_train_pipeline_data,
                '--num_epochs', num_epochs,
                '--batch_size', batch_size],
    inputs=[autoencoder_training_data.as_input(name='Autoencoder_Training_Data'),
            autoencoder_testing_data.as_input(name='Autoencoder_Testing_Data'),
            split_to_train_pipeline_data.as_input('Autoencoder_SplitScale_Outputs')
           ],
    outputs=[train_to_evaluate_pipeline_data],
    compute_target=cpu_cluster,
    source_directory='./pipeline_step_scripts',
    allow_reuse=False,
    runconfig=run_config
)

#Evaluate and register model here
#Compare metrics from current model and register if better than current
#best model
evaluate_and_register_step = PythonScriptStep(
    name='Evaluate and Register Autoencoder',
    script_name='evaluate_and_register.py',
    arguments=['--autoencoder_training_outputs', train_to_evaluate_pipeline_data],
    inputs=[autoencoder_training_data.as_input(name='Autoencoder_Training_Data'),
            autoencoder_testing_data.as_input(name='Autoencoder_Testing_Data'),
            train_to_evaluate_pipeline_data.as_input('Autoencoder_Training_Outputs')],
    compute_target=cpu_cluster,
    source_directory='./pipeline_step_scripts',
    allow_reuse=False,
    runconfig=run_config
)
# -

# ### Create Pipeline
# Create an Azure ML Pipeline by specifying the steps to be executed. Note: based on the dataset dependencies between steps, exection occurs logically such that no step will execute unless all of the necessary input datasets have been generated.

pipeline = Pipeline(workspace=ws, steps=[get_data_step, split_scale_step, train_model_step, evaluate_and_register_step])

# ### Publish Pipeline
# Create a published version of your pipeline that can be triggered via an authenticated REST API request.

published_pipeline = pipeline.publish(name = 'Autoencoder_Training_Registration_Pipeline',
                                     description = 'Pipeline to load/register IoT telemetry data from ADLS Gen2, train a Tensorflow/Keras autoencoder for anomaly detection, and register the trained model if it performs better than the current best model.',
                                     continue_on_step_failure = False)

# ### Optional: Submit a Pipeline Run
# You can create a new experiment  (logical container for pipeline runs) and execute the pipeline. Note: the values of pipeline parameters can be modified when submitting a new run.

experiment = Experiment(ws, 'sample-pipeline-run')
run = experiment.submit(pipeline)
run.wait_for_completion(show_output=True)

StarcoderdataJupyterScriptsDedupFiltered