Batch Pool Setup with Node Initialization

Create a Python script that configures an Azure Batch pool with automatic node initialization to prepare the compute environment before tasks execute.

Requirements

Your script should create a batch pool with the following specifications:

1. Pool Configuration:

   • Pool ID: "analysis-pool"
   • VM size: "STANDARD_D2_V2"
   • Target dedicated nodes: 2
   • Use Virtual Machine Configuration with Ubuntu 20.04 LTS

2. Node Initialization Setup:

   • Configure a start task that runs when nodes join the pool
   • The start task should execute a shell command: apt-get update && apt-get install -y python3-pip
   • The start task must complete successfully before the node accepts tasks (wait for success)
   • Run the start task with elevated (admin) privileges

3. Script Behavior:

   • Accept batch account URL, account name, and account key as command-line arguments
   • Create the pool with the initialization configuration
   • Print a success message when the pool is created
   • Handle authentication using shared key credentials

Input Format

The script should accept three command-line arguments:

1. Batch account URL (e.g., "https://mybatch.eastus.batch.azure.com")
2. Account name
3. Account key

Example usage:

python setup_pool.py https://mybatch.eastus.batch.azure.com myaccount mykey123

Output Format

Print status messages during execution:

Authenticating with Azure Batch...
Creating pool 'analysis-pool'...
Pool created successfully with node initialization configured.

Capabilities

Pool Creation with Start Task

• When called with valid credentials, it creates a pool with ID "analysis-pool" and a start task that executes the specified command @test
• The start task is configured to wait for successful completion before accepting tasks @test
• The start task runs with elevated (admin) privileges @test

@generates

API

import sys
from azure.batch import BatchServiceClient
from azure.batch.batch_auth import SharedKeyCredentials
from azure.batch import models

def create_pool_with_start_task(batch_url: str, account_name: str, account_key: str) -> None:
    """
    Creates an Azure Batch pool with a start task that initializes nodes.

    Args:
        batch_url: The URL of the Azure Batch account
        account_name: The name of the batch account
        account_key: The access key for the batch account
    """
    pass

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python setup_pool.py <batch_url> <account_name> <account_key>")
        sys.exit(1)

    batch_url = sys.argv[1]
    account_name = sys.argv[2]
    account_key = sys.argv[3]

    create_pool_with_start_task(batch_url, account_name, account_key)

Dependencies { .dependencies }

azure-batch { .dependency }

Provides Azure Batch client library for managing batch computing workloads.

@satisfied-by