def test_healthy_pod_service_cluster(self): tpu_map = { 'fake-pod': { 'state': 'READY', 'health': 'HEALTHY', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-32', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' } for ip in range(4)], } } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) service_workers = cr.get_service_workers() expected = [ ServiceWorker(internal_ip=f'10.0.0.{ip}', port='8470', machine_type='v3-32', zone='fake-zone', runtime_version='pytorch-nightly', tpu='fake-pod') for ip in range(4) ] self.assertCountEqual(expected, service_workers)
def test_healthy_vm_list_client_cluster(self): # Arrange list_instances_map = {} instance_resp_map = { 'fake-ig-' + c: gen_fake_instances_get_entry('fake-ig-' + c, 'n1-standard-16', '10.0.0.' + ip, 'RUNNING') for c, ip in zip('abcd', '0123') } compute_service = build_mock_compute_service(instance_resp_map, list_instances_map) self.mock_discovery.side_effect = build_mock_services_fn( compute_service) # Act vms = ['fake-ig-a', 'fake-ig-b', 'fake-ig-c', 'fake-ig-d'] cr = ClusterResolver(['fake-tpu'], vms=vms) vm_cluster = cr.get_client_workers() # Assert expected = [ ClientWorker(internal_ip='10.0.0.' + ip, machine_type='n1-standard-16', zone='fake-zone', hostname='fake-ig-' + c) for c, ip in zip('abcd', '0123') ] self.assertCountEqual(expected, vm_cluster)
def test_healthy_tpuvm_cluster(self): # Using TPUVM flavor of metadata. mock.patch.object(ClusterResolver, 'get_instance_metadata', mock_request_tpuvm_metadata).start() noop_compute_service = build_mock_compute_service({}, {}) self.mock_discovery.side_effect = build_mock_services_fn( noop_compute_service) tpu_map = { 'fake-pod': { 'state': 'READY', 'health': 'HEALTHY', 'runtime_version': 'v2-nightly', 'accelerator_type': 'v3-32', 'api_version': 'V2_ALPHA1', 'network_endpoints': [{ 'ipAddress': f'10.1.0.{index}', 'port': '8470', } for index in range(4)], } } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) cluster = cr.get_cluster() expected_client_workers = [ ClientWorker(internal_ip=f'10.1.0.{index}', machine_type='v3-32', zone='fake-zone', hostname=f'{TPUVM_HOSTNAME_PREFIX}{index}') for index in range(4) ] expected_service_workers = [ ServiceWorker(internal_ip=f'10.1.0.{ip}', port='8470', machine_type='v3-32', zone='fake-zone', runtime_version='v2-nightly', tpu='fake-pod') for ip in range(4) ] expected = Cluster(expected_client_workers, expected_service_workers, client_master_ip='10.1.0.0') self.assertEqual(expected, cluster) mock.patch.object(ClusterResolver, 'get_instance_metadata', mock_request_metadata).start()
def test_unknown_health_pod_service_cluster(self): noop_compute_service = build_mock_compute_service({}, {}) self.mock_discovery.side_effect = build_mock_services_fn( noop_compute_service) tpu_map = { 'fake-pod': { 'state': 'READY', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-32', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' } for ip in range(4)], } } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) self.assertRaisesRegex(RuntimeError, 'TPU fake-pod is not HEALTHY yet.*', cr.get_service_workers)
def test_non_ready_sea_service_cluster(self): noop_compute_service = build_mock_compute_service({}, {}) self.mock_discovery.side_effect = build_mock_services_fn( noop_compute_service) tpu_map = { f'fake-tpu-{ip}': { 'state': 'READY', 'health': 'HEALTHY', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-8', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' }], } for ip in range(3) } tpu_map['fake-tpu-3'] = { 'state': 'CREATING', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-8', } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) self.assertRaisesRegex(RuntimeError, 'TPU fake-tpu-3 is not READY yet.*', cr.get_service_workers)
def test_unhealthy_pod_service_cluster(self): tpu_map = { 'fake-pod': { 'state': 'READY', 'health': 'UNHEALTHY_TENSORFLOW', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-128', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' } for ip in range(16)], } } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) self.assertRaisesRegex(RuntimeError, 'TPU fake-pod is not HEALTHY yet.*', cr.get_service_workers)
def __init__(self, cluster, docker_container=None, docker_image=None, docker_run_flags=None, conda_env=None, env_vars=None): self._cluster = cluster self._initialize() client_master_ip = ClusterResolver.get_instance_metadata( 'instance/network-interfaces/0/ip') self._client_master = next( filter(lambda cw: cw.get_internal_ip() == client_master_ip, self._cluster.get_client_workers())) self.logger = self._get_logger() self.docker_container = docker_container or self.DEFAULT_CONTAINER_NAME self.docker_image = docker_image self.docker_run_flags = list( docker_run_flags) if docker_run_flags else [] self.conda_env = conda_env self.env_vars = list(env_vars) if env_vars else [] for env_var in self.env_vars: if re.match('\w*=\w*', env_var) is None: raise ValueError( ('Environment variable to distribute ({}) should follow ' 'the form: X=Y').format(env_var)) for dist_var in self.DIST_ENV_VARS: if re.match('{}=.*'.format(dist_var), env_var): raise ValueError(( '{} should not be in the training command provided as they' ' will interfere with the values set for distributed' ' training'.format(dist_var)))
def test_healthy_sea_service_cluster(self): noop_compute_service = build_mock_compute_service({}, {}) self.mock_discovery.side_effect = build_mock_services_fn( noop_compute_service) tpu_map = { f'fake-tpu-{ip}': { 'state': 'READY', 'health': 'HEALTHY', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-8', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' }], } for ip in range(256) } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) service_workers = cr.get_service_workers() expected = [ ServiceWorker(internal_ip=f'10.0.0.{ip}', port='8470', machine_type='v3-8', zone='fake-zone', runtime_version='pytorch-nightly', tpu=f'fake-tpu-{ip}') for ip in range(256) ] self.assertCountEqual(expected, service_workers)
def test_bad_cluster(self): list_instances_map = { 'fake-ig': { 'kind': 'compute#instanceGroupsListInstances', 'items': [ gen_fake_ig_list_instances_entry('fake-ig-' + c, 'RUNNING') for c in 'abc' ], }, } instance_resp_map = { 'fake-ig-' + c: gen_fake_instances_get_entry('fake-ig-' + c, 'n1-standard-16', '10.0.0.' + ip, 'RUNNING') for c, ip in zip('abcd', '0123') } compute_service = build_mock_compute_service(instance_resp_map, list_instances_map) self.mock_discovery.side_effect = build_mock_services_fn( compute_service) tpu_map = { 'fake-pod': { 'state': 'READY', 'health': 'HEALTHY', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-32', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' } for ip in range(4)], } } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) self.assertRaisesRegex( RuntimeError, 'The client_workers and service_workers must have a 1:1 mapping', cr.get_cluster)
def test_unhealthy_client_cluster(self): # Arrange list_instances_map = { 'fake-ig': { 'kind': 'compute#instanceGroupsListInstances', 'items': [ gen_fake_ig_list_instances_entry('fake-ig-a', 'RUNNING'), gen_fake_ig_list_instances_entry('fake-ig-b', 'PROVISIONING'), gen_fake_ig_list_instances_entry('fake-ig-c', 'RUNNING'), gen_fake_ig_list_instances_entry('fake-ig-d', 'RUNNING'), ], }, } instance_resp_map = { 'fake-ig-a': gen_fake_instances_get_entry('fake-ig-a', 'n1-standard-16', '10.0.0.0', 'RUNNING'), 'fake-ig-b': gen_fake_instances_get_entry('fake-ig-b', 'n1-standard-16', '10.0.0.1', 'PROVISIONING'), 'fake-ig-c': gen_fake_instances_get_entry('fake-ig-c', 'n1-standard-16', '10.0.0.2', 'RUNNING'), 'fake-ig-d': gen_fake_instances_get_entry('fake-ig-d', 'n1-standard-16', '10.0.0.3', 'RUNNING'), } compute_service = build_mock_compute_service(instance_resp_map, list_instances_map) self.mock_discovery.side_effect = build_mock_services_fn( compute_service) # Act cr = ClusterResolver(['fake-tpu']) # Assert self.assertRaisesRegex(RuntimeError, 'Instance fake-ig-b is not running yet.*', cr.get_client_workers)
def test_empty_instance_group_client_cluster(self): list_instances_map = { 'fake-ig': { 'kind': 'compute#instanceGroupsListInstances', 'items': [], }, } instance_resp_map = { 'fake-ig-a': gen_fake_instances_get_entry('fake-ig-a', 'n1-standard-16', '10.0.0.0', 'RUNNING'), } compute_service = build_mock_compute_service(instance_resp_map, list_instances_map) self.mock_discovery.side_effect = build_mock_services_fn( compute_service) # Act cr = ClusterResolver(['fake-tpu']) # Assert self.assertRaisesRegex(RuntimeError, '.*vms is empty in instance group.*', cr.get_client_workers)
help='Name of the conda environment if running with conda.') parser.add_argument('--env', action='append', type=str, help='List of environment variables to distribute.') parser.add_argument( 'positional', nargs='+', type=str, help='The python command to launch training including model parameters.' ) FLAGS = parser.parse_args() tpuvm_mode = False accel_type = ClusterResolver.get_instance_metadata( 'instance/attributes/accelerator-type') if re.match(r'v[0-9]+-[0-9]+', accel_type): # Only TPUVM will carry the accelerator-type metadata tpuvm_mode = True if (FLAGS.docker_container or FLAGS.docker_image or FLAGS.docker_run_flag) and FLAGS.conda_env: raise ValueError('Docker Setup arguments and Conda Setup' ' arguments are mutually exclusive.') # Resolve VM and TPU clusters. cluster_resolver = ClusterResolver(FLAGS.tpu, vms=FLAGS.vm, tpuvm_mode=tpuvm_mode) cluster = cluster_resolver.get_cluster() executor = DistributedExecutor(cluster,
def test_healthy_cluster(self): list_instances_map = { 'fake-ig': { 'kind': 'compute#instanceGroupsListInstances', 'items': [ gen_fake_ig_list_instances_entry('fake-ig-' + c, 'RUNNING') for c in 'abcd' ], }, } instance_resp_map = { 'fake-ig-' + c: gen_fake_instances_get_entry('fake-ig-' + c, 'n1-standard-16', '10.0.0.' + ip, 'RUNNING') for c, ip in zip('abcd', '0123') } compute_service = build_mock_compute_service(instance_resp_map, list_instances_map) self.mock_discovery.side_effect = build_mock_services_fn( compute_service) tpu_map = { 'fake-pod': { 'state': 'READY', 'health': 'HEALTHY', 'runtime_version': 'pytorch-nightly', 'accelerator_type': 'v3-32', 'network_endpoints': [{ 'ipAddress': f'10.0.0.{ip}', 'port': '8470' } for ip in range(4)], } } self.mock_ctc.side_effect = build_mock_cloud_tpu_client_library( tpu_map) tpus = list(tpu_map.keys()) cr = ClusterResolver(tpus) cluster = cr.get_cluster() expected_client_workers = [ ClientWorker(internal_ip='10.0.0.' + ip, machine_type='n1-standard-16', zone='fake-zone', hostname='fake-ig-' + c) for c, ip in zip('abcd', '0123') ] expected_service_workers = [ ServiceWorker(internal_ip=f'10.0.0.{ip}', port='8470', machine_type='v3-32', zone='fake-zone', runtime_version='pytorch-nightly', tpu='fake-pod') for ip in range(4) ] expected = Cluster(expected_client_workers, expected_service_workers) self.assertEqual(expected, cluster)
help='Name of the conda environment if running with conda.') parser.add_argument('--env', action='append', type=str, help='List of environment variables to distribute.') parser.add_argument( 'positional', nargs='+', type=str, help='The python command to launch training including model parameters.' ) FLAGS = parser.parse_args() if (FLAGS.docker_container or FLAGS.docker_image or FLAGS.docker_run_flag) and FLAGS.conda_env: raise ValueError('Docker Setup arguments and Conda Setup' ' arguments are mutually exclusive.') # Resolve VM and TPU clusters. cluster_resolver = ClusterResolver(FLAGS.tpu, vms=FLAGS.vm) cluster = cluster_resolver.get_cluster() executor = DistributedExecutor(cluster, docker_container=FLAGS.docker_container, docker_image=FLAGS.docker_image, docker_run_flags=FLAGS.docker_run_flag, conda_env=FLAGS.conda_env, env_vars=FLAGS.env) executor.run(FLAGS.positional)
help='Port that XRT local service will be start on.') parser.add_argument( 'positional', nargs='+', type=str, help='The python command to launch training including model parameters.' ) FLAGS = parser.parse_args() if (FLAGS.docker_container or FLAGS.docker_image or FLAGS.docker_run_flag) and FLAGS.conda_env: raise ValueError('Docker Setup arguments and Conda Setup' ' arguments are mutually exclusive.') # Resolve VM and TPU clusters. cluster_resolver = ClusterResolver(FLAGS.tpu, vms=FLAGS.vm) cluster = cluster_resolver.get_cluster() tpuvm_mode = cluster_resolver.get_tpuvm_mode() executor = DistributedExecutor( cluster, docker_container=FLAGS.docker_container, docker_image=FLAGS.docker_image, docker_run_flags=FLAGS.docker_run_flag, conda_env=FLAGS.conda_env, env_vars=FLAGS.env, restart_server=FLAGS.restart_tpuvm_pod_server, tpuvm_mode=tpuvm_mode, tpuvm_server_port=FLAGS.tpuvm_server_port) executor.run(FLAGS.positional)