113 строки
3.3 KiB
Python
113 строки
3.3 KiB
Python
# imports
|
|
import argparse
|
|
|
|
from azureml.core import Workspace
|
|
from azureml.core.compute import ComputeTarget, AmlCompute, AksCompute
|
|
|
|
# setup argparse
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", type=str, default="")
|
|
parser.add_argument("--subscription-id", type=str, default=None)
|
|
parser.add_argument("--workspace-name", type=str, default="default")
|
|
parser.add_argument("--resource-group", type=str, default="azureml-examples")
|
|
parser.add_argument("--location", type=str, default="eastus")
|
|
parser.add_argument("--create-aks", type=bool, default=False)
|
|
parser.add_argument("--create-V100", type=bool, default=False)
|
|
args = parser.parse_args()
|
|
|
|
# constants, variables, parameters, etc.
|
|
amlcomputes = {
|
|
"cpu-cluster": {
|
|
"vm_size": "STANDARD_DS3_V2",
|
|
"min_nodes": 0,
|
|
"max_nodes": 10,
|
|
"idle_seconds_before_scaledown": 1200,
|
|
},
|
|
"gpu-cluster": {
|
|
"vm_size": "STANDARD_NC6",
|
|
"min_nodes": 0,
|
|
"max_nodes": 4,
|
|
"idle_seconds_before_scaledown": 1200,
|
|
},
|
|
"gpu-K80-2": {
|
|
"vm_size": "STANDARD_NC12",
|
|
"min_nodes": 0,
|
|
"max_nodes": 4,
|
|
"idle_seconds_before_scaledown": 1200,
|
|
},
|
|
}
|
|
|
|
v100computes = {
|
|
"gpu-V100-1": {
|
|
"vm_size": "STANDARD_NC6S_V3",
|
|
"min_nodes": 0,
|
|
"max_nodes": 4,
|
|
"idle_seconds_before_scaledown": 1200,
|
|
},
|
|
"gpu-V100-2": {
|
|
"vm_size": "STANDARD_NC12S_V3",
|
|
"min_nodes": 0,
|
|
"max_nodes": 2,
|
|
"idle_seconds_before_scaledown": 1200,
|
|
},
|
|
"gpu-V100-4": {
|
|
"vm_size": "STANDARD_NC24S_V3",
|
|
"min_nodes": 0,
|
|
"max_nodes": 2,
|
|
"idle_seconds_before_scaledown": 1200,
|
|
},
|
|
}
|
|
|
|
akscomputes = {
|
|
"aks-cpu-deploy": {
|
|
"vm_size": "STANDARD_DS3_V2",
|
|
"agent_count": 3,
|
|
},
|
|
"aks-gpu-deploy": {
|
|
"vm_size": "STANDARD_NC6S_V3",
|
|
"agent_count": 3,
|
|
},
|
|
}
|
|
|
|
# create or get Workspace
|
|
try:
|
|
ws = Workspace.from_config(args.config)
|
|
except:
|
|
ws = Workspace.create(
|
|
args.workspace_name,
|
|
subscription_id=args.subscription_id,
|
|
resource_group=args.resource_group,
|
|
location=args.location,
|
|
create_resource_group=True,
|
|
exist_ok=True,
|
|
show_output=True,
|
|
)
|
|
ws.write_config()
|
|
|
|
# create aml compute targets
|
|
for ct_name in amlcomputes:
|
|
if ct_name not in ws.compute_targets:
|
|
compute_config = AmlCompute.provisioning_configuration(**amlcomputes[ct_name])
|
|
ct = ComputeTarget.create(ws, ct_name, compute_config)
|
|
ct.wait_for_completion(show_output=True)
|
|
|
|
# create aml V100 compute targets
|
|
if args.create_V100:
|
|
for ct_name in v100computes:
|
|
if ct_name not in ws.compute_targets:
|
|
compute_config = AmlCompute.provisioning_configuration(
|
|
**v100computes[ct_name]
|
|
)
|
|
ct = ComputeTarget.create(ws, ct_name, compute_config)
|
|
ct.wait_for_completion(show_output=True)
|
|
|
|
# create aks compute targets
|
|
if args.create_aks:
|
|
for ct_name in akscomputes:
|
|
if ct_name not in ws.compute_targets:
|
|
compute_config = AksCompute.provisioning_configuration(
|
|
**akscomputes[ct_name]
|
|
)
|
|
ct = ComputeTarget.create(ws, ct_name, compute_config)
|
|
ct.wait_for_completion(show_output=True)
|