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)