def get_conn(self):
conn = self.get_connection(self.conn_id)
key_path = conn.extra_dejson.get('key_path', False)
if key_path:
if key_path.endswith('.json'):
self.log.info('Getting connection using a JSON key file.')
return get_client_from_auth_file(ContainerInstanceManagementClient,
key_path)
else:
raise AirflowException('Unrecognised extension for key file.')
if os.environ.get('AZURE_AUTH_LOCATION'):
key_path = os.environ.get('AZURE_AUTH_LOCATION')
if key_path.endswith('.json'):
self.log.info('Getting connection using a JSON key file.')
return get_client_from_auth_file(ContainerInstanceManagementClient,
key_path)
else:
raise AirflowException('Unrecognised extension for key file.')
credentials = ServicePrincipalCredentials(
client_id=conn.login,
secret=conn.password,
tenant=conn.extra_dejson['tenantId']
)
subscription_id = conn.extra_dejson['subscriptionId']
return ContainerInstanceManagementClient(credentials, str(subscription_id))
def create_aci_client(secrets: Secrets):
credentials = create_service_principal_credentials(secrets)
return ContainerInstanceManagementClient(
credentials,
secrets.subscription_id,
base_url=None,
)
def get_aci_client():
from azure.common.credentials import ServicePrincipalCredentials
from azure.mgmt.containerinstance import ContainerInstanceManagementClient
credentials = ServicePrincipalCredentials(
client_id = AUTH_CLIENTID,
secret = AUTH_CLIENTSECRET,
tenant = AUTH_TENANT
)
return ContainerInstanceManagementClient(credentials, SUBSCRIPTION_ID)
def setup_connection(endpoint, auth_data):
subscription_id = auth_data['subscription_id']
credentials = ServicePrincipalCredentials(
client_id=auth_data['client_id'],
secret=auth_data['client_secret'],
tenant=auth_data['tenant_id'])
resource_client = ResourceManagementClient(credentials, subscription_id)
network_client = NetworkManagementClient(credentials, subscription_id)
aci_client = ContainerInstanceManagementClient(credentials,
subscription_id)
return (subscription_id, resource_client, network_client, aci_client)
def connect_azure(azure):
# Connect to Azure
azure_context = AzureContext(
subscription_id = azure['azure_subscription_id'],
client_id = azure['azure_client_id'],
client_secret = azure['azure_client_secret'],
tenant = azure['azure_tenant_id']
)
# construct the clients
resource_client = ResourceManagementClient(azure_context.credentials, azure_context.subscription_id)
client = ContainerInstanceManagementClient(azure_context.credentials, azure_context.subscription_id)
return resource_client, client
def _az_connect():
azure_context = AzureContext(
subscription_id=Variable.get("AZURE_SUBSCRIPTION_ID"),
client_id=Variable.get("AZURE_CLIENT_ID"),
client_secret=Variable.get("AZURE_CLIENT_SECRET"),
tenant=Variable.get("AZURE_TENANT_ID"))
# construct the clients
resource_client = ResourceManagementClient(azure_context.credentials,
azure_context.subscription_id)
client = ContainerInstanceManagementClient(azure_context.credentials,
azure_context.subscription_id)
return resource_client, client
def __init__(self, config=None, resource_group_name=None):
"""Initialize a :class:`~.Plugin`.
"""
self.logger = logging.getLogger('Plugin')
logging.basicConfig(level=logging.INFO)
try:
subscription_id = os.environ.get('AZURE_SUBSCRIPTION_ID')
credentials = self._login_azure_app_token()
# Create a Resource Management client
self.resource_client = ResourceManagementClient(credentials, subscription_id)
self.client = ContainerInstanceManagementClient(credentials, subscription_id)
except Exception:
logging.error('Credentials were not valid or were not found.')
raise DriverError
self.resource_group_name = resource_group_name
def main(req: func.HttpRequest) -> func.HttpResponse:
cg_name = req.route_params.get('cg-name')
try:
# tenant_id = os.environ["AZURE_TENANT_ID"]
# client_id = os.environ["AZURE_CLIENT_ID"]
# client_secret = os.environ["AZURE_CLIENT_SECRET"]
subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"]
resource_group = os.environ["AZURE_RESOURCE_GROUP"]
except KeyError as ex:
return func.HttpResponse(body=json.dumps(
{"message": "Error loading environment variable: " + str(ex)}),
headers={"Content-Type": "application/json"},
status_code=500)
credentials = CredentialWrapper()
aci_client = ContainerInstanceManagementClient(
subscription_id=subscription_id, credentials=credentials)
res_client = ResourceManagementClient(subscription_id=subscription_id,
credentials=credentials)
try:
cg = aci_client.container_groups.get(resource_group, cg_name)
except CloudError as ex:
# container group not found
return func.HttpResponse(body=json.dumps({
"message": ex.message,
}),
headers={"Content-Type": "application/json"},
status_code=400)
aci_client.container_groups.delete(resource_group, cg_name)
return func.HttpResponse(body=json.dumps({
"message": f"Container Group {cg_name} deleted",
"name": cg_name
}),
headers={"Content-Type": "application/json"},
status_code=200)
def _get_ci_client(config):
return ContainerInstanceManagementClient(
_get_credentials(config), config.subscription_id
)
def main(req: func.HttpRequest) -> func.HttpResponse:
polling_param = req.params.get("polling")
polling = True if (polling_param != None and polling_param.lower() == "true") else False
try:
body_yaml = req.get_body().decode("utf-8")
cg_definition = yaml.safe_load(body_yaml)
except BaseException as ex:
return func.HttpResponse(
body=json.dumps({
"message": "Error while parsing yaml:\n\n" + str(ex)
}),
headers={
"Content-Type": "application/json"
},
status_code=400
)
try:
# tenant_id = os.environ["AZURE_TENANT_ID"]
# client_id = os.environ["AZURE_CLIENT_ID"]
# client_secret = os.environ["AZURE_CLIENT_SECRET"]
subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"]
resource_group = os.environ["AZURE_RESOURCE_GROUP"]
except KeyError as ex:
return func.HttpResponse(
body=json.dumps({
"message": "Error loading environment variable: " + str(ex)
}),
headers={
"Content-Type": "application/json"
},
status_code=500
)
credentials = DefaultAzureCredential()
aci_client = ContainerInstanceManagementClient(
subscription_id=subscription_id,
credential=credentials
)
res_client = ResourceManagementClient(
subscription_id=subscription_id,
credential=credentials
)
cg_name = cg_definition.get('name', None)
location = cg_definition.get('location', None) or res_client.resource_groups.get(resource_group).location
api_version = cg_definition.get('apiVersion', None) or aci_client._config.api_version
cg_definition['location'] = location
cg_definition['apiVersion'] = api_version
if cg_name == None:
return func.HttpResponse(
body=json.dumps({
"message": "Name property missing in yaml definition"
}),
headers={
"Content-Type": "application/json"
},
status_code=400
)
try:
cg = aci_client.container_groups.get(resource_group, cg_name)
if cg.instance_view.state not in ["Succeeded", "Stopped", "Failed"]:
# Container Group already running
return func.HttpResponse(
body=json.dumps({
"message": f"Container Group {cg_name} is already running",
"name": cg_name,
"state": cg.instance_view.state
}),
headers={
"Content-Type": "application/json"
},
status_code=400
)
except AzureError:
# container group not found, thats fine
pass
try:
# create or update container group
res_client.resources.begin_create_or_update(
resource_group_name=resource_group,
resource_provider_namespace="Microsoft.ContainerInstance",
parent_resource_path="",
resource_type="containerGroups",
resource_name=cg_name,
api_version=api_version,
parameters=cg_definition
)
# start container group and wait until it's truly started
aci_client.container_groups.begin_start(resource_group, cg_name)
cg = aci_client.container_groups.get(resource_group, cg_name)
while cg.instance_view.state not in ["Pending", "Running"]:
aci_client.container_groups.begin_start(resource_group, cg_name)
cg = aci_client.container_groups.get(resource_group, cg_name)
time.sleep(5)
except AzureError as ex:
return func.HttpResponse(
body=json.dumps({
"message": ex.message
}),
headers={
"Content-Type": "application/json"
},
status_code=400
)
# redirect client to status function
if polling == True:
return func.HttpResponse(
body=json.dumps({
"message": f"Container Group {cg_name} started",
"name": cg_name,
"state": cg.instance_view.state
}),
headers={
"Content-Type": "application/json",
"Location": f"https://{os.environ['WEBSITE_HOSTNAME']}/api/status/{cg_name}"
},
status_code=202
)
# just return success
else:
return func.HttpResponse(
body=json.dumps({
"message": f"Container Group {cg_name} started",
"name": cg_name,
"state": cg.instance_view.state
}),
headers={
"Content-Type": "application/json"
},
status_code=200
)
from azureml.core.run import Run, _OfflineRun
# ContainerInstanceManagementClient has not been updated to the new Auth model
# so we use a wrapper instead of DefaultAzureCredential
from cred_wrapper import CredentialWrapper
# from azure.identity import DefaultAzureCredential
# Get workspace
run = Run.get_context()
ws = None
if type(run) == _OfflineRun:
ws = Workspace.from_config()
else:
ws = run.experiment.workspace
# Get where this workspace is located
subscription_id = ws.subscription_id
resource_group = ws.resource_group
# Acquire a credential object
credential = CredentialWrapper() # DefaultAzureCredential()
# https://docs.microsoft.com/en-us/python/api/azure-mgmt-containerinstance/azure.mgmt.containerinstance.containerinstancemanagementclient?view=azure-python
aci_client = ContainerInstanceManagementClient(credential, subscription_id)
# https://azuresdkdocs.blob.core.windows.net/$web/python/azure-mgmt-containerinstance/2.0.0/azure.mgmt.containerinstance.operations.html#azure.mgmt.containerinstance.operations.ContainerGroupsOperations
container_groups = aci_client.container_groups.list_by_resource_group(resource_group)
for container_group in container_groups:
print(f"Stopping {container_group.name}")
aci_client.container_groups.stop(resource_group, container_group.name)
def main(req: func.HttpRequest) -> func.HttpResponse:
cg_name = req.route_params.get('cg-name')
try:
# tenant_id = os.environ["AZURE_TENANT_ID"]
# client_id = os.environ["AZURE_CLIENT_ID"]
# client_secret = os.environ["AZURE_CLIENT_SECRET"]
subscription_id = os.environ["AZURE_SUBSCRIPTION_ID"]
resource_group = os.environ["AZURE_RESOURCE_GROUP"]
except KeyError as ex:
return func.HttpResponse(body=json.dumps(
{"message": "Error loading environment variable: " + str(ex)}),
headers={"Content-Type": "application/json"},
status_code=500)
credentials = CredentialWrapper()
aci_client = ContainerInstanceManagementClient(
subscription_id=subscription_id, credentials=credentials)
res_client = ResourceManagementClient(subscription_id=subscription_id,
credentials=credentials)
try:
cg = aci_client.container_groups.get(resource_group, cg_name)
except CloudError as ex:
# someone probably deleted container group during run
return func.HttpResponse(body=json.dumps({
"message": ex.message,
}),
headers={"Content-Type": "application/json"},
status_code=400)
if cg.instance_view.state not in ["Succeeded", "Stopped", "Failed"]:
return func.HttpResponse(
body=json.dumps({
"message": f"Container Group {cg_name} is in progress",
"name": cg_name,
"state": cg.instance_view.state
}),
headers={
"Content-Type":
"application/json",
"Location":
f"https://{os.environ['WEBSITE_HOSTNAME']}/api/status/{cg_name}"
},
status_code=202)
else:
containers = {}
for container in cg.containers:
logs = aci_client.container.list_logs(resource_group, cg_name,
container.name)
containers[container.name] = {}
containers[container.name][
"start_time"] = container.instance_view and container.instance_view.current_state.start_time and container.instance_view.current_state.start_time.strftime(
"%Y-%m-%dT%H:%M:%S")
containers[container.name][
"finish_time"] = container.instance_view and container.instance_view.current_state.finish_time and container.instance_view.current_state.finish_time.strftime(
"%Y-%m-%dT%H:%M:%S")
containers[container.name][
"state"] = container.instance_view and container.instance_view.current_state.state
containers[container.name][
"exit_code"] = container.instance_view and container.instance_view.current_state.exit_code
containers[container.name][
"detail_status"] = container.instance_view and container.instance_view.current_state.detail_status
containers[container.name]["output"] = logs.content
return func.HttpResponse(
body=json.dumps({
"message": f"Container Group {cg_name} finished",
"name": cg_name,
"state": cg.instance_view.state,
"containers": containers
}),
headers={"Content-Type": "application/json"},
status_code=200 if cg.instance_view.state == "Succeeded" else 500)
import random
import string
import sys
from collections import deque
from config.config import queueConf, azure_context, DATABASE_URI, ACI_CONFIG
from azure.servicebus import ServiceBusService, Message, Queue
from azure.mgmt.resource import ResourceManagementClient
from azure.mgmt.containerinstance import ContainerInstanceManagementClient
from azure.mgmt.containerinstance.models import (
ContainerGroup, Container, ContainerPort, Port, IpAddress,
EnvironmentVariable, ResourceRequirements, ResourceRequests,
ContainerGroupNetworkProtocol, OperatingSystemTypes)
resource_client = ResourceManagementClient(azure_context.credentials,
azure_context.subscription_id)
client = ContainerInstanceManagementClient(azure_context.credentials,
azure_context.subscription_id)
bus_service = ServiceBusService(
service_namespace=queueConf['service_namespace'],
shared_access_key_name=queueConf['saskey_name'],
shared_access_key_value=queueConf['saskey_value'])
BASE_NAMES = deque([
"anders", "wenjun", "robbie", "robin", "allen", "tony", "xiaofeng",
"tingting", "harry", "chen"
])
NAMES_COUNTER = 0
IMAGE = "pskreter/worker-container:latest"
def main():
def main():
SUBSCRIPTION_ID = os.environ.get("SUBSCRIPTION_ID", None)
GROUP_NAME = "testgroupx"
CONTAINER_GROUP = "container_groupxxyyzz"
CONTAINER_NAME = "my-container" # must match the regex '[a-z0-9]([-a-z0-9]*[a-z0-9])?'
# Create client
# # For other authentication approaches, please see: https://pypi.org/project/azure-identity/
resource_client = ResourceManagementClient(
credential=DefaultAzureCredential(), subscription_id=SUBSCRIPTION_ID)
containerinstance_client = ContainerInstanceManagementClient(
credential=DefaultAzureCredential(), subscription_id=SUBSCRIPTION_ID)
# Create resource group
resource_client.resource_groups.create_or_update(GROUP_NAME,
{"location": "eastus"})
# Create container group
container_group = containerinstance_client.container_groups.begin_create_or_update(
GROUP_NAME, CONTAINER_GROUP, {
"location":
"eastus",
"identity": {
"type": "SystemAssigned"
},
"containers": [{
"name":
CONTAINER_NAME,
"command": [],
"environment_variables": [],
"image":
"nginx",
"ports": [{
"port": "80"
}],
"resources": {
"requests": {
"cpu": "1",
"memory_in_gb": "1.5",
"gpu": {
"count": "1",
"sku": "K80"
}
}
},
"volume_mounts": [{
"name": "empty-volume",
"mount_path": "/mnt/mydir"
}]
}],
"diagnostics": {
"log_analytics": {
"workspace_id": "workspaceid",
"workspace_key": "workspaceKey"
}
},
"os_type":
"Linux",
"restart_policy":
"OnFailure",
"volumes": [{
"name": "empty-volume",
"empty_dir": {}
}]
}).result()
print("Create container group:\n{}".format(container_group))
# Get container group
container_group = containerinstance_client.container_groups.get(
GROUP_NAME, CONTAINER_GROUP)
print("Get container group:\n{}".format(container_group))
# Update container group
container_group = containerinstance_client.container_groups.update(
GROUP_NAME, CONTAINER_GROUP,
{"tags": {
"tag1key": "tag1Value",
"tag2key": "tag2Value"
}})
print("Update container group:\n{}".format(container_group))
# Container exec
result = containerinstance_client.containers.execute_command(
GROUP_NAME, CONTAINER_GROUP, CONTAINER_NAME, {
"command": "/bin/bash",
"terminal_size": {
"rows": "12",
"cols": "12"
}
})
print("Container exec:\n{}".format(result))
# Delete container group
container_group = containerinstance_client.container_groups.begin_delete(
GROUP_NAME, CONTAINER_GROUP).result()
print("Delete container group.\n")
# Delete Group
resource_client.resource_groups.begin_delete(GROUP_NAME).result()