[batch] Clean up environment variables for easier local execution of …

…batch tests (#12862)

Trying to make it more ergonomic to simply do `python3 -m pytest
batch/test/test_batch.py::test_job` (now works without any extra
environment variables or configuration).

This involved the following changes:
- Deleted of some env vars that are no longer used / can be easily
consolidated into existing ones
- Gave defaults to those testing env variables for which there are
reasonable defaults. E.g. `DOCKER_ROOT_IMAGE` and
`HAIL_GENETICS_HAIL_IMAGE`.
- Pushed other environment variables for which there are not reasonable
defaults into the tests that need them. If you run a test that requires
`HAIL_CLOUD`, you'll still get an error that that env variable is unset
and you should set it. But, if you just want to run a single test that
doesn't need `HAIL_CLOUD` it won't get in the way.

batch/test/pytest.ini

@@ -0,0 +1,10 @@
+[pytest]
+asyncio_mode = auto
+addopts = --strict-markers
+markers =
+    unchecked_allocator: tests that use the unchecked allocator
+    asyncio: test files that use asyncio
+filterwarnings =
+    error
+    ignore::UserWarning
+    ignore::DeprecationWarning

batch/test/test_accounts.py

@@ -14,11 +14,10 @@
 from hailtop.utils.rich_progress_bar import BatchProgressBar
 
 from .billing_projects import get_billing_project_prefix
+from .utils import DOCKER_ROOT_IMAGE
 
 pytestmark = pytest.mark.asyncio
 
-DOCKER_ROOT_IMAGE = os.environ['DOCKER_ROOT_IMAGE']
-
 
 @pytest.fixture
 async def make_client() -> AsyncGenerator[Callable[[str], Awaitable[BatchClient]], Any]:

batch/test/test_aioclient.py

@@ -1,10 +1,8 @@
-import os
-
 import pytest
 
 from hailtop.batch_client.aioclient import BatchClient
 
-DOCKER_ROOT_IMAGE = os.environ['DOCKER_ROOT_IMAGE']
+from .utils import DOCKER_ROOT_IMAGE
 
 pytestmark = pytest.mark.asyncio
 

batch/test/test_batch.py

@@ -15,22 +15,10 @@
 from hailtop.utils import external_requests_client_session, retry_response_returning_functions, sync_sleep_and_backoff
 
 from .failure_injecting_client_session import FailureInjectingClientSession
-from .utils import legacy_batch_status, smallest_machine_type
+from .utils import DOCKER_ROOT_IMAGE, HAIL_GENETICS_HAIL_IMAGE, legacy_batch_status, smallest_machine_type
 
 deploy_config = get_deploy_config()
 
-DOCKER_PREFIX = os.environ['DOCKER_PREFIX']
-DOCKER_ROOT_IMAGE = os.environ['DOCKER_ROOT_IMAGE']
-UBUNTU_IMAGE = 'ubuntu:20.04'
-
-
-DOMAIN = os.environ.get('HAIL_DOMAIN')
-NAMESPACE = os.environ.get('HAIL_DEFAULT_NAMESPACE')
-SCOPE = os.environ.get('HAIL_SCOPE', 'test')
-CLOUD = os.environ.get('HAIL_CLOUD')
-
-HAIL_GENETICS_HAIL_IMAGE = os.environ['HAIL_GENETICS_HAIL_IMAGE']
-
 
 @pytest.fixture
 def client():
@@ -151,7 +139,7 @@ def test_invalid_resource_requests(client: BatchClient):
         bb.submit()
 
     bb = client.create_batch()
-    resources = {'storage': '10000000Gi', 'machine_type': smallest_machine_type(CLOUD)}
+    resources = {'storage': '10000000Gi', 'machine_type': smallest_machine_type()}
     bb.create_job(DOCKER_ROOT_IMAGE, ['true'], resources=resources)
     with pytest.raises(httpx.ClientResponseError, match='resource requests.*unsatisfiable'):
         bb.submit()
@@ -222,7 +210,7 @@ def test_quota_shared_by_io_and_rootfs(client: BatchClient):
 def test_nonzero_storage(client: BatchClient):
     bb = client.create_batch()
     resources = {'cpu': '0.25', 'memory': '10M', 'storage': '20Gi'}
-    j = bb.create_job(UBUNTU_IMAGE, ['/bin/sh', '-c', 'true'], resources=resources)
+    j = bb.create_job(DOCKER_ROOT_IMAGE, ['/bin/sh', '-c', 'true'], resources=resources)
     b = bb.submit()
     status = j.wait()
     assert status['state'] == 'Success', str((status, b.debug_info()))
@@ -232,7 +220,7 @@ def test_nonzero_storage(client: BatchClient):
 def test_attached_disk(client: BatchClient):
     bb = client.create_batch()
     resources = {'cpu': '0.25', 'memory': '10M', 'storage': '400Gi'}
-    j = bb.create_job(UBUNTU_IMAGE, ['/bin/sh', '-c', 'df -h; fallocate -l 390GiB /io/foo'], resources=resources)
+    j = bb.create_job(DOCKER_ROOT_IMAGE, ['/bin/sh', '-c', 'df -h; fallocate -l 390GiB /io/foo'], resources=resources)
     b = bb.submit()
     status = j.wait()
     assert status['state'] == 'Success', str((status, b.debug_info()))
@@ -359,6 +347,7 @@ def test_fail(client: BatchClient):
 
 def test_unknown_image(client: BatchClient):
     bb = client.create_batch()
+    DOCKER_PREFIX = os.environ['DOCKER_PREFIX']
     j = bb.create_job(f'{DOCKER_PREFIX}/does-not-exist', ['echo', 'test'])
     b = bb.submit()
     status = j.wait()
@@ -607,7 +596,8 @@ def test_cloud_image(client: BatchClient):
     assert status['state'] == 'Success', str((status, b.debug_info()))
 
 
-def test_service_account(client: BatchClient):
+def test_k8s_service_account(client: BatchClient):
+    NAMESPACE = os.environ['HAIL_DEFAULT_NAMESPACE']
     bb = client.create_batch()
     j = bb.create_job(
         os.environ['CI_UTILS_IMAGE'],
@@ -785,6 +775,9 @@ def test_submit_batch_in_job(client: BatchClient):
 
 
 def test_cant_submit_to_default_with_other_ns_creds(client: BatchClient):
+    DOMAIN = os.environ['HAIL_DOMAIN']
+    NAMESPACE = os.environ['HAIL_DEFAULT_NAMESPACE']
+
     remote_tmpdir = get_user_config().get('batch', 'remote_tmpdir')
     script = f'''import hailtop.batch as hb
 backend = hb.ServiceBackend("test", remote_tmpdir="{remote_tmpdir}")
@@ -862,7 +855,7 @@ def test_cannot_contact_other_internal_ips(client: BatchClient):
 
 
 @skip_in_azure
-def test_can_use_google_credentials(client: BatchClient):
+def test_hadoop_can_use_cloud_credentials(client: BatchClient):
     token = os.environ["HAIL_TOKEN"]
     remote_tmpdir = get_user_config().get('batch', 'remote_tmpdir')
     bb = client.create_batch()
@@ -901,7 +894,7 @@ def test_can_use_google_credentials(client: BatchClient):
 def test_user_authentication_within_job(client: BatchClient):
     bb = client.create_batch()
     cmd = ['bash', '-c', 'hailctl auth user']
-    no_token = bb.create_job(os.environ['CI_UTILS_IMAGE'], cmd, mount_tokens=False)
+    no_token = bb.create_job(HAIL_GENETICS_HAILTOP_IMAGE, cmd, mount_tokens=False)
     b = bb.submit()
 
     no_token_status = no_token.wait()
@@ -1046,7 +1039,7 @@ def test_pool_standard_instance_cheapest(client: BatchClient):
 
 def test_job_private_instance_preemptible(client: BatchClient):
     bb = client.create_batch()
-    resources = {'machine_type': smallest_machine_type(CLOUD)}
+    resources = {'machine_type': smallest_machine_type()}
     j = bb.create_job(DOCKER_ROOT_IMAGE, ['true'], resources=resources)
     b = bb.submit()
     status = j.wait()
@@ -1056,7 +1049,7 @@ def test_job_private_instance_preemptible(client: BatchClient):
 
 def test_job_private_instance_nonpreemptible(client: BatchClient):
     bb = client.create_batch()
-    resources = {'machine_type': smallest_machine_type(CLOUD), 'preemptible': False}
+    resources = {'machine_type': smallest_machine_type(), 'preemptible': False}
     j = bb.create_job(DOCKER_ROOT_IMAGE, ['true'], resources=resources)
     b = bb.submit()
     status = j.wait()
@@ -1066,7 +1059,7 @@ def test_job_private_instance_nonpreemptible(client: BatchClient):
 
 def test_job_private_instance_cancel(client: BatchClient):
     bb = client.create_batch()
-    resources = {'machine_type': smallest_machine_type(CLOUD)}
+    resources = {'machine_type': smallest_machine_type()}
     j = bb.create_job(DOCKER_ROOT_IMAGE, ['true'], resources=resources)
     b = bb.submit()
 
@@ -1321,6 +1314,8 @@ def test_submit_update_to_deleted_batch(client: BatchClient):
 
 
 def test_region(client: BatchClient):
+    CLOUD = os.environ['HAIL_CLOUD']
+
     bb = client.create_batch()
     if CLOUD == 'gcp':
         region = 'us-east1'

batch/test/test_dag.py

@@ -10,9 +10,7 @@
 from hailtop.batch_client.client import BatchClient, Job
 from hailtop.config import get_user_config
 
-from .utils import batch_status_job_counter, legacy_batch_status
-
-DOCKER_ROOT_IMAGE = os.environ['DOCKER_ROOT_IMAGE']
+from .utils import DOCKER_ROOT_IMAGE, batch_status_job_counter, legacy_batch_status
 
 
 @pytest.fixture

batch/test/utils.py

@@ -1,3 +1,11 @@
+import os
+
+from hailtop import pip_version
+
+DOCKER_ROOT_IMAGE = os.environ.get('DOCKER_ROOT_IMAGE', 'ubuntu:20.04')
+HAIL_GENETICS_HAIL_IMAGE = os.environ.get('HAIL_GENETICS_HAIL_IMAGE', f'hailgenetics/hail:{pip_version()}')
+
+
 def batch_status_job_counter(batch_status, job_state):
     return len([j for j in batch_status['jobs'] if j['state'] == job_state])
 
@@ -8,7 +16,9 @@ def legacy_batch_status(batch):
     return status
 
 
-def smallest_machine_type(cloud):
+def smallest_machine_type():
+    cloud = os.environ['HAIL_CLOUD']
+
     if cloud == 'gcp':
         return 'n1-standard-1'
     assert cloud == 'azure'

build.yaml

@@ -2447,7 +2447,6 @@ steps:
       export HAIL_TEST_TOKEN_FILE=/user-tokens/tokens.json
       export HAIL_TEST_DEV_TOKEN_FILE=/dev-tokens/tokens.json
       export HAIL_TOKEN="{{ token }}"
-      export HAIL_SCOPE="{{ scope }}"
       export HAIL_CLOUD="{{ global.cloud }}"
       export HAIL_DOMAIN="{{ global.domain }}"
       hailctl config set batch/remote_tmpdir {{ global.test_storage_uri }}/test_batch/{{ token }}/
@@ -2709,7 +2708,6 @@ steps:
       set -ex
       export HAIL_CLOUD={{ global.cloud }}
       export GOOGLE_APPLICATION_CREDENTIALS=/test-gsa-key/key.json
-      export DOCKER_PREFIX="{{ global.docker_prefix }}"
       export DOCKER_ROOT_IMAGE="{{ global.docker_root_image }}"
       export HAIL_GENETICS_HAIL_IMAGE="{{ hailgenetics_hail_image.image }}"
       export HAIL_GENETICS_HAILTOP_IMAGE="{{ hailgenetics_hailtop_image.image }}"

hail/python/test/hailtop/batch/test_batch.py

@@ -9,6 +9,7 @@
 import uuid
 import re
 
+from hailtop import pip_version
 from hailtop.batch import Batch, ServiceBackend, LocalBackend
 from hailtop.batch.exceptions import BatchException
 from hailtop.batch.globals import arg_max
@@ -21,7 +22,7 @@
 
 DOCKER_ROOT_IMAGE = os.environ.get('DOCKER_ROOT_IMAGE', 'ubuntu:20.04')
 PYTHON_DILL_IMAGE = 'hailgenetics/python-dill:3.7-slim'
-HAIL_GENETICS_HAIL_IMAGE = os.environ.get('HAIL_GENETICS_HAIL_IMAGE')
+HAIL_GENETICS_HAIL_IMAGE = os.environ.get('HAIL_GENETICS_HAIL_IMAGE', f'hailgenetics/hail:{pip_version()}')
 
 
 class LocalTests(unittest.TestCase):

hail/python/test/hailtop/batch/test_batch_pool_executor.py

@@ -1,9 +1,7 @@
 import asyncio
 import concurrent.futures
 import time
-import os
 
-import hailtop.batch_client.client
 import pytest
 
 from hailtop.batch import BatchPoolExecutor, ServiceBackend