-
Notifications
You must be signed in to change notification settings - Fork 244
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[hailctl] Add hailctl batch submit #12471
Changes from all commits
84a06fb
42409c7
4e4bfc0
1185df3
301b9f2
5b140a3
e37b6ac
402855f
7671054
ed47258
70a0b06
434ef9f
f44f783
b27465d
db8a50c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,74 @@ | ||
import asyncio | ||
import orjson | ||
import os | ||
|
||
import hailtop.batch as hb | ||
import hailtop.batch_client.client as bc | ||
from hailtop import pip_version | ||
from hailtop.aiotools.copy import copy_from_dict | ||
from hailtop.config import get_remote_tmpdir, get_user_config_path, get_deploy_config | ||
from hailtop.utils import secret_alnum_string, unpack_comma_delimited_inputs | ||
|
||
HAIL_GENETICS_HAIL_IMAGE = os.environ.get('HAIL_GENETICS_HAIL_IMAGE', f'hailgenetics/hail:{pip_version()}') | ||
|
||
|
||
def init_parser(parser): | ||
parser.add_argument('script', type=str, help='Path to script') | ||
parser.add_argument('--name', type=str, default='', help='Batch name') | ||
parser.add_argument('--image-name', type=str, required=False, | ||
help='Name for Docker image. Defaults to hailgenetics/hail') | ||
parser.add_argument('--files', nargs='+', action='append', default=[], | ||
help='Comma-separated list of files or directories to add to the working directory of job') | ||
parser.add_argument('-o', type=str, default='text', choices=['text', 'json']) | ||
|
||
|
||
async def async_main(args): | ||
script = args.script | ||
files = unpack_comma_delimited_inputs(args.files) | ||
user_config = get_user_config_path() | ||
quiet = args.o != 'text' | ||
|
||
remote_tmpdir = get_remote_tmpdir('hailctl batch submit') | ||
tmpdir_path_prefix = secret_alnum_string() | ||
|
||
def cloud_prefix(path): | ||
return f'{remote_tmpdir}/{tmpdir_path_prefix}/{path}' | ||
|
||
b = hb.Batch(name=args.name, backend=hb.ServiceBackend()) | ||
j = b.new_bash_job() | ||
j.image(args.image_name or HAIL_GENETICS_HAIL_IMAGE) | ||
|
||
rel_file_paths = [os.path.relpath(file) for file in files] | ||
local_files_to_cloud_files = [{'from': local, 'to': cloud_prefix(local)} for local in rel_file_paths] | ||
await copy_from_dict(files=[ | ||
{'from': script, 'to': cloud_prefix(script)}, | ||
{'from': str(user_config), 'to': cloud_prefix(user_config)}, | ||
*local_files_to_cloud_files, | ||
]) | ||
for file in local_files_to_cloud_files: | ||
local_file = file['from'] | ||
cloud_file = file['to'] | ||
in_file = b.read_input(cloud_file) | ||
j.command(f'ln -s {in_file} {local_file}') | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. You might need to defend against users mistakenly providing fully qualified paths to There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Changed to use the relative path from the user's current directory, so as to be compatible with scripts reading from relative paths. |
||
|
||
script_file = b.read_input(cloud_prefix(script)) | ||
config_file = b.read_input(cloud_prefix(user_config)) | ||
j.command(f'mkdir -p $HOME/.config/hail && ln -s {config_file} $HOME/.config/hail/config.ini') | ||
|
||
j.env('HAIL_QUERY_BACKEND', 'batch') | ||
|
||
command = 'python3' if script.endswith('.py') else 'bash' | ||
j.command(f'{command} {script_file}') | ||
batch_handle: bc.Batch = b.run(wait=False, disable_progress_bar=quiet) # type: ignore | ||
|
||
if args.o == 'text': | ||
deploy_config = get_deploy_config() | ||
url = deploy_config.external_url('batch', f'/batches/{batch_handle.id}/jobs/1') | ||
print(f'Submitted batch {batch_handle.id}, see {url}') | ||
else: | ||
assert args.o == 'json' | ||
print(orjson.dumps({'id': batch_handle.id}).decode('utf-8')) | ||
|
||
|
||
def main(args, pass_through_args, client): # pylint: disable=unused-argument | ||
asyncio.run(async_main(args)) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Ah this is so delightfully short!