Skip to content

Instantly share code, notes, and snippets.

@itaysk
Created October 7, 2016 19:49
Show Gist options
  • Save itaysk/fba9770f4f91729acc11a4c8f3a81bcd to your computer and use it in GitHub Desktop.
Save itaysk/fba9770f4f91729acc11a4c8f3a81bcd to your computer and use it in GitHub Desktop.
Azure Batch quick start
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment