def _get_azure_batch_client(conf):
from azure.batch import batch_auth, BatchServiceClient
creds = batch_auth.SharedKeyCredentials(
conf[utils.PLATFORM]['batch_account'],
conf[utils.PLATFORM]['batch_key'])
batch_url = f"https://{conf[utils.PLATFORM]['batch_account']}.{conf[utils.PLATFORM]['location']}.batch.azure.com"
return BatchServiceClient(creds, batch_url=batch_url)
def create_batch_client():
credentials = batch_auth.SharedKeyCredentials(config["batch_account_name"],
config["batch_account_key"])
batch_client = batch.BatchServiceClient(
credentials, batch_url=config["batch_account_url"])
return batch_client
def get_conn(self):
"""
Get the batch client connection
:return: Azure batch client
"""
conn = self._connection()
def _get_required_param(name):
"""Extract required parameter from extra JSON, raise exception if not found"""
value = conn.extra_dejson.get(name)
if not value:
raise AirflowException(
'Extra connection option is missing required parameter: `{}`'
.format(name))
return value
batch_account_name = _get_required_param('account_name')
batch_account_key = _get_required_param('account_key')
batch_account_url = _get_required_param('account_url')
credentials = batch_auth.SharedKeyCredentials(batch_account_name,
batch_account_key)
batch_client = BatchServiceClient(credentials,
batch_url=batch_account_url)
return batch_client
def main(req: func.HttpRequest) -> func.HttpResponse:
logging.info('Python HTTP trigger function processed a request.')
start_time = datetime.datetime.now().replace(microsecond=0)
print('Sample start: {}'.format(start_time))
print()
# Create a Batch service client. We'll now be interacting with the Batch
# service in addition to Storage
credentials = batch_auth.SharedKeyCredentials(config._BATCH_ACCOUNT_NAME,
config._BATCH_ACCOUNT_KEY)
batch_client = batch.BatchServiceClient(
credentials, batch_url=config._BATCH_ACCOUNT_URL)
name = req.params.get('name')
if not name:
try:
req_body = req.get_json()
except ValueError:
pass
else:
name = req_body.get('name')
if name:
task_id = add_tasks(batch_client, config._JOB_ID, name)
return func.HttpResponse(f"Task created {task_id}!")
else:
return func.HttpResponse(
"Please pass a name on the query string or in the request body",
status_code=400)
def _get_batch_client(self, config_batch_credentials):
batch_credentials = batch_auth.SharedKeyCredentials(
config_batch_credentials['account'],
config_batch_credentials['account_key'])
return batch_service_client.BatchServiceClient(
batch_credentials,
base_url=config_batch_credentials['account_service_url'])
def get_conn(self):
"""
Get the batch client connection
:return: Azure batch client
"""
conn = self._connection()
def _get_required_param(name):
"""Extract required parameter from extra JSON, raise exception if not found"""
value = conn.extra_dejson.get(name)
if not value:
raise AirflowException(
f'Extra connection option is missing required parameter: `{name}`'
)
return value
batch_account_url = _get_required_param(
'account_url') or _get_required_param(
'extra__azure_batch__account_url')
credentials = batch_auth.SharedKeyCredentials(conn.login,
conn.password)
batch_client = BatchServiceClient(credentials,
batch_url=batch_account_url)
return batch_client
def batch_data_service_factory(cli_ctx, kwargs):
import azure.batch.batch_service_client as batch
import azure.batch.batch_auth as batchauth
account_name = kwargs.pop('account_name', None)
account_key = kwargs.pop('account_key', None)
account_endpoint = kwargs.pop('account_endpoint', None)
kwargs.pop('yes', None)
credentials = None
if not account_key:
from azure.cli.core._profile import Profile
profile = Profile(cli_ctx=cli_ctx)
# in order to use AAD auth in cloud shell mode, we will use mgmt AAD token
# instead of Batch AAD token to auth
if in_cloud_console():
resource = cli_ctx.cloud.endpoints.active_directory_resource_id
else:
resource = cli_ctx.cloud.endpoints.batch_resource_id
credentials, _, _ = profile.get_login_credentials(resource=resource)
else:
credentials = batchauth.SharedKeyCredentials(account_name, account_key)
if not account_endpoint.startswith('https://'):
account_endpoint = 'https://' + account_endpoint
return batch.BatchServiceClient(credentials, base_url=account_endpoint)
def __init__(self, test_method):
super(TestBatchNCJLive, self).__init__(__file__, test_method)
self.account_name = 'test1'
if not self.playback:
self.account_key = os.environ['AZURE_BATCH_ACCESS_KEY']
else:
self.account_key = 'ZmFrZV9hY29jdW50X2tleQ=='
self.account_endpoint = 'https://test1.westus.batch.azure.com/'
storage_account = 'testaccountforbatch'
if not self.playback:
storage_key = os.environ['AZURE_STORAGE_ACCESS_KEY']
else:
storage_key = '1234'
self.blob_client = CloudStorageAccount(storage_account, storage_key)\
.create_block_blob_service()
credentials = batchauth.SharedKeyCredentials(self.account_name,
self.account_key)
self.batch_client = batch.BatchServiceClient(
credentials, base_url=self.account_endpoint)
self.output_blob_container = 'aaatestcontainer'
sas_token = self.blob_client.generate_container_shared_access_signature(
self.output_blob_container,
permission=BlobPermissions(read=True, write=True),
start=datetime.datetime.utcnow(),
expiry=datetime.datetime.utcnow() + datetime.timedelta(days=1))
self.output_container_sas = 'https://{}.blob.core.windows.net/{}?{}'.format(
storage_account, self.output_blob_container, sas_token)
print('Full container sas: {}'.format(self.output_container_sas))
def batch_data_service_factory(cli_ctx, kwargs):
import azure.batch.batch_service_client as batch
import azure.batch.batch_auth as batchauth
account_name = kwargs.pop('account_name', None)
account_key = kwargs.pop('account_key', None)
account_endpoint = kwargs.pop('account_endpoint', None)
kwargs.pop('yes', None)
credentials = None
if not account_key:
from azure.cli.core._profile import Profile
profile = Profile(cli_ctx=cli_ctx)
# in order to use AAD auth in cloud shell mode, we will use mgmt AAD token
# instead of Batch AAD token to auth
if in_cloud_console():
resource = cli_ctx.cloud.endpoints.active_directory_resource_id
else:
resource = cli_ctx.cloud.endpoints.batch_resource_id
credentials, _, _ = profile.get_login_credentials(resource=resource)
else:
# Verify all values are populated and display readable error
if not all([account_name, account_key, account_endpoint]):
raise ValueError(
'usage error: --account-name NAME --account-key KEY --account-endpoint ENDPOINT'
)
credentials = batchauth.SharedKeyCredentials(account_name, account_key)
if not (account_endpoint.startswith('https://')
or account_endpoint.startswith('http://')):
account_endpoint = 'https://' + account_endpoint
return batch.BatchServiceClient(credentials,
batch_url=account_endpoint.rstrip('/'))
def execute_script(global_config, script_config):
"""Executes the sample with the specified configurations.
:param global_config: The global configuration to use.
:type global_config: `configparser.ConfigParser`
:param script_config: The script specific configuration to use.
:type script_config: `configparser.ConfigParser`
"""
# Set up the configuration
batch_account_key = global_config.get('Batch', 'batchaccountkey')
batch_account_name = global_config.get('Batch', 'batchaccountname')
batch_service_url = global_config.get('Batch', 'batchserviceurl')
#coming from the script specified config
pool_id = script_config.get('DEFAULT', 'pool_id')
username = script_config.get('DEFAULT', 'username')
password = script_config.get( #getpass.getpass()
'DEFAULT', 'password')
expirationTime = script_config.getint('DEFAULT', 'expirationTime')
# Print the settings we are running with
common.helpers.print_configuration(global_config)
common.helpers.print_configuration(script_config)
credentials = batchauth.SharedKeyCredentials(batch_account_name,
batch_account_key)
batch_client = batch.BatchServiceClient(credentials,
base_url=batch_service_url)
try:
connectToNodes(batch_client, pool_id, username, password,
expirationTime)
except Exception as e:
print('error/exception: ' + e)
def execute_sample(global_config, sample_config):
"""Executes the sample with the specified configurations.
:param global_config: The global configuration to use.
:type global_config: `configparser.ConfigParser`
:param sample_config: The sample specific configuration to use.
:type sample_config: `configparser.ConfigParser`
"""
# Set up the configuration
batch_account_key = global_config.get('Batch', 'batchaccountkey')
batch_account_name = global_config.get('Batch', 'batchaccountname')
batch_service_url = global_config.get('Batch', 'batchserviceurl')
should_delete_job = sample_config.getboolean(
'DEFAULT',
'shoulddeletejob')
pool_vm_size = sample_config.get(
'DEFAULT',
'poolvmsize')
pool_vm_count = sample_config.getint(
'DEFAULT',
'poolvmcount')
# Print the settings we are running with
common.helpers.print_configuration(global_config)
common.helpers.print_configuration(sample_config)
credentials = batchauth.SharedKeyCredentials(
batch_account_name,
batch_account_key)
batch_client = batch.BatchServiceClient(
credentials,
base_url=batch_service_url)
# Retry 5 times -- default is 3
batch_client.config.retry_policy.retries = 5
job_id = common.helpers.generate_unique_resource_name("HelloWorld")
try:
submit_job_and_add_task(
batch_client,
job_id,
pool_vm_size,
pool_vm_count)
common.helpers.wait_for_tasks_to_complete(
batch_client,
job_id,
datetime.timedelta(minutes=25))
tasks = batch_client.task.list(job_id)
task_ids = [task.id for task in tasks]
common.helpers.print_task_output(batch_client, job_id, task_ids)
finally:
if should_delete_job:
print("Deleting job: ", job_id)
batch_client.job.delete(job_id)
def __init__(self):
"""Return a new AzureBatchEngine object."""
self.credentials = azbatch_auth.SharedKeyCredentials(
current_app.config['BATCH_ACCOUNT_NAME'],
current_app.config['BATCH_ACCOUNT_KEY'])
self.batch_client = azbatch.batch_service_client.BatchServiceClient(
self.credentials, current_app.config['BATCH_ACCOUNT_URL'])
def execute_sample(global_config, sample_config):
"""Executes the sample with the specified configurations.
:param global_config: The global configuration to use.
:type global_config: `configparser.ConfigParser`
:param sample_config: The sample specific configuration to use.
:type sample_config: `configparser.ConfigParser`
"""
credentials = batchauth.SharedKeyCredentials(batch_account_name,
batch_account_key)
#credentials = ServicePrincipalCredentials(
# client_id=aad_client_id,
# secret=aad_client_secret,
# tenant=aad_tenant_id,
# resource="https://batch.core.windows.net/"
#)
batch_client = batch.BatchServiceClient(credentials,
batch_url=batch_service_url)
# Retry 5 times -- default is 3
batch_client.config.retry_policy.retries = 5
block_blob_client = azureblob.BlockBlobService(
account_name=storage_account_name,
account_key=storage_account_key,
endpoint_suffix=storage_account_suffix)
job_id = common.helpers.generate_unique_resource_name(
"poolsandresourcefilesjob")
try:
create_pool(batch_client, block_blob_client, pool_id, pool_vm_size,
pool_vm_count)
submit_job_and_add_task(batch_client, block_blob_client, job_id,
pool_id, storage_account_name)
common.helpers.wait_for_tasks_to_complete(
batch_client, job_id, datetime.timedelta(minutes=25))
tasks = batch_client.task.list(job_id)
task_ids = [task.id for task in tasks]
common.helpers.print_task_output(batch_client, job_id, task_ids)
print("Completed job_id:" + job_id)
finally:
# clean up
if should_delete_container:
block_blob_client.delete_container(_RESOURCE_CONTAINER_NAME,
fail_not_exist=False)
if should_delete_job:
print("Deleting job: ", job_id)
batch_client.job.delete(job_id)
if should_delete_pool:
print("Deleting pool: ", pool_id)
batch_client.pool.delete(pool_id)
def create_batch_service_client():
"""AzureBatchの認証を行いBatchServiceClientを生成する."""
# AzureBatchに対する認証を行う.
credentials = batch_auth.SharedKeyCredentials(cfg.BATCH_ACCOUNT_NAME,
cfg.BATCH_ACCOUNT_KEY)
# batch_service_clientを生成する.
return batch.BatchServiceClient(credentials,
batch_url=cfg.BATCH_ACCOUNT_URL)
def get_batch_client():
credentials = batch_auth.SharedKeyCredentials(
BATCH_ACCOUNT_NAME,
BATCH_ACCOUNT_KEY)
batch_service_client = azure.batch.BatchServiceClient(
credentials,
batch_url=f"https://{BATCH_ACCOUNT_NAME}.{BATCH_LOCATION}.batch.azure.com")
return batch_service_client
def createBatchClient():
credentials = batch_auth.SharedKeyCredentials(
os.environ.get('AZURE_BATCH_ACCOUNT_NAME'),
os.environ.get('AZURE_BATCH_ACCOUNT_KEY'))
batch_client = batch.BatchServiceClient(
credentials, batch_url=os.environ.get('AZURE_BATCH_ACCOUNT_URL'))
return batch_client
def batch_data_service_factory(kwargs):
import azure.batch.batch_service_client as batch
import azure.batch.batch_auth as batchauth
account_name = kwargs.pop('account_name', None)
account_key = kwargs.pop('account_key', None)
account_endpoint = kwargs.pop('account_endpoint', None)
kwargs.pop('yes', None)
credentials = batchauth.SharedKeyCredentials(account_name, account_key)
return batch.BatchServiceClient(credentials, base_url=account_endpoint)
def __init__(self, batch_storage_account):
# Create a Batch service client. We'll now be interacting with the Batch
# service in addition to Storage
self.my_storage = batch_storage_account
configuration = AzureCredentials()
self.account_name = configuration.getBatchAccountName()
self.account_key = configuration.getBatchAccountKey()
self.account_url = configuration.getBatchAccountUrl()
self.credentials = batchauth.SharedKeyCredentials(self.account_name,self.account_key)
self.batch_client = batch.BatchServiceClient(self.credentials,batch_url=self.account_url)
print("API version is: ", self.batch_client.task.api_version)
#self.batch_client.task.api_version = "2020-03-01.11.0"
#print("API version is: ", self.batch_client.task.api_version)
batch_config = AzureBatchConfiguration()
self.pool_count = batch_config.getNodeCount()
self.pool_type = batch_config.getVMSize()
self.pool_os = batch_config.getOSType()
self.pool_publisher = batch_config.getOSPublisher()
self.pool_os_ver = batch_config.getOSVersion()
self.pool_engine_name = batch_config.getEngineName()
batch_json = find_file_path("batch.json", "../")
print("Found batch.json in: {}".format(batch_json))
credential_json = find_file_path("credentials.json", "../")
print("Found credentials.json in: {}".format(credential_json))
task_json = find_file_path("task.json", "../")
print("Found task.json in: {}".format(task_json))
self.my_storage.addApplicationFilePath("engine/"+batch_config.getEngineName())
self.my_storage.addApplicationFilePath("engine/taskengine.py")
#self.my_storage.addApplicationFilePath("batchwrapper/batch.json")
self.my_storage.addApplicationFilePath(batch_json)
self.my_storage.addApplicationFilePath("batchwrapper/__init__.py")
#self.my_storage.addApplicationFilePath("batchwrapper/credentials.json")
self.my_storage.addApplicationFilePath(credential_json)
self.my_storage.addApplicationFilePath(task_json)
self.my_storage.addApplicationFilePath("batchwrapper/config.py")
self.my_storage.uploadApplicationFiles()
def _create_credentials():
# type: (None) -> azure.batch.BatchServiceClient
"""Create authenticated client
:rtype: `azure.batch.BatchServiceClient`
:return: batch_client
"""
ba, url, bakey = os.environ['SHIPYARD_BATCH_ENV'].split(';')
batch_client = azure.batch.BatchServiceClient(
batchauth.SharedKeyCredentials(ba, bakey), batch_url=url)
batch_client.config.add_user_agent('batch-shipyard/tfm')
return batch_client
def check_prokka_tasks():
# Prokka!
prokka_tasks = ProkkaAzureRequest.objects.filter()
credentials = batch_auth.SharedKeyCredentials(settings.BATCH_ACCOUNT_NAME, settings.BATCH_ACCOUNT_KEY)
batch_client = batch.BatchServiceClient(credentials, base_url=settings.BATCH_ACCOUNT_URL)
for prokka_task in prokka_tasks:
prokka_object = ProkkaRequest.objects.get(pk=prokka_task.prokka_request.pk)
batch_job_name = 'prokka-{}'.format(prokka_task.prokka_request.pk)
# Check if tasks related with this amrsummary job have finished.
tasks_completed = True
try:
for cloudtask in batch_client.task.list(batch_job_name):
if cloudtask.state != batchmodels.TaskState.completed:
tasks_completed = False
except: # If something errors first time through job can't get deleted. In that case, give up.
ProkkaRequest.objects.filter(pk=prokka_task.prokka_request.pk).update(status='Error')
# Delete task so we don't keep iterating over it.
ProkkaAzureRequest.objects.filter(id=prokka_task.id).delete()
continue
# If tasks have completed, check if they were successful.
if tasks_completed:
exit_codes_good = True
for cloudtask in batch_client.task.list(batch_job_name):
if cloudtask.execution_info.exit_code != 0:
exit_codes_good = False
# Get rid of job and pool so we don't waste big $$$ and do cleanup/get files downloaded in tasks.
batch_client.job.delete(job_id=batch_job_name)
batch_client.pool.delete(pool_id=batch_job_name)
if exit_codes_good:
# Now need to generate an SAS URL and give access to it/update the download link.
blob_client = BlockBlobService(account_key=settings.AZURE_ACCOUNT_KEY,
account_name=settings.AZURE_ACCOUNT_NAME)
# Download the output container so we can zip it.
download_container(blob_service=blob_client,
container_name=batch_job_name + '-output',
output_dir='olc_webportalv2/media')
output_dir = 'olc_webportalv2/media/{}'.format(batch_job_name)
if os.path.isfile(os.path.join(output_dir, 'batch_config.txt')):
os.remove(os.path.join(output_dir, 'batch_config.txt'))
shutil.make_archive(output_dir, 'zip', output_dir)
prokka_result_container = 'prokka-result-{}'.format(prokka_object.pk)
sas_url = generate_download_link(blob_client=blob_client,
container_name=prokka_result_container,
output_zipfile=output_dir + '.zip',
expiry=8)
prokka_object.download_link = sas_url
prokka_object.status = 'Complete'
prokka_object.save()
shutil.rmtree(output_dir)
os.remove(output_dir + '.zip')
else:
prokka_object.status = 'Error'
prokka_object.save()
ProkkaAzureRequest.objects.filter(id=prokka_task.id).delete()
def execute_script(global_config, script_config):
"""Executes the sample with the specified configurations.
:param global_config: The global configuration to use.
:type global_config: `configparser.ConfigParser`
:param script_config: The script specific configuration to use.
:type script_config: `configparser.ConfigParser`
"""
# Set up the configuration
batch_account_key = global_config.get('Batch', 'batchaccountkey')
batch_account_name = global_config.get('Batch', 'batchaccountname')
batch_service_url = global_config.get('Batch', 'batchserviceurl')
storage_account_key = global_config.get('Storage', 'storageaccountkey')
storage_account_name = global_config.get('Storage', 'storageaccountname')
#coming from the script specified config
pool_id = script_config.get('DEFAULT', 'pool_id')
pool_node_count = script_config.getint('DEFAULT', 'pool_node_count')
vm_size = script_config.get('DEFAULT', 'vm_size')
distro = script_config.get('DEFAULT', 'distribution')
version = script_config.get('DEFAULT', 'version')
# Print the settings we are running with
common.helpers.print_configuration(global_config)
common.helpers.print_configuration(script_config)
credentials = batchauth.SharedKeyCredentials(batch_account_name,
batch_account_key)
batch_client = batch.BatchServiceClient(credentials,
base_url=batch_service_url)
blob_client = azureblob.BlockBlobService(account_name=storage_account_name,
account_key=storage_account_key)
#uploading script to Pool/Node/Blob Container
bcplatform_container_name = 'bcp'
application_file_paths = [os.path.realpath('bcpscript.sh')]
blob_client.create_container(bcplatform_container_name,
fail_on_exist=False)
resource_files = [
common.helpers.upload_file_to_container(blob_client,
bcplatform_container_name,
file_path)
for file_path in application_file_paths
]
try:
changePoolSettings(pool_id, vm_size, distro, version, pool_node_count,
resource_files, batch_client)
except Exception as e:
print('error/exception: ' + e)
def run(config):
batch_account_key = config.get('Batch', 'batchaccountkey')
batch_account_name = config.get('Batch', 'batchaccountname')
batch_service_url = config.get('Batch', 'batchserviceurl')
storage_account_key = config.get('Storage', 'storageaccountkey')
storage_account_name = config.get('Storage', 'storageaccountname')
storage_account_suffix = config.get('Storage', 'storageaccountsuffix')
delete_container = config.getboolean('Slicer', 'deletecontainer')
delete_job = config.getboolean('Slicer', 'deletejob')
delete_pool = config.getboolean('Slicer', 'deletepool')
pool_vm_size = config.get('Slicer', 'poolvmsize')
pool_vm_count = config.getint('Slicer', 'poolvmcount')
credentials = batchauth.SharedKeyCredentials(batch_account_name,
batch_account_key)
batch_client = batch.BatchServiceClient(credentials,
base_url=batch_service_url)
block_blob_client = azureblob.BlockBlobService(
account_name=storage_account_name,
account_key=storage_account_key,
endpoint_suffix=storage_account_suffix)
pool_id = "SlicerPool"
job_id = generate_unique_resource_name("SliceJob")
try:
create_pool(batch_client, block_blob_client, pool_id, pool_vm_size,
pool_vm_count)
submit_job_and_add_task(batch_client, block_blob_client, job_id,
pool_id)
wait_for_tasks_to_complete(batch_client, job_id,
datetime.timedelta(minutes=25))
tasks = batch_client.task.list(job_id)
task_ids = [task.id for task in tasks]
print_task_output(batch_client, job_id, task_ids)
finally:
if delete_container:
block_blob_client.delete_container(CONTAINER_NAME,
fail_not_exist=False)
if delete_job:
print("Deleting job: ", job_id)
batch_client.job.delete(job_id)
if delete_pool:
print("Deleting pool: ", pool_id)
batch_client.pool.delete(pool_id)
def batch_data_service_factory(kwargs):
from azure.cli.command_modules.batch_extensions.version import VERSION
account_name = kwargs['account_name']
account_key = kwargs.pop('account_key', None)
account_endpoint = kwargs['account_endpoint']
credentials = None
if not account_key:
from azure.cli.core._profile import Profile, CLOUD
profile = Profile()
credentials, _, _ = profile.get_login_credentials(
resource=CLOUD.endpoints.batch_resource_id)
else:
credentials = batchauth.SharedKeyCredentials(account_name, account_key)
client = batch.BatchServiceClient(credentials, base_url=account_endpoint)
client.config.add_user_agent('batch-extensions/v{}'.format(VERSION))
return client
def make_batch_client(secrets):
"""
Creates a batch client object
:param str batch_account_key: batch account key
:param str batch_account_name: batch account name
:param str batch_service_url: batch service url
"""
# Validate the given config
credentials = None
if secrets.shared_key:
# Set up SharedKeyCredentials
base_url = secrets.shared_key.batch_service_url
credentials = batch_auth.SharedKeyCredentials(
secrets.shared_key.batch_account_name,
secrets.shared_key.batch_account_key)
else:
# Set up ServicePrincipalCredentials
arm_credentials = ServicePrincipalCredentials(
client_id=secrets.service_principal.client_id,
secret=secrets.service_principal.credential,
tenant=secrets.service_principal.tenant_id,
resource="https://management.core.windows.net/",
)
m = RESOURCE_ID_PATTERN.match(
secrets.service_principal.batch_account_resource_id)
arm_batch_client = BatchManagementClient(arm_credentials,
m.group("subscription"))
account = arm_batch_client.batch_account.get(m.group("resourcegroup"),
m.group("account"))
base_url = "https://{0}/".format(account.account_endpoint)
credentials = ServicePrincipalCredentials(
client_id=secrets.service_principal.client_id,
secret=secrets.service_principal.credential,
tenant=secrets.service_principal.tenant_id,
resource="https://batch.core.windows.net/",
)
# Set up Batch Client
batch_client = batch.BatchServiceClient(credentials, base_url=base_url)
# Set retry policy
batch_client.config.retry_policy.retries = 5
batch_client.config.add_user_agent("aztk/{}".format(__version__))
return batch_client
def newBatchJobSchedule(account, key, URL, job_id, pool_id):
creds = batchAuth.SharedKeyCredentials(account_name=account, key=key)
client_creds = clientAuth.BatchServiceClient(creds, URL)
stop_running = input('How many days do you have this job to run for?: ')
hours = input('How many hours would you like the recurrence interval to be for your schedule?: ')
time_to_run = datetime.datetime.utcnow() + datetime.timedelta(days=int(stop_running))
try:
pool = batchmodels.PoolInformation(pool_id=pool_id)
jobSpec = batchmodels.JobSpecification(pool_info=pool)
schedule = batchmodels.Schedule(do_not_run_after=time_to_run, recurrence_interval=datetime.timedelta(hours=int(hours)))
job = batchmodels.JobScheduleAddParameter(id=job_id, schedule=schedule, job_specification=jobSpec)
client_creds.job_schedule.add(cloud_job_schedule=job)
except Exception as e:
logging.error(msg=e)
def get_conn(self):
"""
Get the Batch client connection
:return: Azure Batch client
"""
conn = self._connection()
batch_account_url = conn.extra_dejson.get(
'extra__azure_batch__account_url')
if not batch_account_url:
raise AirflowException('Batch Account URL parameter is missing.')
credentials = batch_auth.SharedKeyCredentials(conn.login,
conn.password)
batch_client = BatchServiceClient(credentials,
batch_url=batch_account_url)
return batch_client
def _create_batch_service_client(ctx):
# type: (CliContext) -> azure.batch.batch_service_client.BatchServiceClient
"""Create batch service client
:param CliContext ctx: Cli Context
:rtype: azure.batch.batch_service_client.BatchServiceClient
:return: batch service client
"""
bc = settings.credentials_batch(ctx.config)
if util.is_none_or_empty(bc.account_key):
logger.debug('batch account key not specified, using aad auth')
batch_aad = settings.credentials_batch(ctx.config).aad
credentials = aad.create_aad_credentials(ctx, batch_aad)
else:
credentials = batchauth.SharedKeyCredentials(bc.account,
bc.account_key)
batch_client = batchsc.BatchServiceClient(credentials,
base_url=bc.account_service_url)
batch_client.config.add_user_agent('batch-shipyard/{}'.format(__version__))
return batch_client
def __init__(self, test_method):
super(TestBatchExtensionsLive, self).__init__(__file__, test_method)
if self.playback:
self.account_name = 'sdktest2'
self.account_endpoint = 'https://sdktest2.westcentralus.batch.azure.com'
self.account_key = 'abc=='
self.subscription_id = "677f962b-9abf-4423-a27b-0c2f4094dcec"
storage_account = 'sdkteststore2'
storage_key = 'abc=='
else:
self.account_name = os.environ.get('AZURE_BATCH_ACCOUNT', 'test1')
self.account_endpoint = os.environ.get(
'AZURE_BATCH_ENDPOINT',
'https://test1.westus.batch.azure.com/')
self.account_key = os.environ['AZURE_BATCH_ACCESS_KEY']
self.subscription_id = os.environ.get(
'AZURE_BATCH_SUBSCRIPTION_ID',
"677f962b-9abf-4423-a27b-0c2f4094dcec")
storage_account = os.environ.get('AZURE_STORAGE_ACCOUNT',
'testaccountforbatch')
storage_key = os.environ.get('AZURE_STORAGE_ACCESS_KEY', 'abc==')
self.data_dir = os.path.join(os.path.dirname(__file__), 'data')
self.blob_client = CloudStorageAccount(storage_account, storage_key)\
.create_block_blob_service()
credentials = batchauth.SharedKeyCredentials(self.account_name,
self.account_key)
self.batch_client = batch.BatchExtensionsClient(
credentials,
batch_url=self.account_endpoint,
subscription_id=self.subscription_id,
batch_account=self.account_name)
self.output_blob_container = 'aaatestcontainer'
sas_token = self.blob_client.generate_container_shared_access_signature(
container_name=self.output_blob_container,
permission=BlobPermissions(read=True, write=True),
start=datetime.datetime.utcnow(),
expiry=datetime.datetime.utcnow() + datetime.timedelta(days=1))
self.output_container_sas = 'https://{}.blob.core.windows.net/{}?{}'.format(
storage_account, self.output_blob_container, sas_token)
print('Full container sas: {}'.format(self.output_container_sas))
def _create_batch_service_client(ctx):
# type: (CliContext) -> azure.batch.batch_service_client.BatchServiceClient
"""Create batch service client
:param CliContext ctx: Cli Context
:rtype: azure.batch.batch_service_client.BatchServiceClient
:return: batch service client
"""
bc = settings.credentials_batch(ctx.config)
if util.is_none_or_empty(bc.account_key):
if settings.verbose(ctx.config):
logger.debug('using aad auth as batch account key not specified')
batch_aad = settings.credentials_batch(ctx.config).aad
credentials = aad.create_aad_credentials(ctx, batch_aad)
else:
credentials = batchauth.SharedKeyCredentials(bc.account,
bc.account_key)
batch_client = batchsc.BatchServiceClient(credentials,
batch_url=bc.account_service_url)
_modify_client_for_retry_and_user_agent(batch_client)
return batch_client
def batch_data_service_factory(kwargs):
import azure.batch.batch_service_client as batch
import azure.batch.batch_auth as batchauth
account_name = kwargs.pop('account_name', None)
account_key = kwargs.pop('account_key', None)
account_endpoint = kwargs.pop('account_endpoint', None)
kwargs.pop('yes', None)
credentials = None
if not account_key:
from azure.cli.core._profile import Profile, CLOUD
profile = Profile()
credentials, _, _ = profile.get_login_credentials(
resource=CLOUD.endpoints.batch_resource_id)
else:
credentials = batchauth.SharedKeyCredentials(account_name, account_key)
if not account_endpoint.startswith('https://'):
account_endpoint = 'https://' + account_endpoint
return batch.BatchServiceClient(credentials, base_url=account_endpoint)
