aml-private.azcli

#################################################
# Commands to create a secured AML environment
#   in VNets using private endpoints.
#
# Jose Moreno, August 2022
#################################################

# Variables
rg=aml
location=eastus2
vnet_name=amlvnet
vnet_prefix=10.13.76.0/22
vm_subnet_name=compute
vm_subnet_prefix=10.13.76.0/24
azfw_subnet_name=AzureFirewallSubnet
azfw_subnet_prefix=10.13.77.64/26
ep_subnet_name=ep
ep_subnet_prefix=10.13.77.0/27
test_subnet_name=test
test_subnet_prefix=10.13.77.32/27
aks_subnet_name=aks
aks_subnet_prefix=10.13.77.128/25
storage_blob_ep_name=blobep
storage_file_ep_name=fileep
acr_ep_name=acrep
akv_ep_name=akvep
aml_ep_name=wsep
azfw_name=amlfw
test_vm_size=Standard_B1s

# Function to update UDR back to testing PC
function update_myip() {
    myip=$(curl -s4 ifconfig.co)
    echo "Updating IP to $myip..."
    az network route-table route update --route-table-name aml -g $rg --address-prefix "${myip}/32" --name clientIP -o none
}

# Function to update the SSH key of the Linux VM
function update_ssh_key() {
    vm_name=testvm
    user=$(whoami)
    echo "Updating SSH key for VM $vm_name and user $user..."
    az vm user update -g $rg -u $user -n $vm_name --ssh-key-value "$(< ~/.ssh/id_rsa.pub)" -o none
}

# RG and VNet
echo "Creating RG and VNets..."
az group create -n $rg -l $location -o none
az network vnet create -g $rg -n $vnet_name --address-prefix $vnet_prefix --subnet-name $vm_subnet_name --subnet-prefix $vm_subnet_prefix -l $location -o none
az network vnet subnet create -g $rg --vnet-name $vnet_name -n $ep_subnet_name --address-prefix $ep_subnet_prefix -o none
az network vnet subnet create -g $rg --vnet-name $vnet_name -n $ep_subnet_name --address-prefix $ep_subnet_prefix -o none
az network vnet subnet create -g $rg --vnet-name $vnet_name -n $azfw_subnet_name --address-prefix $azfw_subnet_prefix -o none
az network vnet subnet create -g $rg --vnet-name $vnet_name -n $test_subnet_name --address-prefix $test_subnet_prefix -o none
az network vnet subnet create -g $rg --vnet-name $vnet_name -n $aks_subnet_name --address-prefix $aks_subnet_prefix -o none
# az network vnet subnet update -n $compute_subnet_name -g $rg --vnet-name $vnet_name --disable-private-endpoint-network-policies true

# Log Analytics Workspace
logws_name=$(az monitor log-analytics workspace list -g $rg --query '[0].name' -o tsv)
if [[ -z "$logws_name" ]]
then
    echo "Creating new Log Analytics workspace"
    logws_name=log$RANDOM
    az monitor log-analytics workspace create -n $logws_name -g $rg -o none
else
    echo "Log Analytics workspace $logws_name found"
fi
logws_id=$(az resource list -g $rg -n $logws_name --query '[].id' -o tsv)
logws_customerid=$(az monitor log-analytics workspace show -n $logws_name -g $rg --query customerId -o tsv)

# Storage account
storage_account_name=$(az storage account list -g $rg --query '[0].name' -o tsv)
if [[ -z "$storage_account_name" ]]; then
    storage_account_name=aml$RANDOM
    echo "Creating new storage account $storage_account_name..."
    az storage account create -n $storage_account_name -g $rg --sku Standard_LRS --kind StorageV2 -o none
else
    echo "Storage Account $storage_account_name found in resource group"
fi
# Blob private link
echo "Create private endpoint for blob..."
storage_account_id=$(az storage account show -n $storage_account_name -g $rg -o tsv --query id)
az network private-endpoint create -n $storage_blob_ep_name -g $rg --vnet-name $vnet_name --subnet $ep_subnet_name --private-connection-resource-id $storage_account_id --group-id blob --connection-name blob -o none
dns_zone_name=privatelink.blob.core.windows.net
az network private-dns zone create -n $dns_zone_name -g $rg -o none
az network private-dns link vnet create -g $rg -z $dns_zone_name -n blobLink --virtual-network $vnet_name --registration-enabled false -o none
az network private-endpoint dns-zone-group create --endpoint-name $storage_blob_ep_name -g $rg -n blobzonegroup --zone-name zone1 --private-dns-zone $dns_zone_name -o none
# File private link
echo "Create private endpoint for file..."
az network private-endpoint create -n $storage_file_ep_name -g $rg --vnet-name $vnet_name --subnet $ep_subnet_name --private-connection-resource-id $storage_account_id --group-id file --connection-name file -o none
dns_zone_name=privatelink.file.core.windows.net
az network private-dns zone create -n $dns_zone_name -g $rg -o none
az network private-dns link vnet create -g $rg -z $dns_zone_name -n fileLink --virtual-network $vnet_name --registration-enabled false -o none
az network private-endpoint dns-zone-group create --endpoint-name $storage_file_ep_name -g $rg -n filezonegroup --zone-name zone1 --private-dns-zone $dns_zone_name -o none
# Diagnostic settings
storage_account_id=$(az storage account show -n $storage_account_name -g $rg -o tsv --query id)
storage_blob_id="${storage_account_id}/blobServices/default"
storage_file_id="${storage_account_id}/fileServices/default"
az monitor diagnostic-settings create -n blobdiag --storage-account $storage_account_name --resource $storage_blob_id --workspace $logws_id \
    --metrics '[{"category": "Transaction", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false }, "timeGrain": null}]' \
    --logs '[{"category": "StorageRead", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}, 
            {"category": "StorageWrite", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}},
            {"category": "StorageDelete", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}]' -o none
az monitor diagnostic-settings create -n filediag --storage-account $storage_account_name --resource $storage_file_id --workspace $logws_id \
    --metrics '[{"category": "Transaction", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false }, "timeGrain": null}]' \
    --logs '[{"category": "StorageRead", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}, 
            {"category": "StorageWrite", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}},
            {"category": "StorageDelete", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}]' -o none
# Disable public access
echo "Disabling public network access in storage account $storage_account_name..."
az storage account update -n $storage_account_name -g $rg --public-network-access Disabled -o none

# Azure Key Vault
akv_name=$(az keyvault list -g $rg --query '[0].name' -o tsv)
if [[ -z "$akv_name" ]]; then
    akv_name=aml$RANDOM
    echo "Creating new Azure Key Vault $akv_name..."
    az keyvault create -n $akv_name -g $rg -l $location --public-network-access disabled -o none
else
    echo "Azure Key Vault $akv_name found in resource group"
fi
echo "Creating AKV private link..."
akv_id=$(az keyvault show -n $akv_name -g $rg --query id -o tsv)
az network private-endpoint create -n $akv_ep_name -g $rg --vnet-name $vnet_name --subnet $ep_subnet_name --private-connection-resource-id $akv_id --group-id vault --connection-name vault -o none
dns_zone_name=privatelink.file.vaultcore.azure.net
az network private-dns zone create -n $dns_zone_name -g $rg -o none
az network private-dns link vnet create -g $rg -z $dns_zone_name -n akvLink --virtual-network $vnet_name --registration-enabled false -o none
az network private-endpoint dns-zone-group create --endpoint-name $akv_ep_name -g $rg -n akvzonegroup --zone-name zone1 --private-dns-zone $dns_zone_name -o none
# Logs
az monitor diagnostic-settings create -n akvdiag --resource $akv_id --workspace $logws_id \
    --metrics '[{"category": "AllMetrics", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false }, "timeGrain": null}]' \
    --logs '[{"category": "AuditEvent", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}]' -o none
# Disable network access (if not done at creation time)
# az keyvault update -n $akv_name -g $rg --public-network-access disabled -o none

# Azure Container Registry
acr_name=$(az acr list -g $rg --query '[0].name' -o tsv)
if [[ -z "$acr_name" ]]; then
    acr_name=aml$RANDOM
    echo "Creating new Azure Container Registry $acr_name..."
    az acr create -n $acr_name -g $rg -l $location --sku Premium --public-network-enabled false -o none
else
    echo "Azure Container Registry $acr_name found in resource group"
fi
echo "Creating ACR private link..."
acr_id=$(az acr show -n $acr_name -g $rg --query id -o tsv)
az network private-endpoint create -n $acr_ep_name -g $rg --vnet-name $vnet_name --subnet $ep_subnet_name --private-connection-resource-id $acr_id --group-id registry --connection-name acr -o none
dns_zone_name="privatelink.azurecr.io"
az network private-dns zone create -n $dns_zone_name -g $rg -o none
az network private-dns link vnet create -g $rg -z $dns_zone_name -n acrLink --virtual-network $vnet_name --registration-enabled false -o none
az network private-endpoint dns-zone-group create --endpoint-name $acr_ep_name -g $rg -n acrzonegroup --zone-name zone1 --private-dns-zone $dns_zone_name -o none
# Logs
az monitor diagnostic-settings create -n acrdiag --resource $acr_id --workspace $logws_id \
    --metrics '[{"category": "AllMetrics", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false }, "timeGrain": null}]' \
    --logs '[{"category": "ContainerRegistryRepositoryEvents", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}, 
            {"category": "ContainerRegistryLoginEvents", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}]' -o none
# Enable admin user
az acr update --admin-enabled -n $acr_name -g $rg -o none
acr_usr=$(az acr credential show -n "$acr_name" -g "$rg" --query 'username' -o tsv)
acr_pwd=$(az acr credential show -n "$acr_name" -g "$rg" --query 'passwords[0].value' -o tsv)

# AML workspace
aml_ws_name=$(az ml workspace list -g $rg --query '[0].name' -o tsv)
aml_ws_file=/tmp/amlws.yaml
if [[ -z "$aml_ws_name" ]]; then
    aml_ws_name=aml$RANDOM
    cat <<EOF > ${aml_ws_file}
\$schema: https://azuremlschemas.azureedge.net/latest/workspace.schema.json
name: $aml_ws_name
location: $location
display_name: AML workspace with existing resources
description: This configuration specifies a workspace configuration with existing dependent resources
storage_account: $storage_account_id
container_registry: $acr_id
key_vault: $akv_id
tags:
  purpose: demonstration
EOF
    echo "Creating new Azure ML workspace $aml_ws_name..."
    # az ml workspace create -n $aml_ws_name -g $rg -l $location --public-network-access Disabled --image-build-compute cpucompute \
    #     --container-registry $acr_id -o none
    az ml workspace create -g $rg -f $aml_ws_file -o none
else
    echo "Azure ML workspace $aml_ws_name found in resource group"
fi
echo "Creating AML private link..."
aml_ws_id=$(az ml workspace list -g $rg --query '[0].id' -o tsv)
az network private-endpoint create -n $aml_ep_name -g $rg --vnet-name $vnet_name --subnet $ep_subnet_name --private-connection-resource-id $aml_ws_id --group-id amlworkspace --connection-name amlworkspace -o none
dns_zone_name=privatelink.api.azureml.ms
az network private-dns zone create -n $dns_zone_name -g $rg -o none
az network private-dns link vnet create -g $rg -z $dns_zone_name -n amlLink --virtual-network $vnet_name --registration-enabled false -o none
az network private-endpoint dns-zone-group create --endpoint-name $aml_ep_name -g $rg -n amlzonegroup --zone-name zone1 --private-dns-zone $dns_zone_name -o none
echo "Disabling public access to the workspace..."
az ml workspace update -n $aml_ws_name -g $rg --public-network-access Disabled -o none

# Enabling AML Studio (??)
storage_blob_ep_id=$(az network private-endpoint show -n $storage_blob_ep_name -g $rg --query id -o tsv)
storage_file_ep_id=$(az network private-endpoint show -n $storage_file_ep_name -g $rg --query id -o tsv)
aml_ws_appid=$(az ad sp list --display-name $aml_ws_name --query '[0].appId' -o tsv)
echo "Creating Reader assignment for AML workspace's MSI $aml_ws_appid to the blob endpoint $storage_blob_ep_name..."
az role assignment create --assignee $aml_ws_appid --role "Reader" --scope $storage_blob_ep_id -o none
echo "Creating Reader assignment for AML workspace's MSI $aml_ws_appid to the file endpoint $storage_file_ep_name..."
az role assignment create --assignee $aml_ws_appid --role "Reader" --scope $storage_file_ep_id -o none

# Grant access to ACR
echo "Granting pull access on ACR $acr_name to AML workspace $aml_ws_name..."
az role assignment create --assignee $aml_ws_appid --scope $acr_id --role acrpush -o none
echo "Enabling access to trusted services on ACR $acr_name..."
az acr update -n $acr_name --allow-trusted-services true -o none

# Grant access to AKV
# az keyvault set-policy -n $akv_name --object-id $aml_ws_appid -o none \
#     --secret-permissions all --key-permissions all --certificate-permissions all

# Azure Firewall
azfw_policy_name="${azfw_name}-policy"
echo "Creating Azure Firewall Policy..."
az network firewall policy create -n $azfw_policy_name -g $rg -o none
echo "Creating application rule to allow all FQDNs..."
az network firewall policy rule-collection-group create -n amlruleset --policy-name $azfw_policy_name -g $rg --priority 1000 -o none
az network firewall policy rule-collection-group collection add-filter-collection --policy-name $azfw_policy_name --rule-collection-group-name amlruleset -g $rg \
    --name allowall --collection-priority 200 --action Allow --rule-name allowall --rule-type ApplicationRule --description "AllowAll" \
    --target-fqdns '*' --source-addresses '*' --protocols Http=80 Https=443 -o none
az network firewall policy rule-collection-group collection add-filter-collection --policy-name $azfw_policy_name --rule-collection-group-name amlruleset -g $rg \
    --name ntp --collection-priority 110 --action Allow --rule-name allowNTP --rule-type NetworkRule --description "Egress NTP traffic" \
    --destination-addresses '*' --source-addresses "$vnet_prefix" --ip-protocols UDP --destination-ports "123" -o none
az network firewall policy rule-collection-group collection add-filter-collection --policy-name $azfw_policy_name --rule-collection-group-name amlruleset -g $rg \
    --name winactivation --collection-priority 115 --action Allow --rule-name allowWinActivation --rule-type NetworkRule --description "Egress Windows Activation" \
    --destination-addresses '20.118.99.224' '40.83.235.53' '23.102.135.246' --source-addresses "$vnet_prefix" --ip-protocols Tcp --destination-ports "1688" -o none
# az network firewall policy rule-collection-group collection add-filter-collection --policy-name $azfw_policy_name --rule-collection-group-name amlruleset -g $rg \
#     --name allowallnet --collection-priority 150 --action Allow --rule-name AllowAll --rule-type NetworkRule --description "Allow All" \
#     --destination-addresses '*' --source-addresses '*' --ip-protocols All --destination-ports '*' -o none
echo "Creating Azure Firewall..."
az network firewall create -n $azfw_name -g $rg --policy $azfw_policy_name -l $location -o none
echo "Configuring firewall logs and private IP..."
azfw_id=$(az network firewall show -n $azfw_name -g $rg -o tsv --query id)
az monitor diagnostic-settings create -n mydiag --resource $azfw_id --workspace $logws_id \
    --metrics '[{"category": "AllMetrics", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false }, "timeGrain": null}]' \
    --logs '[{"category": "AzureFirewallApplicationRule", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}, 
            {"category": "AzureFirewallNetworkRule", "enabled": true, "retentionPolicy": {"days": 0, "enabled": false}}]' -o none
echo "Configuring Azure Firewall's IP..."
azfw_pip_name="${azfw_name}-pip"
az network public-ip create -g $rg -n $azfw_pip_name --sku standard --allocation-method static -l $location -o none
az network firewall ip-config create -f $azfw_name -n azfw-ipconfig -g $rg --public-ip-address $azfw_pip_name --vnet-name $vnet_name -o none
az network firewall update -n $azfw_name -g $rg -o none
azfw_private_ip=$(az network firewall show -n $azfw_name -g $rg -o tsv --query 'ipConfigurations[0].privateIpAddress') && echo "$azfw_private_ip"

# UDRs to send all traffic through the Azure Firewall
# See https://docs.microsoft.com/en-us/azure/machine-learning/how-to-secure-training-vnet
echo "Creating route table..."
az network route-table create -n aml -g $rg -l $location -o none
my_ip=$(curl -s4 ifconfig.co)
az network route-table route create -n clientIP --route-table-name aml -g $rg --next-hop-type Internet --address-prefix "$my_ip/32" -o none
az network route-table route create -n defaultRoute --route-table-name aml -g $rg --next-hop-type VirtualAppliance --address-prefix "0.0.0.0/0" --next-hop-ip-address $azfw_private_ip -o none
# az network route-table route create -n vnetTraffic --route-table-name aml -g $rg --next-hop-type VirtualAppliance --address-prefix $vnet_prefix --next-hop-ip-address $azfw_private_ip -o none
# az network route-table route delete -n vnetTraffic --route-table-name aml -g $rg -o none
az network route-table route create -n azureBatch --route-table-name aml -g $rg --next-hop-type VirtualAppliance --address-prefix BatchNodeManagement --next-hop-ip-address $azfw_private_ip -o none
az network route-table route create -n AML --route-table-name aml -g $rg --next-hop-type VirtualAppliance --address-prefix AzureMachineLearning --next-hop-ip-address $azfw_private_ip -o none

rt_id=$(az network route-table show -n aml -g $rg --query id -o tsv)
az network vnet subnet update -g $rg --vnet-name $vnet_name -n $vm_subnet_name --route-table $rt_id -o none
az network vnet subnet update -g $rg --vnet-name $vnet_name -n $test_subnet_name --route-table $rt_id -o none

# Test VM (linux)
az vm create -n testvm -g $rg --image ubuntuLTS --generate-ssh-keys --public-ip-address testvm-pip --size $test_vm_size \
    --vnet-name $vnet_name --subnet $test_subnet_name --public-ip-sku Standard -o none
testvm_pip=$(az network public-ip show -n testvm-pip -g $rg --query ipAddress -o tsv)
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "sudo apt update"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "curl -sL https://aka.ms/InstallAzureCLIDeb | sudo bash"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az extension add -n ml"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az login --use-device-code"        # !!!!! Interactive, don't run as a script !!!!!
subscription_id=$(az account show --query id -o tsv)
echo "Setting subscription to $subscription_id..."
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az account set -n $subscription_id"

# Test VM (Windows
win_user=$(whoami)
win_password='Microsoft123!'
win_vm_name=testvmwin
win_vm_sku=Standard_B2ms
win_pip_name=testvmwin-pip
az vm create -n $win_vm_name -g $rg --image win2019datacenter --admin-username $win_user --admin-password $win_password --size $win_vm_sku \
   --vnet-name $vnet_name --subnet $vm_subnet_name --public-ip-address $win_pip_name -o none

# Create compute cluster
compute_type=AmlCompute        # Can be 'ComputeInstance' or 'AmlCompute'
compute_name=$(az ml compute list -w $aml_ws_name -g $rg --query '[0].name' -o tsv)
if [[ -z "$compute_name" ]]; then
    compute_name=cpucompute$RANDOM
    # compute_vm_size=STANDARD_DS3_v2  # $0.23/h
    compute_vm_size=STANDARD_DS11_v2 # $0.15/h
    echo "Disabling private link policies in subnet $subnet_name..."
    az network vnet subnet update -n $vm_subnet_name --vnet-name $vnet_name -g $rg --disable-private-link-service-network-policies -o none      # Required for clusters without PIP
    az network vnet subnet update -n $vm_subnet_name --vnet-name $vnet_name -g $rg --disable-private-endpoint-network-policies -o none          # Required for clusters without PIP
    echo "Creating new compute target $compute_name of type $compute_type..."
    # az ml compute delete -w $aml_ws_name -g $rg -n $compute_name -y   # If you need to delete it
    # ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml compute create -w $aml_ws_name -g $rg -n $compute_name -t AmlCompute --size $compute_vm_size --vnet-name $vnet_name --subnet $vm_subnet_name --min-instances 0 --max-instances 1 -o none"
    # ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml compute create -w $aml_ws_name -g $rg -n $compute_name -t ComputeInstance --size $compute_vm_size --vnet-name $vnet_name --subnet $vm_subnet_name -o none"
    if [[ "$compute_type" == "ComputeInstance" ]]; then
        az ml compute create -w $aml_ws_name -g $rg -n $compute_name -t ComputeInstance --size $compute_vm_size --vnet-name $vnet_name --subnet $vm_subnet_name -o none
    elif [[ "$compute_type" == "AmlCompute" ]]; then
        az ml compute create -w $aml_ws_name -g $rg -n $compute_name -t AmlCompute --size $compute_vm_size --vnet-name $vnet_name --subnet $vm_subnet_name --min-instances 0 --max-instances 1 -o none
    else
        echo "Compute type $compute_type not recognized"
    fi
    echo "Updating cluster to use compute instance $compute_name to build images..."
    az ml workspace update -n $aml_ws_name -g $rg --image-build-compute $compute_name -o none
else
    echo "Compute target $compute_name found in AML workspace $aml_ws_name"
fi

# Dataset
function load_dataset() {
    dataset_name=$1
    dataset_local_file="/tmp/${dataset_name}.csv"
    wget "https://azuremlexamples.blob.core.windows.net/datasets/${dataset_name}.csv" -O $dataset_local_file
    dataset_yaml=/tmp/dataset.yml
    cat <<EOF > ${dataset_yaml}
\$schema: https://azuremlschemas.azureedge.net/latest/asset.schema.json
name: ${dataset_name}
version: 1
path: $dataset_local_file
type: uri_file
description: ${dataset_name} dataset.
EOF
    scp "${dataset_local_file}" "${testvm_pip}:${dataset_local_file}"
    scp "${dataset_yaml}" "${testvm_pip}:${dataset_yaml}"
    ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml data create -w $aml_ws_name -g $rg -f $dataset_yaml -o none"
}
load_dataset iris
load_dataset diabetes

# Environment
conda_env_url='https://raw.githubusercontent.com/MicrosoftLearning/mslearn-aml-cli/master/Allfiles/Labs/01/conda-envs/basic-env-cpu.yml'
env_yaml=/tmp/environment.yml
conda_yaml=/tmp/basic-env-cpu.yml
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "wget $conda_env_url -O $conda_yaml"
cat <<EOF > ${env_yaml}
\$schema: https://azuremlschemas.azureedge.net/latest/environment.schema.json
name: basic-env-scikit2
image: mcr.microsoft.com/azureml/openmpi3.1.2-ubuntu18.04
conda_file: $conda_yaml
description: Environment created from a Docker image plus Conda environment.
EOF
scp "${env_yaml}" "${testvm_pip}:${env_yaml}"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml environment create -w $aml_ws_name -g $rg -f $env_yaml -o none"

# Hello world job
job_file=/tmp/job_hello.yml
cat <<EOF > ${job_file}
\$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
command: echo "hello world"
display_name: hello-world-$compute_name
experiment_name: hello-world-$compute_name
environment:
  image: library/python:latest
compute: azureml:$compute_name
display_name: hello-world-$compute_name
EOF
run_output=$(az ml job create -w $aml_ws_name -g $rg -f $job_file)
run_id=$(echo $run_output | jq -r '.name') && echo "Running job $run_id..."

# Training
# Potentially look at https://github.com/MicrosoftLearning/mslearn-aml-cli/tree/master/Allfiles/Labs/02/basic-job

# Training with AML dataset...
job_file=/tmp/job_diabetes.yml
cat <<EOF > ${job_file}
\$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
code: src
command: >-
  python main.py 
  --diabetes-csv \${{inputs.diabetes}}
inputs:
  diabetes:
    path: azureml:diabetes:1
    mode: ro_mount
environment: azureml:basic-env-scikit@latest
compute: azureml:$compute_name
experiment_name: diabetes-data-example
description: Train a classification model on diabetes data using a registered dataset as input.
display_name: diabetes-$compute_name
EOF
training_script_local_file="/tmp/main.py"
training_script_url="https://raw.githubusercontent.com/MicrosoftLearning/mslearn-aml-cli/master/Allfiles/Labs/02/input-data-job/src/main.py"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "mkdir src"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "wget $training_script_url -O ./src/main.py"
scp "${job_file}" "${testvm_pip}:${job_file}"
job_file_basename=$(basename "$job_file")
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "cp $job_file ./$job_file_basename"
run_output=$(ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml job create -w $aml_ws_name -g $rg -f $job_file_basename")
run_id=$(echo $run_output | jq -r '.name') && echo "Running job $run_id..."

# Training with local data file
training_script=/tmp/diabetes_experiment.py
training_script_basename=$(basename "$training_script")
cat <<EOF > ${training_script}
# Import libraries
from azureml.core import Run
import pandas as pd
import numpy as np
import joblib
import os
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import roc_auc_score
from sklearn.metrics import roc_curve
# Get the experiment run context
run = Run.get_context()
# load the diabetes dataset
print("Loading Data...")
diabetes = pd.read_csv('diabetes.csv')
# Separate features and labels
X, y = diabetes[['Pregnancies','PlasmaGlucose','DiastolicBloodPressure','TricepsThickness','SerumInsulin','BMI','DiabetesPedigree','Age']].values, diabetes['Diabetic'].values
# Split data into training set and test set
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.30, random_state=0)
# Set regularization hyperparameter
reg = 0.01
# Train a logistic regression model
print('Training a logistic regression model with regularization rate of', reg)
run.log('Regularization Rate',  np.float(reg))
model = LogisticRegression(C=1/reg, solver="liblinear").fit(X_train, y_train)
# calculate accuracy
y_hat = model.predict(X_test)
acc = np.average(y_hat == y_test)
print('Accuracy:', acc)
run.log('Accuracy', np.float(acc))
# calculate AUC
y_scores = model.predict_proba(X_test)
auc = roc_auc_score(y_test,y_scores[:,1])
print('AUC: ' + str(auc))
run.log('AUC', np.float(auc))
# Save the trained model in the outputs folder
os.makedirs('outputs', exist_ok=True)
joblib.dump(value=model, filename='outputs/diabetes_model.pkl')
run.complete()
EOF
# Job file:
job_file=/tmp/job_diabetes.yml
cat <<EOF > ${job_file}
\$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
code: src
command: >-
  python ${training_script_basename}
environment: azureml:basic-env-scikit@latest
compute: azureml:$compute_name
experiment_name: diabetes-data-example
description: Train a classification model on diabetes data using a data file uploaded along the script.
display_name: diabetes-localfile-$compute_name
EOF
# Create src folder in local VM, in case it doesnt exist
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "mkdir src"
# Option 1: get training script from the Web
# training_script_url="https://raw.githubusercontent.com/MicrosoftLearning/mslearn-aml-cli/master/Allfiles/Labs/02/basic-job/src/main.py"
# ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "wget $training_script_url -O ./src/main.py"
# Option 2: copy from local file system
scp "${training_script}" "${testvm_pip}:/home/$(whoami)/src/${training_script_basename}"
# Get data from the web
diabetes_data_url='https://github.com/MicrosoftLearning/mslearn-dp100/raw/main/data/diabetes.csv'
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "wget $diabetes_data_url -O ./src/diabetes.csv"
# Copy job file from local file system
job_file_basename=$(basename "$job_file")
scp "${job_file}" "${testvm_pip}:/home/$(whoami)/${job_file_basename}"
# Run job
run_output=$(ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml job create -w $aml_ws_name -g $rg -f $job_file_basename")
run_id=$(echo $run_output | jq -r '.name') && echo "Running job $run_id..."

# Register model
model_name=diabetesmodel
# az ml model create -n $model_name -w $aml_ws_name -g $rg --version 1 --path "azureml://jobs/${run_id}/outputs/artifacts/paths/model/"
az ml model create -n $model_name -w $aml_ws_name -g $rg --version 1 --path "azureml://jobs/${run_id}/outputs/diabetes_model.pkl"

# Deploy to a managed endpoint
# 1. Endpoint
endpoint_name=diabetes$RANDOM
endpoint_file=/tmp/endpoint_diabetes.yml
cat <<EOF > ${endpoint_file}
\$schema: https://azuremlschemas.azureedge.net/exp/managedOnlineEndpoint.schema.json
name: diabetes-endpoint
auth_mode: key
public_network_access: disabled     # This is ingress
# egress_public_network_access: disabled  # Note that egress access is a deployment property, not an endpoint property
EOF
scp "${endpoint_file}" "${testvm_pip}:${endpoint_file}"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml online-endpoint create -n $endpoint_name -w $aml_ws_name -g $rg -f $endpoint_file -o none"
# 2. Deployment
scoring_script_url=https://raw.githubusercontent.com/MicrosoftDocs/pipelines-azureml/master/models/diabetes/score.py
deployment_file=/tmp/deployment_diabetes.yml
cat <<EOF > ${deployment_file}
\$schema: https://azuremlschemas.azureedge.net/latest/managedOnlineDeployment.schema.json
name: diabetes-deployment
endpoint_name: ${endpoint_name}
model: azureml:${model_name}:1
instance_type: Standard_F2s_v2
instance_count: 1
environment: azureml:basic-env-scikit:1
egress_public_network_access: disabled
code_configuration:
  code: /home/$(whoami)/score/
  scoring_script: score.py
EOF
scoring_script=/tmp/score.py
cat <<EOF > ${scoring_script}
import json
import numpy as np
import pickle
from sklearn.linear_model import Ridge
from azureml.core.model import Model
from inference_schema.schema_decorators import input_schema, output_schema
from inference_schema.parameter_types.numpy_parameter_type import NumpyParameterType
def init():
    global model
    model_path = Model.get_model_path('$model_name')

    with open(model_path, 'rb') as file:
        model = pickle.load(file)
input_sample = np.array([[10, 9, 8, 7, 6, 5, 4, 3, 2, 1]])
output_sample = np.array([3726.995])
@input_schema('data', NumpyParameterType(input_sample))
@output_schema(NumpyParameterType(output_sample))
def run(data):
    try:
        result = model.predict(data)
        return result.tolist()
    except Exception as e:
        error = str(e)
        return error
EOF
scp "${deployment_file}" "${testvm_pip}:${deployment_file}"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "mkdir score"
# Option 1: download score.py from Internet
# ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "wget $scoring_script_url -O ./score/score.py"
# ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "sudo sed -i \"s/diabetes-model/$model_name/\" ./score/score.py"  # Fix the scoring script (that version is hardwired to a specific model name)
# Option 2: copy from local file
scp "${scoring_script}" "${testvm_pip}:/home/$(whoami)/score/score.py"
# Create deployment
deployment_name=diabetes$RANDOM
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml online-deployment create -n $deployment_name -w $aml_ws_name -g $rg -f $deployment_file --all-traffic -o none"
# This error happens when using public outbound for the endpoint and the WS is private:
# Your online endpoint deployment failed because your workspace enables a private link feature and it blocks 
#   your online managed endpoint to get data from your workspace. Please enable workspace public network access to avoid this issue
# Unidentified error:
# Creating/updating online deployment diabetes23941 ERROR: (None) ResourceNotReady: User container has crashed or terminated.
#   Please see troubleshooting guide, available here: https://aka.ms/oe-tsg#error-resourcenotready

# Create inference cluster
# https://docs.microsoft.com/en-us/azure/machine-learning/how-to-attach-kubernetes-anywhere?tabs=deploy-extension-with-cli%2Ccli
# 1. Create AKS cluster
aks_subnet_id=$(az network vnet subnet show -n $aks_subnet_name --vnet-name $vnet_name -g $rg --query id -o tsv)
rt_id=$(az network route-table show -n aml -g $rg --query id -o tsv)
az network vnet subnet update --ids $aks_subnet_id --route-table $rt_id -o none   # kubenet/BYO-RT not supported with MSI
aks_node_size=Standard_B2ms    # Some possible values: Standard_B2ms, Standard_D2_v3
network_plugin=azure           # Since BYO-RT not supported with MSI, but AML requires MSI at this point, using Azure CNI
k8s_versions=$(az aks get-versions -l $location -o json)
k8s_version=$(echo $k8s_versions | jq '.orchestrators[]' | jq -rsc 'sort_by(.orchestratorVersion) | reverse[0] | .orchestratorVersion')
echo "Latest supported k8s version in $rg_location is $k8s_version (in preview)"
aks_service_cidr=10.0.0.0/16
aks_name=$(az aks list -g $rg --query '[0].name' -o tsv)
if [[ -z "$aks_name" ]]; then
    echo "Creating AKS cluster..."
    aks_name=amlaks$RANDOM
    az aks create -g $rg -n $aks_name -l $location -o none \
        -c 1 -s $aks_node_size -k $k8s_version --generate-ssh-keys -u $(whoami) \
        --enable-managed-identity \
        --network-plugin $network_plugin --vnet-subnet-id $aks_subnet_id --service-cidr $aks_service_cidr \
        --network-policy 'calico' --load-balancer-sku Standard --outbound-type userDefinedRouting \
        --node-resource-group "$aks_name"-iaas-"$RANDOM" \
        --enable-private-cluster --disable-public-fqdn
else
    echo "AKS cluster $aks_name found in resource group"
fi
az aks enable-addons --addons azure-policy -n $aks_name -g $rg -o none
# 2. Deploy AzureML extension
az k8s-extension create --name amlextension --extension-type Microsoft.AzureML.Kubernetes -o none \
    --config enableTraining=True enableInference=True inferenceRouterServiceType=LoadBalancer internalLoadBalancerProvider=azure allowInsecureConnections=True inferenceLoadBalancerHA=False \
    --cluster-type managedClusters --cluster-name $aks_name --resource-group $rg --scope cluster
# 3. Attach cluster to AML workspace
aks_id=$(az aks show -n $aks_name -g $rg --query id -o tsv) && echo "AKS cluster ID is $aks_id"
aks_ns=aml
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml compute attach -w $aml_ws_name -g $rg --type Kubernetes -n $aks_name --resource-id $aks_id --identity-type SystemAssigned --namespace $aks_ns -o none"

# Deploy model to AKS inferencing cluster
endpoint_name=aksendpoint
endpoint_file=/tmp/aksendpoint_diabetes.yml
scoring_script_url=https://raw.githubusercontent.com/MicrosoftDocs/pipelines-azureml/master/models/diabetes/score.py
cat <<EOF > ${endpoint_file}
name: blue
app_insights_enabled: false
endpoint_name: $endpoint_name
model: 
  path: /home/$(whoami)/${model_name}/model/model.pkl
code_configuration:
  code: /home/$(whoami)/score/
  scoring_script: score.py
environment: 
  conda_file: file:/home/$(whoami)/${model_name}/model/conda.yml
  image: mcr.microsoft.com/azureml/openmpi3.1.2-ubuntu18.04:20210727.v1
EOF
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az ml model download -n $model_name -w $aml_ws_name -g $rg -v 1 -p ./" # This will create the folder $model_name in ~/
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "mkdir score"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "wget $scoring_script_url -O ./score/score.py"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "sudo sed -i \"s/diabetes-model/$model_name/\" ./score/score.py"
scp "${endpoint_file}" "${testvm_pip}:${endpoint_file}"

#################
#  Diagnostics  #
#################

# Jump host
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "ip a"

# Routing
az network route-table route list --route-table aml -g $rg -o table
testvm_nic_id=$(az vm show -n testvm -g $rg --query 'networkProfile.networkInterfaces[0].id' -o tsv)
az network nic show-effective-route-table --ids $testvm_nic_id

# Test DNS resolution
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "nslookup ${storage_account_name}.blob.core.windows.net"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "nslookup ${storage_account_name}.file.core.windows.net"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "nslookup ${acr_name}.azurecr.io"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "nslookup ${akv_name}.file.vaultcore.azure.net"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "nslookup ${ml_ws_name}.api.azureml.ms"

# AML
az ml compute list -w $aml_ws_name -g $rg -o table
az ml job list -w $aml_ws_name -g $rg -o table
az ml data list -w $aml_ws_name -g $rg -o table
az ml environment list -w $aml_ws_name -g $rg -o table
az ml model list -w $aml_ws_name -g $rg -o table
az ml online-endpoint list -w $aml_ws_name -g $rg -o table
az ml online-deployment list -w $aml_ws_name -g $rg -e $endpoint_name -o table
az ml online-deployment get-logs -n diabetes23941 -w $aml_ws_name -g $rg -e $endpoint_name -o table

# ACR
# az acr repository list -n $acr_name -g $rg -u $acr_usr -p $acr_pwd -o table
# ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az acr repository list -n $acr_name -u $acr_usr -p $acr_pwd -o table"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az acr repository list -n $acr_name -o table"
acr_repo_name=$(ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az acr repository list -n $acr_name -o tsv --query '[0]'")
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az acr repository show-tags --repository $acr_repo_name -n $acr_name -o table"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az acr repository show --repository $acr_repo_name -n $acr_name"


# Storage
az storage container list --account-name $storage_account_name --auth-mode login -o table

# AKS
az aks list -g $rg -o table
az k8s-extension show --name amlextension --cluster-type managedClusters --cluster-name $aks_name --resource-group $rg
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "sudo az aks install-cli"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "az aks get-credentials -n $aks_name -g $rg --overwrite --admin"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "kubectl get ns"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "kubectl get svc -A"
ssh -n -o BatchMode=yes -o StrictHostKeyChecking=no $testvm_pip "kubectl get svc -n azureml"

# Logs
logs_query='AzureDiagnostics 
| distinct ResourceType'
logs_query='AzureDiagnostics 
| distinct Category'
az monitor log-analytics query -w $logws_customerid --analytics-query $logs_query -o tsv

# AKV logs
akv_logs_query='AzureDiagnostics 
| where ResourceType == "VAULTS" and Category == "AuditEvent"
| where TimeGenerated >= ago(2h) 
| project TimeGenerated, ResultType, OperationName, CallerIPAddress, msg_s, Message
| take 20'
az monitor log-analytics query -w $logws_customerid --analytics-query $akv_logs_query -o tsv

# Firewall net rule logs
fw_net_logs_query='AzureDiagnostics
| where Category == "AzureFirewallNetworkRule"
| where TimeGenerated >= ago(12h) 
| parse msg_s with * "Action: " Action "." *
| where Action == "Deny"
| project TimeGenerated, msg_s, Action
| take 20 '
az monitor log-analytics query -w $logws_customerid --analytics-query $fw_net_logs_query -o tsv

# Firewall App Rule Logs
fw_app_logs_query='AzureDiagnostics 
| where ResourceType == "AZUREFIREWALLS" 
| where Category == "AzureFirewallApplicationRule" 
| where TimeGenerated >= ago(1h) 
| project TimeGenerated, msg_s
| take 20'
az monitor log-analytics query -w $logws_customerid --analytics-query $fw_app_logs_query -o tsv

####################
#    Start/Stop    #
####################

# Stop all
function stop_lab() {
    echo "Stopping AKS..."
    aks_name=$(az aks list -g $rg --query '[0].name' -o tsv)
    az aks stop -n $aks_name -g $rg --no-wait -o none
    echo "Stopping VMs..."
    stop_vms
    echo "Stopping AML compute..."
    stop_aml_compute
    echo "Stopping Azure Firewall..."
    stop_firewall
}

# Stop VMs
function stop_vms() {
    vm_list=$(az vm list -o tsv -g "$rg" --query "[].name")
    while IFS= read -r vm_name; do
        az vm deallocate -g $rg -n "$vm_name" --no-wait -o none
    done <<< "$vm_list"
}

# Stop AML compute
function stop_aml_compute() {
    compute_list=$(az ml compute list -o tsv -w $aml_ws_name -g "$rg" --query "[].name")
    while IFS= read -r compute_name; do
        az ml compute stop -g $rg -w $aml_ws_name -n "$compute_name" --no-wait -o none
    done <<< "$compute_list"
}

# Start
function start_lab() {
    aks_name=$(az aks list -g $rg --query '[0].name' -o tsv)
    az aks start -n $aks_name -g $rg --no-wait -o none
    start_vms
    start_firewall
    start_aml_compute
}

# Start VMs
function start_vms() {
    vm_list=$(az vm list -o tsv -g "$rg" --query "[].name")
    while IFS= read -r vm_name; do
        az vm start -g $rg -n "$vm_name" --no-wait -o none
    done <<< "$vm_list"
}

# Start AML compute
function start_aml_compute() {
    compute_list=$(az ml compute list -o tsv -w $aml_ws_name -g "$rg" --query "[].name")
    while IFS= read -r compute_name; do
        az ml compute start -g $rg -w $aml_ws_name -n "$compute_name" -o none
    done <<< "$compute_list"
}

# Functions to start/stop the firewall
function stop_firewall() {
    azfw_name=$(az network firewall list -g $rg --query '[0].name' -o tsv)
    azfw_ipconfig_name=$(az network firewall show -n $azfw_name -g $rg --query 'ipConfigurations[0].name' -o tsv)
    az network firewall ip-config delete -f $azfw_name -n $azfw_ipconfig_name -g $rg -o none
    az network firewall update -n $azfw_name -g $rg -o none
}
function start_firewall() {
    azfw_name=$(az network firewall list -g $rg --query '[0].name' -o tsv)
    azfw_ipconfig_name="${azfw_name}-ipconfig"
    az network firewall ip-config create -f $azfw_name -n azfw-ipconfig -g $rg --public-ip-address $azfw_pip_name --vnet-name $vnet_name -o none
    az network firewall update -n $azfw_name -g $rg -o none
}

# stop_lab
# start_lab

####################
# Cleanup !DANGER! #
####################

function cleanup_all() {
    az group delete -n $rg -y
}
# cleanup_all