Azure Batch – Create a Custom Pool using Shared Image Gallery

March 28, 2020Tech Notes 3 minsComments Off on Azure Batch – Create a Custom Pool using Shared Image Gallery

If you have a custom image in a Shared Image Gallery and you want to use it to create a pool in Azure Batch, this document, Use the Shared Image Gallery to create a custom pool, provides a pretty good guidance for it. I followed it to test the scenario and hit two minor issues.

  • As it is mentioned in another doc, AAD authentication is also a prerequisite for using shared image gallery. If you use --shared-key-auth with az batch account login, you would hit an anthentication error with Azure Cli. I raised an issue for the document and hopefully a note will be added to it.
  • There is no sample code to demonstrate how to create a pool with shared image gallery with Python.

So I wrote a simple sample in Python. It is based on the latest version (9.0.0) of Azure Batch package for Python. And it uses a service principal for the AAD authentication. The custom image I used for test was built on top of Ubuntu 18.04-LTS. So the node agent sku is ubuntu 18.04. It needs to be changed accordingly if other os version is used.

# Import the required modules from the
# Azure Batch Client Library for Python
import azure.batch._batch_service_client as batch
import azure.batch.models as batchmodels
from azure.common.credentials import ServicePrincipalCredentials

# Specify Batch account credentials
account = "<batch-account-name>"
batch_url = "<batch-account-url>"
ad_client_id = "<client id of the SP>"
ad_tenant = "<tenant id>"
ad_secret = "<secret of the SP>"

# Pool settings
pool_id = "LinuxNodesSamplePoolPython"
vm_size = "STANDARD_D2_V3"
node_count = 1

# Initialize the Batch client
creds = ServicePrincipalCredentials(
    client_id=ad_client_id,
    secret=ad_secret,
    tenant=ad_tenant,
    resource="https://batch.core.windows.net/"
)
config = batch.BatchServiceClientConfiguration(creds, batch_url)
client = batch.BatchServiceClient(creds, batch_url)

# Create the unbound pool
new_pool = batchmodels.PoolAddParameter(id=pool_id, vm_size=vm_size)
new_pool.target_dedicated = node_count

# Configure the start task for the pool
start_task = batchmodels.StartTask(
    command_line="printenv AZ_BATCH_NODE_STARTUP_DIR"
)
start_task.run_elevated = True
new_pool.start_task = start_task

# Create an ImageReference which specifies the Marketplace
# virtual machine image to install on the nodes.
ir = batchmodels.ImageReference(
    virtual_machine_image_id="<resource id of the image version in sig>"
)

# Create the VirtualMachineConfiguration, specifying
# the VM image reference and the Batch node agent to
# be installed on the node.
vmc = batchmodels.VirtualMachineConfiguration(
    image_reference=ir,
    node_agent_sku_id="batch.node.ubuntu 18.04"
)

# Assign the virtual machine configuration to the pool
new_pool.virtual_machine_configuration = vmc

# Create pool in the Batch service
client.pool.add(new_pool)

Update: I polished the above sample code and pushed it into the document I mentioned at the beginning of this post via a PR. The Python sample code in that document is based on the one in this post.