def testRunInKubernetes(self): self._build_docker_images() temp_spill_dir = tempfile.mkdtemp(prefix='test-mars-k8s-') api_client = k8s_config.new_client_from_config() kube_api = k8s_client.CoreV1Api(api_client) cluster_client = None try: extra_vol_config = HostPathVolumeConfig('mars-src-path', '/mnt/mars', MARS_ROOT) cluster_client = new_cluster( api_client, image=self._docker_image, worker_spill_paths=[temp_spill_dir], extra_volumes=[extra_vol_config], pre_stop_command=['rm', '/tmp/stopping.tmp'], timeout=600, log_when_fail=True) self.assertIsNotNone(cluster_client.endpoint) pod_items = kube_api.list_namespaced_pod( cluster_client.namespace).to_dict() log_processes = [] for item in pod_items['items']: log_processes.append( subprocess.Popen([ 'kubectl', 'logs', '-f', '-n', cluster_client.namespace, item['metadata']['name'] ])) a = mt.ones((100, 100), chunk_size=30) * 2 * 1 + 1 b = mt.ones((100, 100), chunk_size=30) * 2 * 1 + 1 c = (a * b * 2 + 1).sum() r = cluster_client.session.run(c, timeout=600) expected = (np.ones(a.shape) * 2 * 1 + 1)**2 * 2 + 1 assert_array_equal(r, expected.sum()) # turn off service processes with grace to get coverage data procs = [] for item in pod_items['items']: p = subprocess.Popen([ 'kubectl', 'exec', '-n', cluster_client.namespace, item['metadata']['name'], '/srv/graceful_stop.sh' ]) procs.append(p) for p in procs: p.wait() [p.terminate() for p in log_processes] finally: shutil.rmtree(temp_spill_dir) if cluster_client: try: cluster_client.stop(wait=True, timeout=20) except TimeoutError: pass self._remove_docker_image(False)
def _start_kube_cluster(**kwargs): image_name = _build_docker_images() temp_spill_dir = tempfile.mkdtemp(prefix='test-mars-k8s-') api_client = k8s_config.new_client_from_config() kube_api = k8s_client.CoreV1Api(api_client) cluster_client = None try: extra_vol_config = HostPathVolumeConfig('mars-src-path', '/mnt/mars', MARS_ROOT) cluster_client = new_cluster( api_client, image=image_name, worker_spill_paths=[temp_spill_dir], extra_volumes=[extra_vol_config], pre_stop_command=['rm', '/tmp/stopping.tmp'], timeout=600, log_when_fail=True, **kwargs) assert cluster_client.endpoint is not None pod_items = kube_api.list_namespaced_pod( cluster_client.namespace).to_dict() log_processes = [] for item in pod_items['items']: log_processes.append( subprocess.Popen([ 'kubectl', 'logs', '-f', '-n', cluster_client.namespace, item['metadata']['name'] ])) yield # turn off service processes with grace to get coverage data procs = [] pod_items = kube_api.list_namespaced_pod( cluster_client.namespace).to_dict() for item in pod_items['items']: p = subprocess.Popen([ 'kubectl', 'exec', '-n', cluster_client.namespace, item['metadata']['name'], '--', '/srv/graceful_stop.sh' ]) procs.append(p) for p in procs: p.wait() [p.terminate() for p in log_processes] finally: shutil.rmtree(temp_spill_dir) if cluster_client: try: cluster_client.stop(wait=True, timeout=20) except TimeoutError: pass _collect_coverage() _remove_docker_image(image_name, False)
def test_create_timeout(): api_client = k8s_config.new_client_from_config() cluster = None try: extra_vol_config = HostPathVolumeConfig('mars-src-path', '/mnt/mars', MARS_ROOT) with pytest.raises(TimeoutError): cluster = new_cluster(api_client, image='pseudo_image', extra_volumes=[extra_vol_config], timeout=1) finally: if cluster: cluster.stop(wait=True) _collect_coverage()
def testCreateTimeout(self): api_client = k8s_config.new_client_from_config() cluster = None self._docker_image = 'pseudo_image' try: extra_vol_config = HostPathVolumeConfig('mars-src-path', '/mnt/mars', MARS_ROOT) with self.assertRaises(TimeoutError): cluster = new_cluster(api_client, image=self._docker_image, extra_volumes=[extra_vol_config], timeout=1) finally: if cluster: cluster.stop(wait=True)