def test_graceful_decommission_timeout(gcloud_run, workers_arg):
cli.main([
"modify", "test-cluster", workers_arg,
"--graceful-decommission-timeout=1h"
])
assert workers_arg in gcloud_run.call_args[0][0]
assert "--graceful-decommission-timeout=1h" in gcloud_run.call_args[0][0]
def main():
check_for_update()
if len(sys.argv) == 1:
print_help()
sys.exit(0)
else:
module = sys.argv[1]
args = sys.argv[2:]
if module == 'version':
print_version()
elif module == 'dataproc':
from hailtop.hailctl.dataproc import cli # pylint: disable=import-outside-toplevel
cli.main(args)
elif module == 'auth':
from hailtop.hailctl.auth import cli # pylint: disable=import-outside-toplevel
cli.main(args)
elif module == 'dev':
from hailtop.hailctl.dev import cli # pylint: disable=import-outside-toplevel
cli.main(args)
elif module == 'batch':
from hailtop.hailctl.batch import cli # pylint: disable=import-outside-toplevel
cli.main(args)
elif module == 'curl':
from hailtop.hailctl.curl import main # pylint: disable=import-outside-toplevel
main(args)
elif module == 'config':
from hailtop.hailctl.config import cli # pylint: disable=import-outside-toplevel
cli.main(args)
elif module in ('-h', '--help', 'help'):
print_help()
else:
sys.stderr.write(f"ERROR: no such module: {module!r}")
print_help()
sys.exit(1)
def test_connect_zone_required(gcloud_run, gcloud_config):
gcloud_config["compute/zone"] = None
with pytest.raises(Exception):
cli.main(["connect", "test-cluster", "notebook"])
assert gcloud_run.call_count == 0
def test_required_gcloud_version_met(monkeypatch):
monkeypatch.setattr("hailtop.hailctl.dataproc.gcloud.get_version", Mock(return_value=cli.MINIMUM_REQUIRED_GCLOUD_VERSION))
mock_list = Mock()
monkeypatch.setattr(list_clusters, "main", mock_list)
cli.main(["list"])
assert mock_list.called
def test_unable_to_determine_version(monkeypatch):
monkeypatch.setattr("hailtop.hailctl.dataproc.gcloud.get_version", Mock(side_effect=ValueError))
mock_list = Mock()
monkeypatch.setattr(list_clusters, "main", mock_list)
cli.main(["list"])
assert mock_list.called
def test_hailctl_chrome(subprocess, monkeypatch):
monkeypatch.setattr("hailtop.hailctl.dataproc.connect.get_chrome_path",
Mock(side_effect=Exception("Unable to find chrome")))
monkeypatch.setenv("HAILCTL_CHROME", "/path/to/chrome.exe")
cli.main(["connect", "test-cluster", "notebook"])
popen_args = subprocess.Popen.call_args[0][0]
assert popen_args[0] == "/path/to/chrome.exe"
def test_creator_label(gcloud_run, gcloud_config):
gcloud_config["account"] = "*****@*****.**"
cli.main(["start", "my-cluster"])
assert "--labels=creator=test-user_hail_is" in gcloud_run.call_args[0][0]
gcloud_config["account"] = None
cli.main(["start", "my-cluster"])
assert not any(arg.startswith("--labels=") and "creator=" in arg for arg in gcloud_run.call_args[0][0])
def test_master_tags(gcloud_run):
cli.main(["start", "test-cluster", "--master-tags=foo"])
assert gcloud_run.call_count == 2
assert gcloud_run.call_args_list[0][0][0][:4] == [
"dataproc", "clusters", "create", "test-cluster"
]
assert gcloud_run.call_args_list[1][0][0] == [
"compute", "instances", "add-tags", "test-cluster-m", "--tags", "foo"
]
def test_modify_wheel_remote_wheel(gcloud_run):
cli.main(["modify", "test-cluster", "--wheel=gs://some-bucket/hail.whl"])
assert gcloud_run.call_count == 1
gcloud_args = gcloud_run.call_args[0][0]
assert gcloud_args[:3] == ["compute", "ssh", "test-cluster-m"]
remote_command = gcloud_args[gcloud_args.index("--") + 1]
assert remote_command == ("sudo gsutil cp gs://some-bucket/hail.whl /tmp/ && " +
"sudo /opt/conda/default/bin/pip uninstall -y hail && " +
"sudo /opt/conda/default/bin/pip install --no-dependencies /tmp/hail.whl")
def test_wheel_and_update_hail_version_mutually_exclusive(gcloud_run, capsys):
with pytest.raises(SystemExit):
cli.main([
"modify", "test-cluster", "--wheel=./hail.whl",
"--update-hail-version"
])
assert gcloud_run.call_count == 0
assert "argument --update-hail-version: not allowed with argument --wheel" in capsys.readouterr(
).err
def test_cluster_and_service_required(gcloud_run):
with pytest.raises(SystemExit):
cli.main(["connect"])
assert gcloud_run.call_count == 0
with pytest.raises(SystemExit):
cli.main(["connect", "notebook"])
assert gcloud_run.call_count == 0
def test_cluster_and_script_required(gcloud_run):
with pytest.raises(SystemExit):
cli.main(["submit"])
assert gcloud_run.call_count == 0
with pytest.raises(SystemExit):
cli.main(["submit", "test-cluster"])
assert gcloud_run.call_count == 0
def test_requester_pays_project_configuration(gcloud_run, gcloud_config, requester_pays_arg):
gcloud_config["project"] = "foo-project"
cli.main(["start", "test-cluster", requester_pays_arg])
properties = next(arg for arg in gcloud_run.call_args[0][0] if arg.startswith("--properties="))
assert "spark:spark.hadoop.fs.gs.requester.pays.project.id=foo-project" in properties
cli.main(["start", "--project=bar-project", "test-cluster", requester_pays_arg])
properties = next(arg for arg in gcloud_run.call_args[0][0] if arg.startswith("--properties="))
assert "spark:spark.hadoop.fs.gs.requester.pays.project.id=bar-project" in properties
def test_required_gcloud_version_unmet(monkeypatch, capsys):
monkeypatch.setattr("hailtop.hailctl.dataproc.gcloud.get_version", Mock(return_value=(200, 0, 0)))
mock_list = Mock()
monkeypatch.setattr(list_clusters, "main", mock_list)
with pytest.raises(SystemExit):
cli.main(["list"])
assert "hailctl dataproc requires Google Cloud SDK (gcloud) version" in capsys.readouterr().err
assert not mock_list.called
def test_modify_wheel_remote_wheel(gcloud_run):
cli.main(["modify", "test-cluster", "--wheel=gs://some-bucket/hail.whl"])
assert gcloud_run.call_count == 1
gcloud_args = gcloud_run.call_args[0][0]
assert gcloud_args[:3] == ["compute", "ssh", "test-cluster-m"]
remote_command = gcloud_args[gcloud_args.index("--") + 1]
assert remote_command == (
"sudo gsutil cp gs://some-bucket/hail.whl /tmp/ && " +
"sudo /opt/conda/default/bin/pip uninstall -y hail && " +
"sudo /opt/conda/default/bin/pip install --no-dependencies /tmp/hail.whl && "
+ "unzip /tmp/hail.whl && " +
"grep 'Requires-Dist: ' hail*dist-info/METADATA | sed 's/Requires-Dist: //' | sed 's/ (//' | sed 's/)//' | grep -v 'pyspark' | xargs /opt/conda/default/bin/pip install"
)
def test_modify_wheel_local_wheel(gcloud_run):
cli.main(["modify", "test-cluster", "--wheel=./local-hail.whl"])
assert gcloud_run.call_count == 2
copy_gcloud_args = gcloud_run.call_args_list[0][0][0]
assert copy_gcloud_args[:2] == ["compute", "scp"]
assert copy_gcloud_args[-2:] == ["./local-hail.whl", "test-cluster-m:/tmp/"]
install_gcloud_args = gcloud_run.call_args_list[1][0][0]
assert install_gcloud_args[:3] == ["compute", "ssh", "test-cluster-m"]
remote_command = install_gcloud_args[install_gcloud_args.index("--") + 1]
assert remote_command == ("sudo /opt/conda/default/bin/pip uninstall -y hail && " +
"sudo /opt/conda/default/bin/pip install --no-dependencies /tmp/local-hail.whl")
def test_update_hail_version(gcloud_run, monkeypatch, deploy_metadata):
monkeypatch.setattr("hailtop.hailctl.dataproc.modify.get_deploy_metadata", lambda: deploy_metadata)
cli.main(["modify", "test-cluster", "--update-hail-version"])
assert gcloud_run.call_count == 1
gcloud_args = gcloud_run.call_args[0][0]
assert gcloud_args[:3] == ["compute", "ssh", "test-cluster-m"]
remote_command = gcloud_args[gcloud_args.index("--") + 1]
assert remote_command == (
"sudo gsutil cp gs://hail-common/hailctl/dataproc/test-version/hail-test-version-py3-none-any.whl /tmp/ && " +
"sudo /opt/conda/default/bin/pip uninstall -y hail && " +
"sudo /opt/conda/default/bin/pip install --no-dependencies /tmp/hail-test-version-py3-none-any.whl"
)
def main():
check_for_update()
if len(sys.argv) == 1:
print_help()
sys.exit(0)
else:
module = sys.argv[1]
args = sys.argv[2:]
if module == 'version':
print_version()
elif module == 'dataproc':
from hailtop.hailctl.dataproc import cli
cli.main(args)
elif module == 'auth':
from hailtop.hailctl.auth import cli
cli.main(args)
elif module == 'dev':
from hailtop.hailctl.dev import cli
cli.main(args)
elif module == 'batch':
from hailtop.hailctl.batch import cli
cli.main(args)
elif module in ('-h', '--help', 'help'):
print_help()
else:
sys.stderr.write(f"ERROR: no such module: {module!r}")
print_help()
sys.exit(1)
def test_connect(gcloud_run, subprocess):
cli.main(["connect", "test-cluster", "notebook"])
gcloud_args = gcloud_run.call_args[0][0]
assert gcloud_args[:2] == ["compute", "ssh"]
assert gcloud_args[2][(gcloud_args[2].find("@") + 1):] == "test-cluster-m"
assert "--ssh-flag=-D 10000" in gcloud_args
assert "--ssh-flag=-N" in gcloud_args
assert "--ssh-flag=-f" in gcloud_args
assert "--ssh-flag=-n" in gcloud_args
popen_args = subprocess.Popen.call_args[0][0]
assert popen_args[0] == "chromium"
assert popen_args[1].startswith("http://localhost")
assert "--proxy-server=socks5://localhost:10000" in popen_args
assert "--host-resolver-rules=MAP * 0.0.0.0 , EXCLUDE localhost" in popen_args
assert "--proxy-bypass-list=<-loopback>" in popen_args
assert any(arg.startswith("--user-data-dir=") for arg in popen_args)
def test_modify_wheel_local_wheel(gcloud_run):
cli.main(["modify", "test-cluster", "--wheel=./local-hail.whl"])
assert gcloud_run.call_count == 2
copy_gcloud_args = gcloud_run.call_args_list[0][0][0]
assert copy_gcloud_args[:2] == ["compute", "scp"]
assert copy_gcloud_args[-2:] == [
"./local-hail.whl", "test-cluster-m:/tmp/"
]
install_gcloud_args = gcloud_run.call_args_list[1][0][0]
assert install_gcloud_args[:3] == ["compute", "ssh", "test-cluster-m"]
remote_command = install_gcloud_args[install_gcloud_args.index("--") + 1]
assert remote_command == (
"sudo /opt/conda/default/bin/pip uninstall -y hail && " +
"sudo /opt/conda/default/bin/pip install --no-dependencies /tmp/local-hail.whl && "
+ "unzip /tmp/local-hail.whl && " +
"grep 'Requires-Dist: ' hail*dist-info/METADATA | sed 's/Requires-Dist: //' | sed 's/ (//' | sed 's/)//' | grep -v 'pyspark' | xargs /opt/conda/default/bin/pip install"
)
def test_vep_defaults_to_larger_worker_boot_disk(gcloud_run):
cli.main(["start", "test-cluster", "--vep=GRCh37"])
assert "--worker-boot-disk-size=200GB" in gcloud_run.call_args[0][0]
assert "--secondary-worker-boot-disk-size=200GB" in gcloud_run.call_args[
0][0]
def test_vep_defaults_to_highmem_master_machine(gcloud_run):
cli.main(["start", "test-cluster", "--vep=GRCh37"])
assert "--master-machine-type=n1-highmem-8" in gcloud_run.call_args[0][0]
def test_boot_disk_size_configuration(gcloud_run, machine_arg):
cli.main(["start", machine_arg, "test-cluster"])
assert f"{machine_arg}GB" in gcloud_run.call_args[0][0]
def test_machine_type_configuration(gcloud_run, machine_arg):
cli.main(["start", machine_arg, "test-cluster"])
assert machine_arg in gcloud_run.call_args[0][0]
def test_secondary_workers_configuration(gcloud_run, workers_arg):
cli.main(["start", workers_arg, "test-cluster"])
assert "--num-secondary-workers=8" in gcloud_run.call_args[0][0]
def test_cluster_name_required(capsys, gcloud_run):
with pytest.raises(SystemExit):
cli.main(["start"])
assert "arguments are required: name" in capsys.readouterr().err
assert gcloud_run.call_count == 0
def test_workers_configuration(gcloud_run):
cli.main(["start", "--num-workers=4", "test-cluster"])
assert "--num-workers=4" in gcloud_run.call_args[0][0]
def test_cluster_location(gcloud_run, location_arg):
cli.main(["start", location_arg, "test-cluster"])
assert location_arg in gcloud_run.call_args[0][0]
def test_cluster_project(gcloud_run):
cli.main(["start", "--project", "foo", "test-cluster"])
assert "--project=foo" in gcloud_run.call_args[0][0]
def test_master_tags_dry_run(gcloud_run):
cli.main(["start", "test-cluster", "--master-tags=foo", "--dry-run"])
assert gcloud_run.call_count == 0
