# coding: utf-8 # ------------------------------------------------------------------------- # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. See License.txt in the project root for # license information. # -------------------------------------------------------------------------- # pylint: disable=line-too-long import time import requests import six from azure.mgmt.batchai import BatchAIManagementClient from azure.mgmt.batchai import models from azure.storage.file import FileService from azure_devtools.scenario_tests import AzureTestError from devtools_testutils import AzureMgmtTestCase, FakeStorageAccount, AzureMgmtPreparer, StorageAccountPreparer from devtools_testutils.resource_testcase import RESOURCE_GROUP_PARAM, ResourceGroupPreparer class Helpers(object): MINUTE = 60 # Node allocation should not take longer than 15 minutes. NODE_STARTUP_TIMEOUT_SEC = 15 * MINUTE # Auto scaling should not take longer than 30 minutes. AUTO_SCALE_TIMEOUT_SEC = 30 * MINUTE # User name and password for admin user configured on compute cluster and file servers. ADMIN_USER_NAME = 'demoUser' ADMIN_USER_PASSWORD = 'Dem0Pa$$w0rd' # Name of the Azure Files which will be create for each cluster. AZURE_FILES_NAME = 'share' # Name of the directory for mounting Azure Files on cluster. AZURE_FILES_MOUNTING_PATH = 'azfiles' # Job's output directory configuration. JOB_OUTPUT_DIRECTORY_PATH = '$AZ_BATCHAI_MOUNT_ROOT/{0}/'.format(AZURE_FILES_MOUNTING_PATH) JOB_OUTPUT_DIRECTORY_ID = 'OUTPUTS' # Environment variable used by jobs to access the output directory. JOB_OUTPUT_DIRECTORY_PATH_ENV = '$AZ_BATCHAI_OUTPUT_{0}'.format(JOB_OUTPUT_DIRECTORY_ID) # Polling interval for checking nodes allocation, jobs completion and so on. _POLL_INTERVAL_SEC = 20 # Fake storage account returned during running tests against recorded session. FAKE_STORAGE = FakeStorageAccount(name='psdk', id='fakeid') # ID of output directory containing job's standard output STANDARD_OUTPUT_DIRECTORY_ID = 'stdouterr' # Name of the folder created by BatchAI for storing job's standard output/error folder. STDOUTERR_FOLDER_NAME = 'stdouterr' # Name of the folder created by BatchAI for storing output directories. OUTPUT_DIRECTORIES_FOLDER_NAME = 'outputs' # Location to run tests. LOCATION = 'northeurope' # Regular expression to validate IP address (we don't need strict validation, just a smoke test enough). RE_ID_ADDRESS = '\d+(?:\.\d+){3}' # The workspace to create clusters. DEFAULT_WORKSPACE_NAME = 'workspace' # Experiment to create jobs. DEFAULT_EXPERIMENT_NAME = 'experiment' @staticmethod def sleep_before_next_poll(is_live): """Sleep for polling interval :param int is_live: True if running in live mode. :return: Slept time in sec. """ if is_live: time.sleep(Helpers._POLL_INTERVAL_SEC) return Helpers._POLL_INTERVAL_SEC @staticmethod def _create_file_share(storage_account, storage_account_key): """Creates Azure Files in the storage account to be mounted into a cluster :param str storage_account: name of the storage account. :param str storage_account_key: storage account key. """ if storage_account == Helpers.FAKE_STORAGE.name: return service = FileService(storage_account, storage_account_key) service.create_share(Helpers.AZURE_FILES_NAME) @staticmethod def create_file_server(client, location, resource_group, nfs_name, subnet_id=None): """Creates NFS :param BatchAIManagementClient client: client instance. :param str location: location. :param str resource_group: resource group name. :param str nfs_name: file server name. :param models.ResourceId subnet_id: id of the subnet. :return models.FileServer: created file server. """ client.workspaces.create(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, location).result() return client.file_servers.create( resource_group, Helpers.DEFAULT_WORKSPACE_NAME, nfs_name, models.FileServerCreateParameters( vm_size='STANDARD_D1', ssh_configuration=models.SshConfiguration( user_account_settings=models.UserAccountSettings( admin_user_name=Helpers.ADMIN_USER_NAME, admin_user_password=Helpers.ADMIN_USER_PASSWORD, ) ), data_disks=models.DataDisks( disk_size_in_gb=10, disk_count=2, storage_account_type='Standard_LRS'), subnet=subnet_id)).result() @staticmethod def create_cluster(client, location, resource_group, cluster_name, vm_size, target_nodes, storage_account, storage_account_key, file_servers=None, file_systems=None, subnet_id=None, setup_task_cmd=None, setup_task_env=None, setup_task_secrets=None): """Creates a cluster with given parameters and mounted Azure Files :param BatchAIManagementClient client: client instance. :param str location: location. :param str resource_group: resource group name. :param str cluster_name: name of the cluster. :param str vm_size: vm size. :param int target_nodes: number of nodes. :param str storage_account: name of the storage account. :param str storage_account_key: storage account key. :param list(models.FileServerReference) file_servers: file servers. :param list(models.UnmanagedFileServerReference) file_systems: file systems. :param str setup_task_cmd: start task cmd line. :param dict[str, str] setup_task_env: environment variables for start task. :param dict[str, str] setup_task_secrets: environment variables with secret values for start task, server doesn't return values for these environment variables in get cluster responses. :param str subnet_id: virtual network subnet id. :return models.Cluster: the created cluster """ Helpers._create_file_share(storage_account, storage_account_key) setup_task = None if setup_task_cmd: setup_task = models.SetupTask( command_line=setup_task_cmd, environment_variables=[models.EnvironmentVariable(name=k, value=v) for k, v in setup_task_env.items()], secrets=[models.EnvironmentVariableWithSecretValue(name=k, value=v) for k, v in setup_task_secrets.items()], std_out_err_path_prefix='$AZ_BATCHAI_MOUNT_ROOT/{0}'.format(Helpers.AZURE_FILES_MOUNTING_PATH)) client.workspaces.create(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, location).result() return client.clusters.create( resource_group, Helpers.DEFAULT_WORKSPACE_NAME, cluster_name, parameters=models.ClusterCreateParameters( vm_size=vm_size, scale_settings=models.ScaleSettings( manual=models.ManualScaleSettings(target_node_count=target_nodes)), node_setup=models.NodeSetup( mount_volumes=models.MountVolumes( azure_file_shares=[models.AzureFileShareReference( azure_file_url='https://{0}.file.core.windows.net/{1}'.format(storage_account, Helpers.AZURE_FILES_NAME), relative_mount_path=Helpers.AZURE_FILES_MOUNTING_PATH, account_name=storage_account, credentials=models.AzureStorageCredentialsInfo( account_key=storage_account_key ), )], file_servers=file_servers, unmanaged_file_systems=file_systems ), setup_task=setup_task ), subnet=subnet_id, user_account_settings=models.UserAccountSettings( admin_user_name=Helpers.ADMIN_USER_NAME, admin_user_password=Helpers.ADMIN_USER_PASSWORD ), vm_priority='lowpriority' )).result() @staticmethod def create_custom_job(client, resource_group, cluster_id, job_name, nodes, cmd, job_preparation_cmd=None, container=None): """Creates custom toolkit job :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str cluster_id: resource Id of the cluster. :param str job_name: job name. :param int nodes: number of nodes to execute the job. :param str cmd: command line to run. :param str or None job_preparation_cmd: Job preparation command line. :param models.ContainerSettings or None container: container settings to run the job. :return models.Job: the created job. """ job_preparation = None if job_preparation_cmd: job_preparation = models.JobPreparation(command_line=job_preparation_cmd) client.experiments.create(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, Helpers.DEFAULT_EXPERIMENT_NAME).result() return client.jobs.create( resource_group, Helpers.DEFAULT_WORKSPACE_NAME, Helpers.DEFAULT_EXPERIMENT_NAME, job_name, parameters=models.JobCreateParameters( cluster=models.ResourceId(id=cluster_id), node_count=nodes, std_out_err_path_prefix='$AZ_BATCHAI_MOUNT_ROOT/{0}'.format(Helpers.AZURE_FILES_MOUNTING_PATH), output_directories=[models.OutputDirectory( id=Helpers.JOB_OUTPUT_DIRECTORY_ID, path_prefix=Helpers.JOB_OUTPUT_DIRECTORY_PATH, path_suffix="files")], input_directories=[models.InputDirectory( id='INPUT', path='$AZ_BATCHAI_MOUNT_ROOT/{0}/input'.format(Helpers.AZURE_FILES_MOUNTING_PATH))], container_settings=container, job_preparation=job_preparation, custom_toolkit_settings=models.CustomToolkitSettings( command_line=cmd ) ) ).result() @staticmethod def wait_for_nodes(is_live, client, resource_group, cluster_name, target, timeout_sec): """Wait for target number of nodes in a cluster become idle. :param bool is_live: True if running in live mode. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str cluster_name: cluster name. :param int target: target number of node. :param int timeout_sec: Timeout in seconds. :return int: Number of idle nodes. """ wait_time = 0 while True: cluster = client.clusters.get(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, cluster_name) counts = cluster.node_state_counts # type: models.NodeStateCounts if counts.idle_node_count == target and cluster.allocation_state != models.AllocationState.resizing: return counts.idle_node_count if wait_time < timeout_sec: wait_time += Helpers.sleep_before_next_poll(is_live) else: print("Cluster:") print(cluster.serialize()) return counts.idle_node_count @staticmethod def wait_for_job_completion(is_live, client, resource_group, job_name, timeout_sec): """Wait for job completion. :param bool is_live: True if running in live mode. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str job_name: job name. :param int timeout_sec: Timeout in seconds. :return models.ExecutionState: The job's execution state. """ wait_time = 0 while True: job = client.jobs.get(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, Helpers.DEFAULT_EXPERIMENT_NAME, job_name) if job.execution_state in [models.ExecutionState.succeeded, models.ExecutionState.failed]: return job.execution_state if wait_time < timeout_sec: wait_time += Helpers.sleep_before_next_poll(is_live) else: print("Job:") print(job.serialize()) return job.execution_state @staticmethod def wait_for_job_start_running(is_live, client, resource_group, job_name, timeout_sec): """Wait for job start running. :param bool is_live: True if running in live mode. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str job_name: job name. :param int timeout_sec: Timeout in seconds. :return models.ExecutionState: The job's execution state. """ wait_time = 0 while True: job = client.jobs.get(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, Helpers.DEFAULT_EXPERIMENT_NAME, job_name) if job.execution_state != models.ExecutionState.queued: return job.execution_state if wait_time < timeout_sec: wait_time += Helpers.sleep_before_next_poll(is_live) else: print("Job:") print(job.serialize()) return job.execution_state @staticmethod def print_file_server(file_server): """Output information about file server :param models.FileServer file_server: file server. """ print("File server:") print(file_server.serialize()) @staticmethod def wait_for_file_server(is_live, client, resource_group, file_server_name, timeout_sec): """Wait for file server to become available. :param bool is_live: True if running in live mode. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str file_server_name: cluster name. :param int timeout_sec: Timeout in seconds. :return bool: True is the file server become available. """ wait_time = 0 while True: server = client.file_servers.get(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, file_server_name) # type: models.FileServer if server.provisioning_state == models.FileServerProvisioningState.succeeded: return True if wait_time < timeout_sec: wait_time += Helpers.sleep_before_next_poll(is_live) else: Helpers.print_file_server(server) return False @staticmethod def assert_remote_login_info_reported_for_nodes(test, client, resource_group, cluster_name, expected): """Checks that given number of nodes are reported for the cluster :param AzureMgmtTestCase test: test instance. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str cluster_name: cluster :param int expected: expected number of nodes. """ nodes = list(client.clusters.list_remote_login_information(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, cluster_name)) test.assertEqual(len(nodes), expected) # Check if there is a reasonable information about nodes. for n in nodes: # type: models.RemoteLoginInformation test.assertIsNotNone(n.ip_address) test.assertIsNotNone(n.node_id) test.assertRegexpMatches(n.ip_address, Helpers.RE_ID_ADDRESS) test.assertGreater(len(n.node_id), 0) test.assertGreater(n.port, 0) @staticmethod def get_node_ids(client, resource_group, cluster_name): """Checks that given number of nodes are reported for the cluster :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str cluster_name: cluster :return list(str): list of node Ids """ return [n.node_id for n in list(client.clusters.list_remote_login_information(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, cluster_name))] @staticmethod def assert_job_files_in_path_are(test, client, resource_group, job_name, output_directory_id, path, expected): """Checks that the given task has expected output :param AzureMgmtTestCase test: test instance. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str job_name: job name. :param str output_directory_id: output directory id. :param str path: a path inside of output directory. :param dict(str, str or None) expected: expected content, directories must have None value. """ files = client.jobs.list_output_files(resource_group, Helpers.DEFAULT_WORKSPACE_NAME, Helpers.DEFAULT_EXPERIMENT_NAME, job_name, models.JobsListOutputFilesOptions(outputdirectoryid=output_directory_id, directory=path)) actual = dict() execution_log_found = False for f in files: if (output_directory_id == Helpers.STANDARD_OUTPUT_DIRECTORY_ID and f.name.startswith('execution') and f.name.endswith('.log')): execution_log_found = True continue actual[f.name] = None if f.file_type == models.FileType.file: v = requests.get(f.download_url).content actual[f.name] = v if isinstance(v, six.string_types) else v.decode() test.assertEquals(sorted(actual.keys()), sorted(expected.keys())) for k, v in expected.items(): a = actual[k] if a is None and v is None: # both are directories continue if v is None: test.fail('Expected {0} to be a directory, got a file'.format(k)) if a is None: test.fail('Expected {0} to be a file, got a directory'.format(k)) if isinstance(v, six.string_types): test.assertEquals(v, a, k + "expected {0} got {1}".format(v, a)) else: test.assertRegexpMatches(actual.get(k), v, k) if output_directory_id == Helpers.STANDARD_OUTPUT_DIRECTORY_ID and not execution_log_found: test.fail("No execution log was generated for the job.") @staticmethod def assert_job_files_are(test, client, resource_group, job_name, output_directory_id, expected): """Checks that the given task has expected output in the root of its output directory :param AzureMgmtTestCase test: test instance. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param str job_name: job name. :param str output_directory_id: output directory id. :param dict(str, str) expected: expected content, directories must have None value. """ Helpers.assert_job_files_in_path_are(test, client, resource_group, job_name, output_directory_id, '.', expected) @staticmethod def assert_existing_clusters_are(test, client, resource_group, expected): """Checks if there are expected set of clusters in the given resource group reported :param AzureMgmtTestCase test: test instance. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param list[str] expected: list of cluster names. """ actual = [c.name for c in list(client.clusters.list_by_workspace(resource_group, Helpers.DEFAULT_WORKSPACE_NAME))] test.assertListEqual(sorted(expected), sorted(actual)) @staticmethod def assert_existing_file_servers_are(test, client, resource_group, expected): """Checks if there are expected set of file servers in the given resource group reported :param AzureMgmtTestCase test: test instance. :param BatchAIManagementClient client: client instance. :param str resource_group: resource group name. :param list[str] expected: list of file servers names. """ actual = [s.name for s in list(client.file_servers.list_by_workspace(resource_group, Helpers.DEFAULT_WORKSPACE_NAME))] test.assertListEqual(sorted(expected), sorted(actual)) @staticmethod def assert_file_in_file_share(test, storage_account, storage_account_key, directory, filename, expected_content): """Checks if there is a file with given name and content exists in the Azure File share. :param AzureMgmtTestCase test: test instance. :param str storage_account: storage account name. :param str storage_account_key: storage account key. :param str directory: folder. :param str filename: filename. :param unicode expected_content: expected content. """ if not test.is_live: return service = FileService(storage_account, storage_account_key) actual = service.get_file_to_text(Helpers.AZURE_FILES_NAME, directory, filename).content test.assertEqual(expected_content, actual) @staticmethod def create_batchai_client(preparer): """Creates a Batch AI management client for tests. To create a custom version of the client (e.g. for integration environment), create custom_client.py file with create() method returning the instance of the client. :param AzureMgmtPreparer preparer: an instance of AzureMgmtPreparer :returns BatchAIManagementClient: an instance of Batch AI management client """ try: from custom_client import create as create_custom_client return create_custom_client() except ImportError: return preparer.create_mgmt_client(BatchAIManagementClient) class ClusterPreparer(AzureMgmtPreparer): """Batch AI cluster preparer""" def __init__(self, location=Helpers.LOCATION, vm_size='STANDARD_D1', target_nodes=1, wait=True, name_prefix='cluster', parameter_name='cluster', storage_account_parameter='storage_account', resource_group_parameter_name=RESOURCE_GROUP_PARAM): super(ClusterPreparer, self).__init__(name_prefix, 24) self.client = None # type: BatchAIManagementClient or None self.location = location self.resource_group_parameter_name = resource_group_parameter_name self.parameter_name = parameter_name self.storage_key = '' self.vm_size = vm_size self.target_nodes = target_nodes self.storage_account_param = storage_account_parameter self.wait = wait def create_resource(self, name, **kwargs): if self.is_live: self.client = Helpers.create_batchai_client(self) group = self._get_resource_group(**kwargs) self.resource = Helpers.create_cluster(self.client, self.location, group.name, name, self.vm_size, self.target_nodes, self._get_storage_account(**kwargs).name, self._get_storage_account_key(**kwargs)) self.client.experiments.create(group.name, Helpers.DEFAULT_WORKSPACE_NAME, Helpers.DEFAULT_EXPERIMENT_NAME).result() if self.wait: Helpers.wait_for_nodes(self.is_live, self.client, group.name, name, self.target_nodes, Helpers.NODE_STARTUP_TIMEOUT_SEC) else: self.resource = models.Cluster() self.resource.id = models.ResourceId(id='fake') return {self.parameter_name: self.resource} def remove_resource(self, name, **kwargs): if self.is_live: group = self._get_resource_group(**kwargs) self.client.clusters.delete(group.name, Helpers.DEFAULT_WORKSPACE_NAME, name).result() def _get_resource_group(self, **kwargs): try: return kwargs.get(self.resource_group_parameter_name) except KeyError: template = 'To create a cluster a resource group is required. Please add ' \ 'decorator @{} in front of this cluster preparer.' raise AzureTestError(template.format(ResourceGroupPreparer.__name__)) def _get_storage_account(self, **kwargs): try: return kwargs.get(self.storage_account_param) except KeyError: template = 'To create a cluster a storage account is required. Please add ' \ 'decorator @{} in front of this cluster preparer.' raise AzureTestError(template.format(StorageAccountPreparer.__name__)) def _get_storage_account_key(self, **kwargs): try: return kwargs.get(self.storage_account_param + '_key') except KeyError: template = 'To create a cluster a storage account is required. Please add ' \ 'decorator @{} in front of this cluster preparer.' raise AzureTestError(template.format(StorageAccountPreparer.__name__))