def setUp(self):
self.backend = ServiceBackend()
self.bucket_name = get_user_config().get('batch', 'bucket')
self.gcs_input_dir = f'gs://{self.bucket_name}/batch-tests/resources'
token = uuid.uuid4()
self.gcs_output_path = f'/batch-tests/{token}'
self.gcs_output_dir = f'gs://{self.bucket_name}{self.gcs_output_path}'
in_cluster_key_file = '/test-gsa-key/key.json'
if os.path.exists(in_cluster_key_file):
credentials = google.oauth2.service_account.Credentials.from_service_account_file(
in_cluster_key_file)
else:
credentials = None
gcs_client = google.cloud.storage.Client(project='hail-vdc',
credentials=credentials)
bucket = gcs_client.bucket(self.bucket_name)
if not bucket.blob('batch-tests/resources/hello.txt').exists():
bucket.blob('batch-tests/resources/hello.txt').upload_from_string(
'hello world')
if not bucket.blob('batch-tests/resources/hello spaces.txt').exists():
bucket.blob('batch-tests/resources/hello spaces.txt'
).upload_from_string('hello')
if not bucket.blob(
'batch-tests/resources/hello (foo) spaces.txt').exists():
bucket.blob('batch-tests/resources/hello (foo) spaces.txt'
).upload_from_string('hello')
def setUp(self):
self.backend = ServiceBackend()
remote_tmpdir = get_user_config().get('batch', 'remote_tmpdir')
if not remote_tmpdir.endswith('/'):
remote_tmpdir += '/'
self.remote_tmpdir = remote_tmpdir
if remote_tmpdir.startswith('gs://'):
self.bucket = re.fullmatch(
'gs://(?P<bucket_name>[^/]+).*',
remote_tmpdir).groupdict()['bucket_name']
else:
assert remote_tmpdir.startswith('hail-az://')
storage_account, container_name = re.fullmatch(
'hail-az://(?P<storage_account>[^/]+)/(?P<container_name>[^/]+).*',
remote_tmpdir).groups()
self.bucket = f'{storage_account}/{container_name}'
self.cloud_input_dir = f'{self.remote_tmpdir}batch-tests/resources'
token = uuid.uuid4()
self.cloud_output_path = f'/batch-tests/{token}'
self.cloud_output_dir = f'{self.remote_tmpdir}{self.cloud_output_path}'
in_cluster_key_file = '/test-gsa-key/key.json'
if not os.path.exists(in_cluster_key_file):
in_cluster_key_file = None
router_fs = RouterAsyncFS(
'gs',
gcs_kwargs={
'project': 'hail-vdc',
'credentials_file': in_cluster_key_file
},
azure_kwargs={'credential_file': in_cluster_key_file})
def sync_exists(url):
return async_to_blocking(router_fs.exists(url))
def sync_write(url, data):
return async_to_blocking(router_fs.write(url, data))
if not sync_exists(
f'{self.remote_tmpdir}batch-tests/resources/hello.txt'):
sync_write(f'{self.remote_tmpdir}batch-tests/resources/hello.txt',
b'hello world')
if not sync_exists(
f'{self.remote_tmpdir}batch-tests/resources/hello spaces.txt'):
sync_write(
f'{self.remote_tmpdir}batch-tests/resources/hello spaces.txt',
b'hello')
if not sync_exists(
f'{self.remote_tmpdir}batch-tests/resources/hello (foo) spaces.txt'
):
sync_write(
f'{self.remote_tmpdir}batch-tests/resources/hello (foo) spaces.txt',
b'hello')
def test_large_command(self):
backend = ServiceBackend(
remote_tmpdir=f'gs://{HAIL_TEST_GCS_BUCKET}/temporary-files')
b = Batch(backend=backend)
j1 = b.new_job()
long_str = secrets.token_urlsafe(15 * 1024)
j1.command(f'echo "{long_str}"')
b.run()
def test_service_backend_bucket_parameter(self):
backend = ServiceBackend(bucket=HAIL_TEST_GCS_BUCKET)
b = Batch(backend=backend)
j1 = b.new_job()
j1.command(f'echo hello > {j1.ofile}')
j2 = b.new_job()
j2.command(f'cat {j1.ofile}')
b.run()
def test_service_backend_remote_tempdir_with_no_trailing_slash(self):
backend = ServiceBackend(
remote_tmpdir=f'gs://{HAIL_TEST_GCS_BUCKET}/temporary-files/')
b = Batch(backend=backend)
j1 = b.new_job()
j1.command(f'echo hello > {j1.ofile}')
j2 = b.new_job()
j2.command(f'cat {j1.ofile}')
b.run()
def test_big_batch_which_uses_slow_path(self):
backend = ServiceBackend(
remote_tmpdir=f'{self.remote_tmpdir}/temporary-files')
b = Batch(backend=backend)
# 8 * 256 * 1024 = 2 MiB > 1 MiB max bunch size
for i in range(8):
j1 = b.new_job()
long_str = secrets.token_urlsafe(256 * 1024)
j1.command(f'echo "{long_str}"')
batch = b.run()
assert not batch.submission_info.used_fast_create
batch_status = batch.status()
assert batch_status['state'] == 'success', str((batch.debug_info()))
class ServiceTests(unittest.TestCase):
def setUp(self):
self.backend = ServiceBackend()
self.bucket_name = get_user_config().get('batch', 'bucket')
self.gcs_input_dir = f'gs://{self.bucket_name}/batch-tests/resources'
token = uuid.uuid4()
self.gcs_output_path = f'/batch-tests/{token}'
self.gcs_output_dir = f'gs://{self.bucket_name}{self.gcs_output_path}'
in_cluster_key_file = '/test-gsa-key/key.json'
if os.path.exists(in_cluster_key_file):
credentials = google.oauth2.service_account.Credentials.from_service_account_file(
in_cluster_key_file)
else:
credentials = None
gcs_client = google.cloud.storage.Client(project='hail-vdc',
credentials=credentials)
bucket = gcs_client.bucket(self.bucket_name)
if not bucket.blob('batch-tests/resources/hello.txt').exists():
bucket.blob('batch-tests/resources/hello.txt').upload_from_string(
'hello world')
if not bucket.blob('batch-tests/resources/hello spaces.txt').exists():
bucket.blob('batch-tests/resources/hello spaces.txt'
).upload_from_string('hello')
if not bucket.blob(
'batch-tests/resources/hello (foo) spaces.txt').exists():
bucket.blob('batch-tests/resources/hello (foo) spaces.txt'
).upload_from_string('hello')
def tearDown(self):
self.backend.close()
def batch(self, requester_pays_project=None, default_python_image=None):
return Batch(backend=self.backend,
default_image=DOCKER_ROOT_IMAGE,
attributes={
'foo': 'a',
'bar': 'b'
},
requester_pays_project=requester_pays_project,
default_python_image=default_python_image)
def test_single_task_no_io(self):
b = self.batch()
j = b.new_job()
j.command('echo hello')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_task_input(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_task_input_resource_group(self):
b = self.batch()
input = b.read_input_group(foo=f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.storage('10Gi')
j.command(f'cat {input.foo}')
j.command(f'cat {input}.foo')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_task_output(self):
b = self.batch()
j = b.new_job(attributes={'a': 'bar', 'b': 'foo'})
j.command(f'echo hello > {j.ofile}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_task_write_output(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile,
f'{self.gcs_output_dir}/test_single_task_output.txt')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_task_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_task_write_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
b.write_output(
j.output,
f'{self.gcs_output_dir}/test_single_task_write_resource_group')
b.write_output(
j.output.foo,
f'{self.gcs_output_dir}/test_single_task_write_resource_group_file.txt'
)
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_multiple_dependent_tasks(self):
output_file = f'{self.gcs_output_dir}/test_multiple_dependent_tasks.txt'
b = self.batch()
j = b.new_job()
j.command(f'echo "0" >> {j.ofile}')
for i in range(1, 3):
j2 = b.new_job()
j2.command(f'echo "{i}" > {j2.tmp1}')
j2.command(f'cat {j.ofile} {j2.tmp1} > {j2.ofile}')
j = j2
b.write_output(j.ofile, output_file)
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_specify_cpu(self):
b = self.batch()
j = b.new_job()
j.cpu('0.5')
j.command(f'echo "hello" > {j.ofile}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_specify_memory(self):
b = self.batch()
j = b.new_job()
j.memory('100M')
j.command(f'echo "hello" > {j.ofile}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_scatter_gather(self):
b = self.batch()
for i in range(3):
j = b.new_job(name=f'foo{i}')
j.command(f'echo "{i}" > {j.ofile}')
merger = b.new_job()
merger.command('cat {files} > {ofile}'.format(files=' '.join([
j.ofile for j in sorted(
b.select_jobs('foo'), key=lambda x: x.name, reverse=True)
]),
ofile=merger.ofile))
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_file_name_space(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello (foo) spaces.txt')
j = b.new_job()
j.command(f'cat {input} > {j.ofile}')
b.write_output(j.ofile,
f'{self.gcs_output_dir}/hello (foo) spaces.txt')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_dry_run(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile,
f'{self.gcs_output_dir}/test_single_job_output.txt')
b.run(dry_run=True)
def test_verbose(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
b.write_output(input, f'{self.gcs_output_dir}/hello.txt')
res = b.run(verbose=True)
assert res.status()['state'] == 'success', debug_info(res)
def test_gcsfuse(self):
path = f'/{self.bucket_name}{self.gcs_output_path}'
b = self.batch()
head = b.new_job()
head.command(f'mkdir -p {path}; echo head > {path}/gcsfuse_test_1')
head.gcsfuse(self.bucket_name, f'/{self.bucket_name}', read_only=False)
tail = b.new_job()
tail.command(f'cat {path}/gcsfuse_test_1')
tail.gcsfuse(self.bucket_name, f'/{self.bucket_name}', read_only=True)
tail.depends_on(head)
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_gcsfuse_read_only(self):
path = f'/{self.bucket_name}{self.gcs_output_path}'
b = self.batch()
j = b.new_job()
j.command(f'mkdir -p {path}; echo head > {path}/gcsfuse_test_1')
j.gcsfuse(self.bucket_name, f'/{self.bucket_name}', read_only=True)
res = b.run()
assert res.status()['state'] == 'failure', debug_info(res)
def test_gcsfuse_implicit_dirs(self):
path = f'/{self.bucket_name}{self.gcs_output_path}'
b = self.batch()
head = b.new_job()
head.command(
f'mkdir -p {path}/gcsfuse/; echo head > {path}/gcsfuse/data')
head.gcsfuse(self.bucket_name, f'/{self.bucket_name}', read_only=False)
tail = b.new_job()
tail.command(f'cat {path}/gcsfuse/data')
tail.gcsfuse(self.bucket_name, f'/{self.bucket_name}', read_only=True)
tail.depends_on(head)
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_requester_pays(self):
b = self.batch(requester_pays_project='hail-vdc')
input = b.read_input('gs://hail-services-requester-pays/hello')
j = b.new_job()
j.command(f'cat {input}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_benchmark_lookalike_workflow(self):
b = self.batch()
setup_jobs = []
for i in range(10):
j = b.new_job(f'setup_{i}').cpu(0.25)
j.command(f'echo "foo" > {j.ofile}')
setup_jobs.append(j)
jobs = []
for i in range(500):
j = b.new_job(f'create_file_{i}').cpu(0.25)
j.command(
f'echo {setup_jobs[i % len(setup_jobs)].ofile} > {j.ofile}')
j.command(f'echo "bar" >> {j.ofile}')
jobs.append(j)
combine = b.new_job(f'combine_output').cpu(0.25)
for tasks in grouped(arg_max(), jobs):
combine.command(
f'cat {" ".join(shq(j.ofile) for j in jobs)} >> {combine.ofile}'
)
b.write_output(combine.ofile,
f'{self.gcs_output_dir}/pipeline_benchmark_test.txt')
# too slow
# assert b.run().status()['state'] == 'success'
def test_envvar(self):
b = self.batch()
j = b.new_job()
j.env('SOME_VARIABLE', '123abcdef')
j.command('[ $SOME_VARIABLE = "123abcdef" ]')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_job_with_shell(self):
msg = 'hello world'
b = self.batch()
j = b.new_job(shell='/bin/sh')
j.command(f'echo "{msg}"')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_single_job_with_nonsense_shell(self):
b = self.batch()
j = b.new_job(shell='/bin/ajdsfoijasidojf')
j.command(f'echo "hello"')
res = b.run()
assert res.status()['state'] == 'failure', debug_info(res)
def test_single_job_with_intermediate_failure(self):
b = self.batch()
j = b.new_job()
j.command(f'echoddd "hello"')
j2 = b.new_job()
j2.command(f'echo "world"')
res = b.run()
assert res.status()['state'] == 'failure', debug_info(res)
def test_input_directory(self):
b = self.batch()
input1 = b.read_input(self.gcs_input_dir)
input2 = b.read_input(self.gcs_input_dir.rstrip('/') + '/')
j = b.new_job()
j.command(f'ls {input1}/hello.txt')
j.command(f'ls {input2}/hello.txt')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
def test_python_job(self):
b = self.batch(
default_python_image='gcr.io/hail-vdc/python-dill:3.7-slim')
head = b.new_job()
head.command(f'echo "5" > {head.r5}')
head.command(f'echo "3" > {head.r3}')
def read(path):
with open(path, 'r') as f:
i = f.read()
return int(i)
def multiply(x, y):
return x * y
def reformat(x, y):
return {'x': x, 'y': y}
middle = b.new_python_job()
r3 = middle.call(read, head.r3)
r5 = middle.call(read, head.r5)
r_mult = middle.call(multiply, r3, r5)
middle2 = b.new_python_job()
r_mult = middle2.call(multiply, r_mult, 2)
r_dict = middle2.call(reformat, r3, r5)
tail = b.new_job()
tail.command(
f'cat {r3.as_str()} {r5.as_repr()} {r_mult.as_str()} {r_dict.as_json()}'
)
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
assert res.get_job_log(
4)['main'] == "3\n5\n30\n{\"x\": 3, \"y\": 5}\n", debug_info(res)
def test_python_job_w_resource_group_unpack_individually(self):
b = self.batch(
default_python_image='gcr.io/hail-vdc/python-dill:3.7-slim')
head = b.new_job()
head.declare_resource_group(count={
'r5': '{root}.r5',
'r3': '{root}.r3'
})
head.command(f'echo "5" > {head.count.r5}')
head.command(f'echo "3" > {head.count.r3}')
def read(path):
with open(path, 'r') as f:
r = int(f.read())
return r
def multiply(x, y):
return x * y
def reformat(x, y):
return {'x': x, 'y': y}
middle = b.new_python_job()
r3 = middle.call(read, head.count.r3)
r5 = middle.call(read, head.count.r5)
r_mult = middle.call(multiply, r3, r5)
middle2 = b.new_python_job()
r_mult = middle2.call(multiply, r_mult, 2)
r_dict = middle2.call(reformat, r3, r5)
tail = b.new_job()
tail.command(
f'cat {r3.as_str()} {r5.as_repr()} {r_mult.as_str()} {r_dict.as_json()}'
)
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
assert res.get_job_log(
4)['main'] == "3\n5\n30\n{\"x\": 3, \"y\": 5}\n", debug_info(res)
def test_python_job_w_resource_group_unpack_jointly(self):
b = self.batch(
default_python_image='gcr.io/hail-vdc/python-dill:3.7-slim')
head = b.new_job()
head.declare_resource_group(count={
'r5': '{root}.r5',
'r3': '{root}.r3'
})
head.command(f'echo "5" > {head.count.r5}')
head.command(f'echo "3" > {head.count.r3}')
def read_rg(root):
with open(root['r3'], 'r') as f:
r3 = int(f.read())
with open(root['r5'], 'r') as f:
r5 = int(f.read())
return (r3, r5)
def multiply(r):
x, y = r
return x * y
middle = b.new_python_job()
r = middle.call(read_rg, head.count)
r_mult = middle.call(multiply, r)
tail = b.new_job()
tail.command(f'cat {r_mult.as_str()}')
res = b.run()
assert res.status()['state'] == 'success', debug_info(res)
assert res.get_job_log(3)['main'] == "15\n", debug_info(res)
class BatchTests(unittest.TestCase):
def setUp(self):
self.backend = ServiceBackend()
bucket_name = get_userinfo()['bucket_name']
token = uuid.uuid4()
self.gcs_input_dir = f'gs://{bucket_name}/batch-tests/resources'
self.gcs_output_dir = f'gs://{bucket_name}/batch-tests/{token}'
in_cluster_key_file = '/test-gsa-key/key.json'
if os.path.exists(in_cluster_key_file):
credentials = google.oauth2.service_account.Credentials.from_service_account_file(
in_cluster_key_file)
else:
credentials = None
gcs_client = google.cloud.storage.Client(project='hail-vdc', credentials=credentials)
bucket = gcs_client.bucket(bucket_name)
if not bucket.blob('batch-tests/resources/hello (foo) spaces.txt').exists():
bucket.blob('batch-tests/resources/hello.txt').upload_from_string(
'hello world')
bucket.blob('batch-tests/resources/hello spaces.txt').upload_from_string(
'hello')
bucket.blob('batch-tests/resources/hello (foo) spaces.txt').upload_from_string(
'hello')
def tearDown(self):
self.backend.close()
def batch(self):
return Batch(backend=self.backend,
default_image='google/cloud-sdk:237.0.0-alpine',
attributes={'foo': 'a', 'bar': 'b'})
def test_single_task_no_io(self):
b = self.batch()
j = b.new_job()
j.command('echo hello')
assert b.run().status()['state'] == 'success'
def test_single_task_input(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
assert b.run().status()['state'] == 'success'
def test_single_task_input_resource_group(self):
b = self.batch()
input = b.read_input_group(foo=f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.storage('0.25Gi')
j.command(f'cat {input.foo}')
assert b.run().status()['state'] == 'success'
def test_single_task_output(self):
b = self.batch()
j = b.new_job(attributes={'a': 'bar', 'b': 'foo'})
j.command(f'echo hello > {j.ofile}')
assert b.run().status()['state'] == 'success'
def test_single_task_write_output(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile, f'{self.gcs_output_dir}/test_single_task_output.txt')
assert b.run().status()['state'] == 'success'
def test_single_task_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
assert b.run().status()['state'] == 'success'
def test_single_task_write_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
b.write_output(j.output, f'{self.gcs_output_dir}/test_single_task_write_resource_group')
b.write_output(j.output.foo, f'{self.gcs_output_dir}/test_single_task_write_resource_group_file.txt')
assert b.run().status()['state'] == 'success'
def test_multiple_dependent_tasks(self):
output_file = f'{self.gcs_output_dir}/test_multiple_dependent_tasks.txt'
b = self.batch()
j = b.new_job()
j.command(f'echo "0" >> {j.ofile}')
for i in range(1, 3):
j2 = b.new_job()
j2.command(f'echo "{i}" > {j2.tmp1}')
j2.command(f'cat {j.ofile} {j2.tmp1} > {j2.ofile}')
j = j2
b.write_output(j.ofile, output_file)
assert b.run().status()['state'] == 'success'
def test_specify_cpu(self):
b = self.batch()
j = b.new_job()
j.cpu('0.5')
j.command(f'echo "hello" > {j.ofile}')
assert b.run().status()['state'] == 'success'
def test_specify_memory(self):
b = self.batch()
j = b.new_job()
j.memory('100M')
j.command(f'echo "hello" > {j.ofile}')
assert b.run().status()['state'] == 'success'
def test_scatter_gather(self):
b = self.batch()
for i in range(3):
j = b.new_job(name=f'foo{i}')
j.command(f'echo "{i}" > {j.ofile}')
merger = b.new_job()
merger.command('cat {files} > {ofile}'.format(files=' '.join([j.ofile for j in sorted(b.select_jobs('foo'),
key=lambda x: x.name,
reverse=True)]),
ofile=merger.ofile))
assert b.run().status()['state'] == 'success'
def test_file_name_space(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello (foo) spaces.txt')
j = b.new_job()
j.command(f'cat {input} > {j.ofile}')
b.write_output(j.ofile, f'{self.gcs_output_dir}/hello (foo) spaces.txt')
assert b.run().status()['state'] == 'success'
def test_dry_run(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile, f'{self.gcs_output_dir}/test_single_job_output.txt')
b.run(dry_run=True)
def test_verbose(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
b.write_output(input, f'{self.gcs_output_dir}/hello.txt')
assert b.run(verbose=True).status()['state'] == 'success'
def test_benchmark_lookalike_workflow(self):
b = self.batch()
setup_jobs = []
for i in range(10):
j = b.new_job(f'setup_{i}').cpu(0.1)
j.command(f'echo "foo" > {j.ofile}')
setup_jobs.append(j)
jobs = []
for i in range(500):
j = b.new_job(f'create_file_{i}').cpu(0.1)
j.command(f'echo {setup_jobs[i % len(setup_jobs)].ofile} > {j.ofile}')
j.command(f'echo "bar" >> {j.ofile}')
jobs.append(j)
combine = b.new_job(f'combine_output').cpu(0.1)
for tasks in grouped(arg_max(), jobs):
combine.command(f'cat {" ".join(shq(j.ofile) for j in jobs)} >> {combine.ofile}')
b.write_output(combine.ofile, f'{self.gcs_output_dir}/pipeline_benchmark_test.txt')
def setUp(self):
self.backend = ServiceBackend('test')
self.gcs_input_dir = 'gs://hail-services/batch-testing/resources'
self.gcs_output_dir = os.environ.get('SCRATCH').rstrip('/') + '/output'
class BatchTests(unittest.TestCase):
def setUp(self):
self.backend = ServiceBackend('test')
self.gcs_input_dir = 'gs://hail-services/batch-testing/resources'
self.gcs_output_dir = os.environ.get('SCRATCH').rstrip('/') + '/output'
def tearDown(self):
self.backend.close()
def batch(self):
return Batch(backend=self.backend,
default_image='google/cloud-sdk:237.0.0-alpine',
attributes={
'foo': 'a',
'bar': 'b'
})
def test_single_task_no_io(self):
b = self.batch()
j = b.new_job()
j.command('echo hello')
assert b.run().status()['state'] == 'success'
def test_single_task_input(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
assert b.run().status()['state'] == 'success'
def test_single_task_input_resource_group(self):
b = self.batch()
input = b.read_input_group(foo=f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.storage('0.25Gi')
j.command(f'cat {input.foo}')
assert b.run().status()['state'] == 'success'
def test_single_task_output(self):
b = self.batch()
j = b.new_job(attributes={'a': 'bar', 'b': 'foo'})
j.command(f'echo hello > {j.ofile}')
assert b.run().status()['state'] == 'success'
def test_single_task_write_output(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile,
f'{self.gcs_output_dir}/test_single_task_output.txt')
assert b.run().status()['state'] == 'success'
def test_single_task_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
assert b.run().status()['state'] == 'success'
def test_single_task_write_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
b.write_output(
j.output,
f'{self.gcs_output_dir}/test_single_task_write_resource_group')
b.write_output(
j.output.foo,
f'{self.gcs_output_dir}/test_single_task_write_resource_group_file.txt'
)
assert b.run().status()['state'] == 'success'
def test_multiple_dependent_tasks(self):
output_file = f'{self.gcs_output_dir}/test_multiple_dependent_tasks.txt'
b = self.batch()
j = b.new_job()
j.command(f'echo "0" >> {j.ofile}')
for i in range(1, 3):
j2 = b.new_job()
j2.command(f'echo "{i}" > {j2.tmp1}')
j2.command(f'cat {j.ofile} {j2.tmp1} > {j2.ofile}')
j = j2
b.write_output(j.ofile, output_file)
assert b.run().status()['state'] == 'success'
def test_specify_cpu(self):
b = self.batch()
j = b.new_job()
j.cpu('0.5')
j.command(f'echo "hello" > {j.ofile}')
assert b.run().status()['state'] == 'success'
def test_specify_memory(self):
b = self.batch()
j = b.new_job()
j.memory('100M')
j.command(f'echo "hello" > {j.ofile}')
assert b.run().status()['state'] == 'success'
def test_scatter_gather(self):
b = self.batch()
for i in range(3):
j = b.new_job(name=f'foo{i}')
j.command(f'echo "{i}" > {j.ofile}')
merger = b.new_job()
merger.command('cat {files} > {ofile}'.format(files=' '.join([
j.ofile for j in sorted(
b.select_jobs('foo'), key=lambda x: x.name, reverse=True)
]),
ofile=merger.ofile))
assert b.run().status()['state'] == 'success'
def test_file_name_space(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello (foo) spaces.txt')
j = b.new_job()
j.command(f'cat {input} > {j.ofile}')
b.write_output(j.ofile,
f'{self.gcs_output_dir}/hello (foo) spaces.txt')
assert b.run().status()['state'] == 'success'
def test_dry_run(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile,
f'{self.gcs_output_dir}/test_single_job_output.txt')
b.run(dry_run=True)
def test_verbose(self):
b = self.batch()
input = b.read_input(f'{self.gcs_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
b.write_output(input, f'{self.gcs_output_dir}/hello.txt')
assert b.run(verbose=True).status()['state'] == 'success'
def test_benchmark_lookalike_workflow(self):
b = self.batch()
setup_jobs = []
for i in range(10):
j = b.new_job(f'setup_{i}').cpu(0.1)
j.command(f'echo "foo" > {j.ofile}')
setup_jobs.append(j)
jobs = []
for i in range(500):
j = b.new_job(f'create_file_{i}').cpu(0.1)
j.command(
f'echo {setup_jobs[i % len(setup_jobs)].ofile} > {j.ofile}')
j.command(f'echo "bar" >> {j.ofile}')
jobs.append(j)
combine = b.new_job(f'combine_output').cpu(0.1)
for tasks in grouped(arg_max(), jobs):
combine.command(
f'cat {" ".join(shq(j.ofile) for j in jobs)} >> {combine.ofile}'
)
b.write_output(combine.ofile,
f'{self.gcs_output_dir}/pipeline_benchmark_test.txt')
class ServiceTests(unittest.TestCase):
def setUp(self):
self.backend = ServiceBackend()
remote_tmpdir = get_user_config().get('batch', 'remote_tmpdir')
if not remote_tmpdir.endswith('/'):
remote_tmpdir += '/'
self.remote_tmpdir = remote_tmpdir
if remote_tmpdir.startswith('gs://'):
self.bucket = re.fullmatch(
'gs://(?P<bucket_name>[^/]+).*',
remote_tmpdir).groupdict()['bucket_name']
else:
assert remote_tmpdir.startswith('hail-az://')
storage_account, container_name = re.fullmatch(
'hail-az://(?P<storage_account>[^/]+)/(?P<container_name>[^/]+).*',
remote_tmpdir).groups()
self.bucket = f'{storage_account}/{container_name}'
self.cloud_input_dir = f'{self.remote_tmpdir}batch-tests/resources'
token = uuid.uuid4()
self.cloud_output_path = f'/batch-tests/{token}'
self.cloud_output_dir = f'{self.remote_tmpdir}{self.cloud_output_path}'
in_cluster_key_file = '/test-gsa-key/key.json'
if not os.path.exists(in_cluster_key_file):
in_cluster_key_file = None
router_fs = RouterAsyncFS(
'gs',
gcs_kwargs={
'project': 'hail-vdc',
'credentials_file': in_cluster_key_file
},
azure_kwargs={'credential_file': in_cluster_key_file})
def sync_exists(url):
return async_to_blocking(router_fs.exists(url))
def sync_write(url, data):
return async_to_blocking(router_fs.write(url, data))
if not sync_exists(
f'{self.remote_tmpdir}batch-tests/resources/hello.txt'):
sync_write(f'{self.remote_tmpdir}batch-tests/resources/hello.txt',
b'hello world')
if not sync_exists(
f'{self.remote_tmpdir}batch-tests/resources/hello spaces.txt'):
sync_write(
f'{self.remote_tmpdir}batch-tests/resources/hello spaces.txt',
b'hello')
if not sync_exists(
f'{self.remote_tmpdir}batch-tests/resources/hello (foo) spaces.txt'
):
sync_write(
f'{self.remote_tmpdir}batch-tests/resources/hello (foo) spaces.txt',
b'hello')
def tearDown(self):
self.backend.close()
def batch(self,
requester_pays_project=None,
default_python_image=None,
cancel_after_n_failures=None):
return Batch(backend=self.backend,
default_image=DOCKER_ROOT_IMAGE,
attributes={
'foo': 'a',
'bar': 'b'
},
requester_pays_project=requester_pays_project,
default_python_image=default_python_image,
cancel_after_n_failures=cancel_after_n_failures)
def test_single_task_no_io(self):
b = self.batch()
j = b.new_job()
j.command('echo hello')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_task_input(self):
b = self.batch()
input = b.read_input(f'{self.cloud_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_task_input_resource_group(self):
b = self.batch()
input = b.read_input_group(foo=f'{self.cloud_input_dir}/hello.txt')
j = b.new_job()
j.storage('10Gi')
j.command(f'cat {input.foo}')
j.command(f'cat {input}.foo')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_task_output(self):
b = self.batch()
j = b.new_job(attributes={'a': 'bar', 'b': 'foo'})
j.command(f'echo hello > {j.ofile}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_task_write_output(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile,
f'{self.cloud_output_dir}/test_single_task_output.txt')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_task_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_task_write_resource_group(self):
b = self.batch()
j = b.new_job()
j.declare_resource_group(output={'foo': '{root}.foo'})
j.command(f'echo "hello" > {j.output.foo}')
b.write_output(
j.output,
f'{self.cloud_output_dir}/test_single_task_write_resource_group')
b.write_output(
j.output.foo,
f'{self.cloud_output_dir}/test_single_task_write_resource_group_file.txt'
)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_multiple_dependent_tasks(self):
output_file = f'{self.cloud_output_dir}/test_multiple_dependent_tasks.txt'
b = self.batch()
j = b.new_job()
j.command(f'echo "0" >> {j.ofile}')
for i in range(1, 3):
j2 = b.new_job()
j2.command(f'echo "{i}" > {j2.tmp1}')
j2.command(f'cat {j.ofile} {j2.tmp1} > {j2.ofile}')
j = j2
b.write_output(j.ofile, output_file)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_specify_cpu(self):
b = self.batch()
j = b.new_job()
j.cpu('0.5')
j.command(f'echo "hello" > {j.ofile}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_specify_memory(self):
b = self.batch()
j = b.new_job()
j.memory('100M')
j.command(f'echo "hello" > {j.ofile}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_scatter_gather(self):
b = self.batch()
for i in range(3):
j = b.new_job(name=f'foo{i}')
j.command(f'echo "{i}" > {j.ofile}')
merger = b.new_job()
merger.command('cat {files} > {ofile}'.format(files=' '.join([
j.ofile for j in sorted(
b.select_jobs('foo'), key=lambda x: x.name, reverse=True)
]),
ofile=merger.ofile))
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_file_name_space(self):
b = self.batch()
input = b.read_input(f'{self.cloud_input_dir}/hello (foo) spaces.txt')
j = b.new_job()
j.command(f'cat {input} > {j.ofile}')
b.write_output(j.ofile,
f'{self.cloud_output_dir}/hello (foo) spaces.txt')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_dry_run(self):
b = self.batch()
j = b.new_job()
j.command(f'echo hello > {j.ofile}')
b.write_output(j.ofile,
f'{self.cloud_output_dir}/test_single_job_output.txt')
b.run(dry_run=True)
def test_verbose(self):
b = self.batch()
input = b.read_input(f'{self.cloud_input_dir}/hello.txt')
j = b.new_job()
j.command(f'cat {input}')
b.write_output(input, f'{self.cloud_output_dir}/hello.txt')
res = b.run(verbose=True)
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_cloudfuse(self):
assert self.bucket
path = f'/{self.bucket}{self.cloud_output_path}'
b = self.batch()
head = b.new_job()
head.command(f'mkdir -p {path}; echo head > {path}/cloudfuse_test_1')
head.cloudfuse(self.bucket, f'/{self.bucket}', read_only=False)
tail = b.new_job()
tail.command(f'cat {path}/cloudfuse_test_1')
tail.cloudfuse(self.bucket, f'/{self.bucket}', read_only=True)
tail.depends_on(head)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_cloudfuse_read_only(self):
assert self.bucket
path = f'/{self.bucket}{self.cloud_output_path}'
b = self.batch()
j = b.new_job()
j.command(f'mkdir -p {path}; echo head > {path}/cloudfuse_test_1')
j.cloudfuse(self.bucket, f'/{self.bucket}', read_only=True)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'failure', str(
(res_status, res.debug_info()))
def test_cloudfuse_implicit_dirs(self):
assert self.bucket
path = f'/{self.bucket}{self.cloud_output_path}'
b = self.batch()
head = b.new_job()
head.command(
f'mkdir -p {path}/cloudfuse/; echo head > {path}/cloudfuse/data')
head.cloudfuse(self.bucket, f'/{self.bucket}', read_only=False)
tail = b.new_job()
tail.command(f'cat {path}/cloudfuse/data')
tail.cloudfuse(self.bucket, f'/{self.bucket}', read_only=True)
tail.depends_on(head)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_cloudfuse_empty_string_bucket_fails(self):
assert self.bucket
b = self.batch()
j = b.new_job()
with self.assertRaises(BatchException):
j.cloudfuse('', '/empty_bucket')
with self.assertRaises(BatchException):
j.cloudfuse(self.bucket, '')
@skip_in_azure
def test_requester_pays(self):
b = self.batch(requester_pays_project='hail-vdc')
input = b.read_input('gs://hail-services-requester-pays/hello')
j = b.new_job()
j.command(f'cat {input}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_benchmark_lookalike_workflow(self):
b = self.batch()
setup_jobs = []
for i in range(10):
j = b.new_job(f'setup_{i}').cpu(0.25)
j.command(f'echo "foo" > {j.ofile}')
setup_jobs.append(j)
jobs = []
for i in range(500):
j = b.new_job(f'create_file_{i}').cpu(0.25)
j.command(
f'echo {setup_jobs[i % len(setup_jobs)].ofile} > {j.ofile}')
j.command(f'echo "bar" >> {j.ofile}')
jobs.append(j)
combine = b.new_job(f'combine_output').cpu(0.25)
for tasks in grouped(arg_max(), jobs):
combine.command(
f'cat {" ".join(shq(j.ofile) for j in jobs)} >> {combine.ofile}'
)
b.write_output(combine.ofile,
f'{self.cloud_output_dir}/pipeline_benchmark_test.txt')
# too slow
# assert b.run().status()['state'] == 'success'
def test_envvar(self):
b = self.batch()
j = b.new_job()
j.env('SOME_VARIABLE', '123abcdef')
j.command('[ $SOME_VARIABLE = "123abcdef" ]')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_job_with_shell(self):
msg = 'hello world'
b = self.batch()
j = b.new_job(shell='/bin/sh')
j.command(f'echo "{msg}"')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_single_job_with_nonsense_shell(self):
b = self.batch()
j = b.new_job(shell='/bin/ajdsfoijasidojf')
j.command(f'echo "hello"')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'failure', str(
(res_status, res.debug_info()))
def test_single_job_with_intermediate_failure(self):
b = self.batch()
j = b.new_job()
j.command(f'echoddd "hello"')
j2 = b.new_job()
j2.command(f'echo "world"')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'failure', str(
(res_status, res.debug_info()))
def test_input_directory(self):
b = self.batch()
input1 = b.read_input(self.cloud_input_dir)
input2 = b.read_input(self.cloud_input_dir.rstrip('/') + '/')
j = b.new_job()
j.command(f'ls {input1}/hello.txt')
j.command(f'ls {input2}/hello.txt')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
def test_python_job(self):
b = self.batch(default_python_image=PYTHON_DILL_IMAGE)
head = b.new_job()
head.command(f'echo "5" > {head.r5}')
head.command(f'echo "3" > {head.r3}')
def read(path):
with open(path, 'r') as f:
i = f.read()
return int(i)
def multiply(x, y):
return x * y
def reformat(x, y):
return {'x': x, 'y': y}
middle = b.new_python_job()
r3 = middle.call(read, head.r3)
r5 = middle.call(read, head.r5)
r_mult = middle.call(multiply, r3, r5)
middle2 = b.new_python_job()
r_mult = middle2.call(multiply, r_mult, 2)
r_dict = middle2.call(reformat, r3, r5)
tail = b.new_job()
tail.command(
f'cat {r3.as_str()} {r5.as_repr()} {r_mult.as_str()} {r_dict.as_json()}'
)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
assert res.get_job_log(
4)['main'] == "3\n5\n30\n{\"x\": 3, \"y\": 5}\n", str(
res.debug_info())
def test_python_job_w_resource_group_unpack_individually(self):
b = self.batch(default_python_image=PYTHON_DILL_IMAGE)
head = b.new_job()
head.declare_resource_group(count={
'r5': '{root}.r5',
'r3': '{root}.r3'
})
head.command(f'echo "5" > {head.count.r5}')
head.command(f'echo "3" > {head.count.r3}')
def read(path):
with open(path, 'r') as f:
r = int(f.read())
return r
def multiply(x, y):
return x * y
def reformat(x, y):
return {'x': x, 'y': y}
middle = b.new_python_job()
r3 = middle.call(read, head.count.r3)
r5 = middle.call(read, head.count.r5)
r_mult = middle.call(multiply, r3, r5)
middle2 = b.new_python_job()
r_mult = middle2.call(multiply, r_mult, 2)
r_dict = middle2.call(reformat, r3, r5)
tail = b.new_job()
tail.command(
f'cat {r3.as_str()} {r5.as_repr()} {r_mult.as_str()} {r_dict.as_json()}'
)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
assert res.get_job_log(
4)['main'] == "3\n5\n30\n{\"x\": 3, \"y\": 5}\n", str(
res.debug_info())
def test_python_job_w_resource_group_unpack_jointly(self):
b = self.batch(default_python_image=PYTHON_DILL_IMAGE)
head = b.new_job()
head.declare_resource_group(count={
'r5': '{root}.r5',
'r3': '{root}.r3'
})
head.command(f'echo "5" > {head.count.r5}')
head.command(f'echo "3" > {head.count.r3}')
def read_rg(root):
with open(root['r3'], 'r') as f:
r3 = int(f.read())
with open(root['r5'], 'r') as f:
r5 = int(f.read())
return (r3, r5)
def multiply(r):
x, y = r
return x * y
middle = b.new_python_job()
r = middle.call(read_rg, head.count)
r_mult = middle.call(multiply, r)
tail = b.new_job()
tail.command(f'cat {r_mult.as_str()}')
res = b.run()
res_status = res.status()
assert res_status['state'] == 'success', str(
(res_status, res.debug_info()))
job_log_3 = res.get_job_log(3)
assert job_log_3['main'] == "15\n", str((job_log_3, res.debug_info()))
def test_python_job_w_non_zero_ec(self):
b = self.batch(default_python_image=PYTHON_DILL_IMAGE)
j = b.new_python_job()
def error():
raise Exception("this should fail")
j.call(error)
res = b.run()
res_status = res.status()
assert res_status['state'] == 'failure', str(
(res_status, res.debug_info()))
def test_fail_fast(self):
b = self.batch(cancel_after_n_failures=1)
j1 = b.new_job()
j1.command('false')
j2 = b.new_job()
j2.command('sleep 300')
res = b.run()
job_status = res.get_job(2).status()
assert job_status['state'] == 'Cancelled', str(
(job_status, res.debug_info()))
def test_service_backend_remote_tempdir_with_trailing_slash(self):
backend = ServiceBackend(
remote_tmpdir=f'{self.remote_tmpdir}/temporary-files/')
b = Batch(backend=backend)
j1 = b.new_job()
j1.command(f'echo hello > {j1.ofile}')
j2 = b.new_job()
j2.command(f'cat {j1.ofile}')
b.run()
def test_service_backend_remote_tempdir_with_no_trailing_slash(self):
backend = ServiceBackend(
remote_tmpdir=f'{self.remote_tmpdir}/temporary-files')
b = Batch(backend=backend)
j1 = b.new_job()
j1.command(f'echo hello > {j1.ofile}')
j2 = b.new_job()
j2.command(f'cat {j1.ofile}')
b.run()
def test_large_command(self):
backend = ServiceBackend(
remote_tmpdir=f'{self.remote_tmpdir}/temporary-files')
b = Batch(backend=backend)
j1 = b.new_job()
long_str = secrets.token_urlsafe(15 * 1024)
j1.command(f'echo "{long_str}"')
b.run()
def test_big_batch_which_uses_slow_path(self):
backend = ServiceBackend(
remote_tmpdir=f'{self.remote_tmpdir}/temporary-files')
b = Batch(backend=backend)
# 8 * 256 * 1024 = 2 MiB > 1 MiB max bunch size
for i in range(8):
j1 = b.new_job()
long_str = secrets.token_urlsafe(256 * 1024)
j1.command(f'echo "{long_str}"')
batch = b.run()
assert not batch.submission_info.used_fast_create
batch_status = batch.status()
assert batch_status['state'] == 'success', str((batch.debug_info()))