azure_batch_pool_job_task_parallel_workload_quickstart.py

1import azure.batch as batch
2import azure.batch.models as batchmodels
3from azure.common.credentials import ServicePrincipalCredentials
4from azure.identity import DefaultAzureCredential
5
6# Update the following with your Batch account details
7BATCH_ACCOUNT_NAME = 'YOUR_BATCH_ACCOUNT_NAME'
8BATCH_ACCOUNT_KEY = 'YOUR_BATCH_ACCOUNT_KEY'
9BATCH_ACCOUNT_URL = 'https://YOUR_BATCH_ACCOUNT_NAME.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
17def create_pool(batch_client, pool_id):
18    print(f"Creating pool: {pool_id}")
19    
20    # Configure the VM configuration
21    # Use a supported image reference (Ubuntu 20.04)
22    image_ref = batchmodels.ImageReference(
23        publisher='microsoft-azure-batch',
24        offer='ubuntu-server-container',
25        sku='20-04-lts',
26        version='latest'
27    )
28    
29    vm_config = batchmodels.VirtualMachineConfiguration(
30        image_reference=image_ref,
31        node_agent_sku_id='batch.node.ubuntu 20.04'
32    )
33    
34    # Create the pool
35    new_pool = batchmodels.PoolAddParameter(
36        id=pool_id,
37        vm_size=VM_SIZE,
38        virtual_machine_configuration=vm_config,
39        target_dedicated_nodes=NODE_COUNT
40    )
41    
42    batch_client.pool.add(new_pool)
43
44def create_job(batch_client, job_id, pool_id):
45    print(f"Creating job: {job_id}")
46    job = batchmodels.JobAddParameter(
47        id=job_id,
48        pool_info=batchmodels.PoolInformation(pool_id=pool_id)
49    )
50    batch_client.job.add(job)
51
52def add_tasks(batch_client, job_id):
53    print(f"Adding tasks to job: {job_id}")
54    tasks = []
55    for i in range(5):
56        task_id = f"Task_{i}"
57        command = f"/bin/bash -c 'echo Hello from Azure Batch task {i}'"
58        tasks.append(batchmodels.TaskAddParameter(id=task_id, command_line=command))
59    
60    batch_client.task.add_collection(job_id, tasks)
61
62def main():
63    # Authenticate using Shared Key
64    credentials = batch.batch_auth.SharedKeyCredentials(BATCH_ACCOUNT_NAME, BATCH_ACCOUNT_KEY)
65    batch_client = batch.BatchServiceClient(credentials, base_url=BATCH_ACCOUNT_URL)
66
67    # 1. Create Pool
68    try:
69        create_pool(batch_client, POOL_ID)
70    except batchmodels.BatchErrorException as e:
71        if e.error.code != 'PoolExists':
72            raise
73        print(f"Pool {POOL_ID} already exists.")
74
75    # 2. Create Job
76    try:
77        create_job(batch_client, JOB_ID, POOL_ID)
78    except batchmodels.BatchErrorException as e:
79        if e.error.code != 'JobExists':
80            raise
81        print(f"Job {JOB_ID} already exists.")
82
83    # 3. Add Tasks
84    add_tasks(batch_client, JOB_ID)
85    
86    print("Tasks submitted. You can monitor progress in the Azure Portal.")
87
88if __name__ == '__main__':
89    main()