Hail 118

Makefile

@@ -81,6 +81,19 @@ check-pip-requirements:
 		ci \
 		memory
 
+.PHONY: check-linux-pip-requirements
+check-linux-pip-requirements:
+	./check_linux_pip_requirements.sh \
+		hail/python/hailtop \
+		hail/python \
+		hail/python/dev \
+		gear \
+		web_common \
+		auth \
+		batch \
+		ci \
+		memory
+
 .PHONY: install-dev-requirements
 install-dev-requirements:
 	python3 -m pip install \
@@ -146,6 +159,13 @@ base-image: hail-ubuntu-image docker/Dockerfile.base
 	./docker-build.sh . docker/Dockerfile.base.out $(BASE_IMAGE)
 	echo $(BASE_IMAGE) > $@
 
+hail-run-image: base-image hail/Dockerfile.hail-run hail/python/pinned-requirements.txt hail/python/dev/pinned-requirements.txt docker/core-site.xml
+	$(eval BASE_IMAGE := $(DOCKER_PREFIX)/hail-run:$(TOKEN))
+	$(MAKE) -C hail wheel
+	python3 ci/jinja2_render.py '{"base_image":{"image":"'$$(cat base-image)'"}}' hail/Dockerfile.hail-run hail/Dockerfile.hail-run.out
+	./docker-build.sh . hail/Dockerfile.hail-run.out $(BASE_IMAGE)
+	echo $(BASE_IMAGE) > $@
+
 private-repo-hailgenetics-hail-image: hail-ubuntu-image docker/hailgenetics/hail/Dockerfile $(shell git ls-files hail/src/main hail/python)
 	$(eval PRIVATE_REPO_HAILGENETICS_HAIL_IMAGE := $(DOCKER_PREFIX)/hailgenetics/hail:$(TOKEN))
 	$(MAKE) -C hail wheel

auth/auth/auth.py

@@ -513,12 +513,12 @@ async def post_create_user(request, userdata):  # pylint: disable=unused-argumen
 @auth.rest_authenticated_developers_only
 async def rest_get_users(request, userdata):  # pylint: disable=unused-argument
     db: Database = request.app['db']
-    users = await db.select_and_fetchall(
-        '''
-SELECT id, username, login_id, state, is_developer, is_service_account FROM users;
+    _query = '''
+SELECT id, username, login_id, state, is_developer, is_service_account, hail_identity
+FROM users;
 '''
-    )
-    return json_response([user async for user in users])
+    users = [x async for x in db.select_and_fetchall(_query)]
+    return json_response(users)
 
 
 @routes.get('/api/v1alpha/users/{user}')
@@ -529,7 +529,7 @@ async def rest_get_user(request, userdata):  # pylint: disable=unused-argument
 
     user = await db.select_and_fetchone(
         '''
-SELECT id, username, login_id, state, is_developer, is_service_account FROM users
+SELECT id, username, login_id, state, is_developer, is_service_account, hail_identity FROM users
 WHERE username = %s;
 ''',
         (username,),

auth/deployment.yaml

@@ -126,11 +126,7 @@ spec:
   selector:
     matchLabels:
       app: auth
-{% if deploy %}
-  replicas: 3
-{% else %}
-  replicas: 1
-{% endif %}
+  replicas: 5
   template:
     metadata:
       labels:
@@ -253,6 +249,7 @@ spec:
          secret:
            optional: false
            secretName: ssl-config-auth
+{% if deploy %}
 ---
 apiVersion: autoscaling/v2beta1
 kind: HorizontalPodAutoscaler
@@ -263,13 +260,8 @@ spec:
     apiVersion: apps/v1
     kind: Deployment
     name: auth
-{% if deploy %}
   minReplicas: 3
   maxReplicas: 10
-{% else %}
-  minReplicas: 1
-  maxReplicas: 3
-{% endif %}
   metrics:
    - type: Resource
      resource:
@@ -281,14 +273,11 @@ kind: PodDisruptionBudget
 metadata:
   name: auth
 spec:
-{% if deploy %}
   minAvailable: 2
-{% else %}
-  minAvailable: 0
-{% endif %}
   selector:
     matchLabels:
       app: auth
+{% endif %}
 ---
 apiVersion: v1
 kind: Service

auth/pinned-requirements.txt

@@ -4,12 +4,12 @@
 #
 #    pip-compile --output-file=hail/auth/pinned-requirements.txt hail/auth/requirements.txt
 #
-cachetools==5.3.0
+cachetools==5.3.1
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
     #   -c hail/auth/../hail/python/pinned-requirements.txt
     #   google-auth
-certifi==2022.12.7
+certifi==2023.5.7
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
     #   -c hail/auth/../hail/python/dev/pinned-requirements.txt
@@ -22,7 +22,7 @@ charset-normalizer==3.1.0
     #   -c hail/auth/../hail/python/pinned-requirements.txt
     #   -c hail/auth/../web_common/pinned-requirements.txt
     #   requests
-google-auth==2.17.3
+google-auth==2.19.1
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
     #   -c hail/auth/../hail/python/pinned-requirements.txt
@@ -37,7 +37,9 @@ idna==3.4
     #   -c hail/auth/../web_common/pinned-requirements.txt
     #   requests
 oauthlib==3.2.2
-    # via requests-oauthlib
+    # via
+    #   -c hail/auth/../hail/python/pinned-requirements.txt
+    #   requests-oauthlib
 pyasn1==0.5.0
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
@@ -49,14 +51,16 @@ pyasn1-modules==0.3.0
     #   -c hail/auth/../gear/pinned-requirements.txt
     #   -c hail/auth/../hail/python/pinned-requirements.txt
     #   google-auth
-requests==2.28.2
+requests==2.31.0
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
     #   -c hail/auth/../hail/python/dev/pinned-requirements.txt
     #   -c hail/auth/../hail/python/pinned-requirements.txt
     #   requests-oauthlib
 requests-oauthlib==1.3.1
-    # via google-auth-oauthlib
+    # via
+    #   -c hail/auth/../hail/python/pinned-requirements.txt
+    #   google-auth-oauthlib
 rsa==4.9
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
@@ -68,9 +72,10 @@ six==1.16.0
     #   -c hail/auth/../hail/python/dev/pinned-requirements.txt
     #   -c hail/auth/../hail/python/pinned-requirements.txt
     #   google-auth
-urllib3==1.26.15
+urllib3==1.26.16
     # via
     #   -c hail/auth/../gear/pinned-requirements.txt
     #   -c hail/auth/../hail/python/dev/pinned-requirements.txt
     #   -c hail/auth/../hail/python/pinned-requirements.txt
+    #   google-auth
     #   requests

batch/Dockerfile.worker

@@ -6,9 +6,8 @@ RUN hail-apt-get-install \
     iptables \
     openjdk-8-jre-headless \
     liblapack3 \
-    libyajl-dev \
-    wget \
-    xfsprogs
+    xfsprogs \
+    libyajl-dev
 
 RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.8 1
 
@@ -22,8 +21,8 @@ RUN echo "APT::Acquire::Retries \"5\";" > /etc/apt/apt.conf.d/80-retries && \
 
 {% elif global.cloud == "azure" %}
 # https://github.com/Azure/azure-storage-fuse/issues/603
-RUN hail-apt-get-install ca-certificates pkg-config libfuse-dev cmake libcurl4-gnutls-dev libgnutls28-dev uuid-dev libgcrypt20-dev wget && \
-    wget https://packages.microsoft.com/config/ubuntu/20.04/packages-microsoft-prod.deb && \
+RUN hail-apt-get-install ca-certificates pkg-config libfuse-dev cmake libcurl4-gnutls-dev libgnutls28-dev uuid-dev libgcrypt20-dev && \
+    curl -LO https://packages.microsoft.com/config/ubuntu/20.04/packages-microsoft-prod.deb && \
     dpkg -i packages-microsoft-prod.deb && \
     apt-get update && \
     hail-apt-get-install blobfuse
@@ -47,7 +46,7 @@ ENV PYSPARK_PYTHON python3
 
 COPY docker/core-site.xml ${SPARK_HOME}/conf/core-site.xml
 
-RUN wget https://github.com/jvm-profiling-tools/async-profiler/releases/download/v2.9/async-profiler-2.9-linux-x64.tar.gz -qO- | tar -zxvf -
+RUN curl -L https://github.com/jvm-profiling-tools/async-profiler/releases/download/v2.9/async-profiler-2.9-linux-x64.tar.gz | tar -zxvf -
 
 # Build crun in separate build step
 FROM base AS crun_builder

batch/batch/cloud/azure/worker/worker_api.py

@@ -5,12 +5,11 @@
 
 import aiohttp
 
-from gear.cloud_config import get_azure_config
 from hailtop import httpx
 from hailtop.aiocloud import aioazure
-from hailtop.utils import check_exec_output, request_retry_transient_errors, time_msecs
+from hailtop.utils import check_exec_output, retry_transient_errors, time_msecs
 
-from ....worker.worker_api import CloudWorkerAPI
+from ....worker.worker_api import CloudWorkerAPI, ContainerRegistryCredentials
 from ..instance_config import AzureSlimInstanceConfig
 from .credentials import AzureUserCredentials
 from .disk import AzureDisk
@@ -40,20 +39,24 @@ def create_disk(self, instance_name: str, disk_name: str, size_in_gb: int, mount
     def get_cloud_async_fs(self) -> aioazure.AzureAsyncFS:
         return aioazure.AzureAsyncFS(credentials=self.azure_credentials)
 
-    def get_compute_client(self) -> aioazure.AzureComputeClient:
-        azure_config = get_azure_config()
-        return aioazure.AzureComputeClient(azure_config.subscription_id, azure_config.resource_group)
-
     def user_credentials(self, credentials: Dict[str, str]) -> AzureUserCredentials:
         return AzureUserCredentials(credentials)
 
-    async def worker_access_token(self, session: httpx.ClientSession) -> Dict[str, str]:
+    async def worker_container_registry_credentials(self, session: httpx.ClientSession) -> ContainerRegistryCredentials:
         # https://docs.microsoft.com/en-us/azure/container-registry/container-registry-authentication?tabs=azure-cli#az-acr-login-with---expose-token
         return {
             'username': '00000000-0000-0000-0000-000000000000',
             'password': await self.acr_refresh_token.token(session),
         }
 
+    async def user_container_registry_credentials(
+        self, user_credentials: AzureUserCredentials
+    ) -> ContainerRegistryCredentials:
+        return {
+            'username': user_credentials.username,
+            'password': user_credentials.password,
+        }
+
     def instance_config_from_config_dict(self, config_dict: Dict[str, str]) -> AzureSlimInstanceConfig:
         return AzureSlimInstanceConfig.from_dict(config_dict)
 
@@ -112,6 +115,9 @@ async def unmount_cloudfuse(self, mount_base_path_data: str):
             os.remove(self._blobfuse_credential_files[mount_base_path_data])
             del self._blobfuse_credential_files[mount_base_path_data]
 
+    async def close(self):
+        pass
+
     def __str__(self):
         return f'subscription_id={self.subscription_id} resource_group={self.resource_group}'
 
@@ -137,18 +143,16 @@ class AadAccessToken(LazyShortLivedToken):
     async def _fetch(self, session: httpx.ClientSession) -> Tuple[str, int]:
         # https://docs.microsoft.com/en-us/azure/active-directory/managed-identities-azure-resources/how-to-use-vm-token#get-a-token-using-http
         params = {'api-version': '2018-02-01', 'resource': 'https://management.azure.com/'}
-        async with await request_retry_transient_errors(
-            session,
-            'GET',
+        resp_json = await retry_transient_errors(
+            session.get_read_json,
             'http://169.254.169.254/metadata/identity/oauth2/token',
             headers={'Metadata': 'true'},
             params=params,
             timeout=aiohttp.ClientTimeout(total=60),  # type: ignore
-        ) as resp:
-            resp_json = await resp.json()
-            access_token: str = resp_json['access_token']
-            expiration_time_ms = int(resp_json['expires_on']) * 1000
-            return access_token, expiration_time_ms
+        )
+        access_token: str = resp_json['access_token']
+        expiration_time_ms = int(resp_json['expires_on']) * 1000
+        return access_token, expiration_time_ms
 
 
 class AcrRefreshToken(LazyShortLivedToken):
@@ -164,14 +168,13 @@ async def _fetch(self, session: httpx.ClientSession) -> Tuple[str, int]:
             'service': self.acr_url,
             'access_token': await self.aad_access_token.token(session),
         }
-        async with await request_retry_transient_errors(
-            session,
-            'POST',
+        resp_json = await retry_transient_errors(
+            session.post_read_json,
             f'https://{self.acr_url}/oauth2/exchange',
             headers={'Content-Type': 'application/x-www-form-urlencoded'},
             data=data,
             timeout=aiohttp.ClientTimeout(total=60),  # type: ignore
-        ) as resp:
-            refresh_token: str = (await resp.json())['refresh_token']
-            expiration_time_ms = time_msecs() + 60 * 60 * 1000  # token expires in 3 hours so we refresh after 1 hour
-            return refresh_token, expiration_time_ms
+        )
+        refresh_token: str = resp_json['refresh_token']
+        expiration_time_ms = time_msecs() + 60 * 60 * 1000  # token expires in 3 hours so we refresh after 1 hour
+        return refresh_token, expiration_time_ms

batch/batch/cloud/gcp/driver/create_instance.py

@@ -268,10 +268,14 @@ def scheduling() -> dict:
 iptables --append FORWARD --destination $INTERNAL_GATEWAY_IP --jump ACCEPT
 # And this worker
 iptables --append FORWARD --destination $IP_ADDRESS --jump ACCEPT
-# Forbid outgoing requests to cluster-internal IP addresses
+# Allow traffic going to the internet
 INTERNET_INTERFACE=$(ip link list | grep ens | awk -F": " '{{ print $2 }}')
 iptables --append FORWARD --out-interface $INTERNET_INTERFACE ! --destination 10.128.0.0/16 --jump ACCEPT
 
+# [private]
+# Allow all traffic from the private job network
+iptables --append FORWARD --source 172.20.0.0/16 --jump ACCEPT
+
 {make_global_config_str}
 
 # retry once

batch/batch/cloud/gcp/worker/credentials.py

@@ -13,14 +13,6 @@ def __init__(self, data: Dict[str, str]):
     def cloud_env_name(self) -> str:
         return 'GOOGLE_APPLICATION_CREDENTIALS'
 
-    @property
-    def username(self):
-        return '_json_key'
-
-    @property
-    def password(self) -> str:
-        return self._key
-
     @property
     def mount_path(self):
         return '/gsa-key/key.json'

batch/batch/cloud/gcp/worker/worker_api.py

@@ -6,9 +6,9 @@
 
 from hailtop import httpx
 from hailtop.aiocloud import aiogoogle
-from hailtop.utils import check_exec_output, request_retry_transient_errors
+from hailtop.utils import check_exec_output, retry_transient_errors
 
-from ....worker.worker_api import CloudWorkerAPI
+from ....worker.worker_api import CloudWorkerAPI, ContainerRegistryCredentials
 from ..instance_config import GCPSlimInstanceConfig
 from .credentials import GCPUserCredentials
 from .disk import GCPDisk
@@ -46,22 +46,23 @@ def create_disk(self, instance_name: str, disk_name: str, size_in_gb: int, mount
     def get_cloud_async_fs(self) -> aiogoogle.GoogleStorageAsyncFS:
         return aiogoogle.GoogleStorageAsyncFS(session=self._google_session)
 
-    def get_compute_client(self) -> aiogoogle.GoogleComputeClient:
-        return self._compute_client
-
     def user_credentials(self, credentials: Dict[str, str]) -> GCPUserCredentials:
         return GCPUserCredentials(credentials)
 
-    async def worker_access_token(self, session: httpx.ClientSession) -> Dict[str, str]:
-        async with await request_retry_transient_errors(
-            session,
-            'POST',
+    async def worker_container_registry_credentials(self, session: httpx.ClientSession) -> ContainerRegistryCredentials:
+        token_dict = await retry_transient_errors(
+            session.post_read_json,
             'http://169.254.169.254/computeMetadata/v1/instance/service-accounts/default/token',
             headers={'Metadata-Flavor': 'Google'},
             timeout=aiohttp.ClientTimeout(total=60),  # type: ignore
-        ) as resp:
-            access_token = (await resp.json())['access_token']
-            return {'username': 'oauth2accesstoken', 'password': access_token}
+        )
+        access_token = token_dict['access_token']
+        return {'username': 'oauth2accesstoken', 'password': access_token}
+
+    async def user_container_registry_credentials(
+        self, user_credentials: GCPUserCredentials
+    ) -> ContainerRegistryCredentials:
+        return {'username': '_json_key', 'password': user_credentials.key}
 
     def instance_config_from_config_dict(self, config_dict: Dict[str, str]) -> GCPSlimInstanceConfig:
         return GCPSlimInstanceConfig.from_dict(config_dict)
@@ -118,5 +119,8 @@ async def unmount_cloudfuse(self, mount_base_path_data: str):
             os.remove(self._gcsfuse_credential_files[mount_base_path_data])
             del self._gcsfuse_credential_files[mount_base_path_data]
 
+    async def close(self):
+        await self._compute_client.close()
+
     def __str__(self):
         return f'project={self.project} zone={self.zone}'