azure_batch_pool_job_tasks_quickstart_with_ubuntu_vms.py

1import azure.batch as batch
2import azure.batch.models as batchmodels
3from msrest.authentication import BasicTokenAuthentication
4from azure.common.credentials import ServicePrincipalCredentials
5
6# Replace with your Batch account credentials
7BATCH_ACCOUNT_NAME = 'your_batch_account_name'
8BATCH_ACCOUNT_KEY = 'your_batch_account_key'
9BATCH_ACCOUNT_URL = 'https://your_batch_account_region.batch.azure.com'
10
11# Pool and Job configuration
12POOL_ID = 'PythonQuickstartPool'
13JOB_ID = 'PythonQuickstartJob'
14NODE_COUNT = 2
15VM_SIZE = 'standard_a1_v2'
16
17# Create a Batch client
18credentials = batch.batch_auth.SharedKeyCredentials(BATCH_ACCOUNT_NAME, BATCH_ACCOUNT_KEY)
19batch_client = batch.BatchServiceClient(credentials, base_url=BATCH_ACCOUNT_URL)
20
21def create_pool(batch_client, pool_id):
22    print(f"Creating pool: {pool_id}")
23    # Configure the Ubuntu image for the VMs
24    sku_to_use = '20_04-lts-gen2'
25    publisher = 'canonical'
26    offer = '0001-com-ubuntu-server-focal'
27    
28    image_ref = batchmodels.ImageReference(
29        publisher=publisher,
30        offer=offer,
31        sku=sku_to_use,
32        version='latest'
33    )
34    
35    config = batchmodels.VirtualMachineConfiguration(
36        image_reference=image_ref,
37        node_agent_sku_id='batch.node.ubuntu 20.04'
38    )
39    
40    pool = batchmodels.PoolAddParameter(
41        id=pool_id,
42        vm_size=VM_SIZE,
43        virtual_machine_configuration=config,
44        target_dedicated_nodes=NODE_COUNT
45    )
46    
47    batch_client.pool.add(pool)
48
49def create_job(batch_client, job_id, pool_id):
50    print(f"Creating job: {job_id}")
51    job = batchmodels.JobAddParameter(
52        id=job_id,
53        pool_info=batchmodels.PoolInformation(pool_id=pool_id)
54    )
55    batch_client.job.add(job)
56
57def add_tasks(batch_client, job_id):
58    print(f"Adding tasks to job: {job_id}")
59    tasks = []
60    for i in range(5):
61        task_id = f'task_{i}'
62        command = f'/bin/bash -c "echo Hello from task {i}"'
63        tasks.append(batchmodels.TaskAddParameter(id=task_id, command_line=command))
64    
65    batch_client.task.add_collection(job_id, tasks)
66
67if __name__ == '__main__':
68    try:
69        create_pool(batch_client, POOL_ID)
70        create_job(batch_client, JOB_ID, POOL_ID)
71        add_tasks(batch_client, JOB_ID)
72        print("Success! Pool, Job, and Tasks have been submitted.")
73    except batchmodels.BatchErrorException as err:
74        print(f"An error occurred: {err.message}")