Azure Batch quick start

azure_batch_quick_start.py

import os
import azure.batch.batch_service_client as batch
import azure.batch.batch_auth as batchauth
import azure.batch.models as batchmodels

#---------------------parameters---------------------#
#--------------this is the part you edit-------------#

batch_account_name = "<your batch account name here>"
batch_account_url = "<your batch account url here>"
batch_account_key = "<your batch key here>"

pool_id = "poo11"
pool_vm_size = "STANDARD_A1"
pool_vm_count = 2

job_id = "job1"
task_id = "task1"

pool_start_task_commands = [ 
]

#list of files to download for the start task
#The files must be publicly accessible! (or the required auth info should be contained in the url)
pool_start_task_resource_files_urls = [
]

job_preparation_task_commands = [
        'sudo apt install tree' #note that we later set the startup task to run elevated 
]

#list of files to download for the start task
#The files must be publicly accessible! (or the required auth info should be contained in the url)
job_preparation_task_resource_files_urls = [
]

task_commands = [
        'python task.py'
]

#list of files to download for the task
#The files must be publicly accessible! (or the required auth info should be contained in the url)
task_resource_files_urls = [
        "https://<your-storage-account>/task.py", #this is the actual piece of code that we want to run
        "https://<yout-storage-account>/1.txt" #this is just a data file that the task might need
]

#---------------------object model---------------------#
#------------------better to stay away-----------------#

def wrap_commands_in_shell(commands):
        return "/bin/bash -c 'set -e; set -o pipefail; {}; wait'".format(';'.join(commands))

pool_start_task_resource_files = [batchmodels.ResourceFile(file_path=os.path.basename(file_url), blob_source=file_url) for file_url in pool_start_task_resource_files_urls]

pool_creation_information = batch.models.PoolAddParameter(
	id=pool_id,
	virtual_machine_configuration=batchmodels.VirtualMachineConfiguration(
                image_reference=batchmodels.ImageReference(
        	        publisher="Canonical",
        	        offer="UbuntuServer",
        	        sku="14.04.4-LTS",
        	        version="latest"
                ),
                node_agent_sku_id="batch.node.ubuntu 14.04"
        ),
	vm_size=pool_vm_size,
	target_dedicated=pool_vm_count,
	start_task=batch.models.StartTask(
		command_line=wrap_commands_in_shell(pool_start_task_commands),
		run_elevated=True,
		wait_for_success=True,
                resource_files = pool_start_task_resource_files,
        )
)

job_preparation_task_resource_files = [batchmodels.ResourceFile(file_path=os.path.basename(file_url), blob_source=file_url) for file_url in job_preparation_task_resource_files_urls]

job_creation_information = batch.models.JobAddParameter(
        id=job_id, 
        pool_info=batch.models.PoolInformation(pool_id=pool_id), 
        job_preparation_task=batch.models.JobPreparationTask(
                command_line=wrap_commands_in_shell(job_preparation_task_commands),
                run_elevated=True,
                wait_for_success=True,
                resource_files=job_preparation_task_resource_files        
        )
)		

task_resource_files = [batchmodels.ResourceFile(file_path=os.path.basename(file_url), blob_source=file_url) for file_url in task_resource_files_urls]

#because we want each node to run a task, we need to create as many tasks as there are nodes
tasks_creation_information = list()
for i in range(pool_vm_count):
    tasks_creation_information.append(batch.models.TaskAddParameter(
                                                        "%s%s" %(task_id, i),
                                                        wrap_commands_in_shell(task_commands), 
                                                        resource_files = task_resource_files, 
                                                        run_elevated=False,
                                        )
                                    )

#---------------------api invoke---------------------#
#---------comment out parts on repeated runs---------#

credentials = batchauth.SharedKeyCredentials(batch_account_name, batch_account_key)
batch_client = batch.BatchServiceClient(credentials, base_url=batch_account_url)

batch_client.pool.add(pool_creation_information)
batch_client.job.add(job_creation_information)
batch_client.task.add_collection(job_id, tasks_creation_information)