[batch] support multiple spark versions for jvm jobs (#15675)

## Change Description

JVM jobs may now declare which spark version their jar was built against
via an optional 'spark_version' field in the job spec. The front end and
worker default it to 3.5.3, the version every released client is
compatible
with, so existing clients are unaffected and need no changes.

- The worker no longer bakes pyspark into its image. It downloads an
archive
of the requested version's spark jars from cloud storage
(SPARK_ARCHIVE_URL_PREFIX, plumbed from the new
HAIL_QUERY_SPARK_ARCHIVE_SUBFOLDER env var) and extracts it to
/spark/\<version>, reusing the query-jar download mechanism via a
generic
download_artifact.
- scripts/upload-spark-jars.sh builds the archive from the pyspark sdist
and uploads it to the namespace's jar prefix; the new upload_spark_jars
CI step uploads the archive for env/SPARK_VERSION.
- The front end validates at job creation that an archive exists for the
requested spark version, making a missing archive a 400 at submission
rather than a worker-side failure.
- JVMEntryway keeps one shared spark classloader per `(spark classpath,
jar classpath)`
and is now plain java to avoid dependence on scala; the worker image
runs jre 21. Other images that never run hail code (tls,
infra bootstrap) use jre 25.
- Bump INSTANCE_VERSION: the entryway wire protocol now includes the
spark
classpath.

## Security Assessment

- This change potentially impacts the Hail Batch instance as deployed by
Broad Institute in GCP

### Impact Rating

- This change has a low security impact

### Impact Description

This change adds one new user-controllable input: the optional
`spark_version`
field on JVM job specs, which the front end interpolates into a
cloud-storage
URL and the worker uses as a path component under `/spark/`. It is
validated
with `re.fullmatch(r'\d+\.\d+\.\d+')` at job creation (before any use),
so no
path traversal or URL manipulation is possible, and it can only select
among
spark-jar archives that batch's own CI uploaded to the batch-owned
storage
prefix — the same trust model as the existing
ACCEPTABLE_QUERY_JAR_URL_PREFIX
check for query jars.

The worker now fetches spark jars from that batch-owned prefix instead
of
shipping them in the worker image; the provenance of code executed by
JVM jobs
is unchanged (CI-built artifacts in batch-controlled storage).

The JRE 11→25 upgrade and the JVMEntryway Scala→Java conversion do not
change
authentication, authorization, or secret handling.

### Appsec Review

- [ ] Required: The impact has been assessed and approved by appsec

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

Author: ehigham

Makefile

@@ -62,7 +62,7 @@ check-hail: check-hail-fast pylint-hailtop
 
 .PHONY: check-batch
 check-batch: check-batch-fast pylint-batch
-	cd batch/jvm-entryway && $(MILL) $(MILLOPTS) checkFormat + fix --check
+	cd batch/jvm-entryway && $(MILL) $(MILLOPTS) palantirformat --check
 
 .PHONY: check-services
 check-services: $(CHECK_SERVICES_MODULES)

batch/Dockerfile.worker

@@ -6,7 +6,7 @@ RUN hail-apt-get-install \
     iproute2 \
     iptables \
     ca-certificates-java \
-    openjdk-11-jre-headless \
+    openjdk-21-jre-headless \
     liblapack3 \
     xfsprogs \
     libyajl-dev
@@ -45,18 +45,13 @@ COPY batch/pinned-requirements.txt batch-requirements.txt
 RUN hail-pip-install \
     -r hailtop-requirements.txt \
     -r gear-requirements.txt \
-    -r batch-requirements.txt \
-    pyspark==3.5.3
+    -r batch-requirements.txt
 
 RUN curl -sLO https://github.com/async-profiler/async-profiler/releases/download/v4.3/async-profiler-4.3-linux-x64.tar.gz \
  && echo '69a16462c34c06ff55618f41653cffad1f8946822d30842512a3e0e774841c06  async-profiler-4.3-linux-x64.tar.gz' | sha256sum -c --status - \
  && tar -C /opt -zxvf async-profiler-4.3-linux-x64.tar.gz
 
-ENV SPARK_HOME=$VIRTUAL_ENV/lib/python$UV_PYTHON/site-packages/pyspark
-ENV PATH="$PATH:$SPARK_HOME/sbin:$SPARK_HOME/bin:/opt/async-profiler-4.3-linux-x64/bin"
-ENV PYSPARK_PYTHON=python3
-
-COPY docker/core-site.xml ${SPARK_HOME}/conf/core-site.xml
+ENV PATH="$PATH:/opt/async-profiler-4.3-linux-x64/bin"
 
 # Build crun in separate build step
 FROM base AS crun_builder

batch/batch/cloud/azure/driver/create_instance.py

@@ -12,7 +12,7 @@
 from ....file_store import FileStore
 from ....instance_config import InstanceConfig
 from ...resource_utils import unreserved_worker_data_disk_size_gib
-from ...utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX
+from ...utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX, SPARK_ARCHIVE_URL_PREFIX
 from ..resource_utils import azure_machine_type_to_parts
 
 log = logging.getLogger('create_instance')
@@ -135,6 +135,7 @@ def create_vm_config(
 WORKER_DATA_DISK_NAME="{worker_data_disk_name}"
 UNRESERVED_WORKER_DATA_DISK_SIZE_GB="{unreserved_disk_storage_gb}"
 ACCEPTABLE_QUERY_JAR_URL_PREFIX="{ACCEPTABLE_QUERY_JAR_URL_PREFIX}"
+SPARK_ARCHIVE_URL_PREFIX="{SPARK_ARCHIVE_URL_PREFIX}"
 
 # format worker data disk
 sudo mkfs.xfs -f -m reflink=1 -n ftype=1 {disk_location}
@@ -273,6 +274,7 @@ def create_vm_config(
 -e INTERNET_INTERFACE=$INTERNET_INTERFACE \
 -e UNRESERVED_WORKER_DATA_DISK_SIZE_GB=$UNRESERVED_WORKER_DATA_DISK_SIZE_GB \
 -e ACCEPTABLE_QUERY_JAR_URL_PREFIX=$ACCEPTABLE_QUERY_JAR_URL_PREFIX \
+-e SPARK_ARCHIVE_URL_PREFIX=$SPARK_ARCHIVE_URL_PREFIX \
 -e INTERNAL_GATEWAY_IP=$INTERNAL_GATEWAY_IP \
 -v /var/run/docker.sock:/var/run/docker.sock \
 -v /var/run/netns:/var/run/netns:shared \

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

@@ -12,7 +12,7 @@
 from ....file_store import FileStore
 from ....instance_config import InstanceConfig
 from ...resource_utils import unreserved_worker_data_disk_size_gib
-from ...utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX
+from ...utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX, SPARK_ARCHIVE_URL_PREFIX
 from ..resource_utils import GPUConfig, gcp_local_ssd_count, gcp_machine_type_to_parts, machine_type_to_gpu
 
 log = logging.getLogger('create_instance')
@@ -185,6 +185,7 @@ def scheduling() -> dict:
 NUM_LOCAL_SSDS="{num_local_ssds}"
 UNRESERVED_WORKER_DATA_DISK_SIZE_GB="{unreserved_disk_storage_gb}"
 ACCEPTABLE_QUERY_JAR_URL_PREFIX="{ACCEPTABLE_QUERY_JAR_URL_PREFIX}"
+SPARK_ARCHIVE_URL_PREFIX="{SPARK_ARCHIVE_URL_PREFIX}"
 
 CORES=$(nproc)
 NAMESPACE=$(curl -s -H "Metadata-Flavor: Google" "http://metadata.google.internal/computeMetadata/v1/instance/attributes/namespace")
@@ -356,6 +357,7 @@ def scheduling() -> dict:
 -e INTERNET_INTERFACE=$INTERNET_INTERFACE \
 -e UNRESERVED_WORKER_DATA_DISK_SIZE_GB=$UNRESERVED_WORKER_DATA_DISK_SIZE_GB \
 -e ACCEPTABLE_QUERY_JAR_URL_PREFIX=$ACCEPTABLE_QUERY_JAR_URL_PREFIX \
+-e SPARK_ARCHIVE_URL_PREFIX=$SPARK_ARCHIVE_URL_PREFIX \
 -e INTERNAL_GATEWAY_IP=$INTERNAL_GATEWAY_IP \
 -v /var/run/docker.sock:/var/run/docker.sock \
 -v /var/run/netns:/var/run/netns:shared \

batch/batch/cloud/terra/azure/driver/driver.py

@@ -41,7 +41,7 @@
     azure_worker_memory_per_core_mib,
     azure_worker_properties_to_machine_type,
 )
-from ....utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX
+from ....utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX, SPARK_ARCHIVE_URL_PREFIX
 from ..instance_config import TerraAzureSlimInstanceConfig
 
 log = logging.getLogger('driver')
@@ -160,6 +160,7 @@ def create_vm_config(
 
       UNRESERVED_WORKER_DATA_DISK_SIZE_GB=50
       ACCEPTABLE_QUERY_JAR_URL_PREFIX={shq(ACCEPTABLE_QUERY_JAR_URL_PREFIX)}
+      SPARK_ARCHIVE_URL_PREFIX={shq(SPARK_ARCHIVE_URL_PREFIX)}
 
       sudo mkdir -p /host/batch
       sudo mkdir -p /host/logs
@@ -236,6 +237,7 @@ def create_vm_config(
       -e WORKSPACE_ID=$WORKSPACE_ID \
       -e UNRESERVED_WORKER_DATA_DISK_SIZE_GB=$UNRESERVED_WORKER_DATA_DISK_SIZE_GB \
       -e ACCEPTABLE_QUERY_JAR_URL_PREFIX=$ACCEPTABLE_QUERY_JAR_URL_PREFIX \
+      -e SPARK_ARCHIVE_URL_PREFIX=$SPARK_ARCHIVE_URL_PREFIX \
       -e REGION=$REGION \
       -v /var/run/docker.sock:/var/run/docker.sock \
       -v /var/run/netns:/var/run/netns:shared \

batch/batch/cloud/utils.py

@@ -30,22 +30,23 @@ def possible_cloud_locations(cloud: str) -> Set[str]:
     return gcp_config.regions
 
 
-def _acceptable_query_jar_url_prefix() -> str:
+def _query_storage_url_prefix(subfolder_envvar: str) -> str:
     cloud = os.environ['CLOUD']
     query_storage_uri = os.environ['HAIL_QUERY_STORAGE_URI']
-    jar_subfolder = os.environ['HAIL_QUERY_ACCEPTABLE_JAR_SUBFOLDER']
-    acceptable_query_jar_url_prefix = query_storage_uri + jar_subfolder
+    subfolder = os.environ[subfolder_envvar]
+    url_prefix = query_storage_uri + subfolder
 
-    assert jar_subfolder[0] == '/', (query_storage_uri, jar_subfolder)
-    assert query_storage_uri[-1] != '/', (query_storage_uri, jar_subfolder)
+    assert subfolder[0] == '/', (query_storage_uri, subfolder)
+    assert query_storage_uri[-1] != '/', (query_storage_uri, subfolder)
 
     if cloud == 'gcp':
-        assert GoogleStorageAsyncFS.valid_url(acceptable_query_jar_url_prefix)
+        assert GoogleStorageAsyncFS.valid_url(url_prefix)
     else:
         assert cloud == 'azure'
-        assert AzureAsyncFS.valid_url(acceptable_query_jar_url_prefix)
+        assert AzureAsyncFS.valid_url(url_prefix)
 
-    return acceptable_query_jar_url_prefix
+    return url_prefix
 
 
-ACCEPTABLE_QUERY_JAR_URL_PREFIX = _acceptable_query_jar_url_prefix()
+ACCEPTABLE_QUERY_JAR_URL_PREFIX = _query_storage_url_prefix('HAIL_QUERY_ACCEPTABLE_JAR_SUBFOLDER')
+SPARK_ARCHIVE_URL_PREFIX = _query_storage_url_prefix('HAIL_QUERY_SPARK_ARCHIVE_SUBFOLDER')

batch/batch/front_end/front_end.py

@@ -96,7 +96,7 @@
     memory_to_worker_type,
     valid_machine_types,
 )
-from ..cloud.utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX
+from ..cloud.utils import ACCEPTABLE_QUERY_JAR_URL_PREFIX, SPARK_ARCHIVE_URL_PREFIX
 from ..exceptions import (
     BatchOperationAlreadyCompletedError,
     BatchUserError,
@@ -110,6 +110,7 @@
 from ..file_store import FileStore
 from ..globals import (
     BATCH_FORMAT_VERSION,
+    DEFAULT_SPARK_VERSION,
     HTTP_CLIENT_MAX_SIZE,
     RESERVED_STORAGE_GB_PER_CORE,
     complete_states,
@@ -1312,6 +1313,10 @@ async def _create_jobs(
                 if not jar_url.startswith(ACCEPTABLE_QUERY_JAR_URL_PREFIX):
                     raise web.HTTPBadRequest(reason=f'unacceptable JAR url: {jar_url}')
 
+            spark_version = spec['process'].setdefault('spark_version', DEFAULT_SPARK_VERSION)
+            if not await app[AppKeys.SPARK_ARCHIVE_EXISTENCE_CACHE].lookup(spark_version):
+                raise web.HTTPBadRequest(reason=f'no spark jars archive exists for spark version {spark_version}')
+
         req_memory_bytes: Optional[int]
         if machine_type is None:
             if 'cpu' not in resources:
@@ -3835,6 +3840,7 @@ def log(self, request, response, time):
 
 class AppKeys(CommonAiohttpAppKeys):
     QOB_JAR_RESOLUTION_CACHE = web.AppKey('qob_jar_resolution_cache', TimeLimitedMaxSizeCache[Tuple[str, str], str])
+    SPARK_ARCHIVE_EXISTENCE_CACHE = web.AppKey('spark_archive_existence_cache', TimeLimitedMaxSizeCache[str, bool])
 
 
 async def on_startup(app):
@@ -3923,6 +3929,13 @@ async def resolve_qob_jar_url(git_revision: str) -> Optional[str]:
         resolve_qob_jar_url, int(1e10), 100, AppKeys.QOB_JAR_RESOLUTION_CACHE._name
     )
 
+    async def spark_archive_exists(spark_version: str) -> bool:
+        return await fs.exists(SPARK_ARCHIVE_URL_PREFIX + '/spark-' + spark_version + '.tar.gz')
+
+    app[AppKeys.SPARK_ARCHIVE_EXISTENCE_CACHE] = TimeLimitedMaxSizeCache(
+        spark_archive_exists, int(1e10), 100, AppKeys.SPARK_ARCHIVE_EXISTENCE_CACHE._name
+    )
+
     app['task_manager'].ensure_future(periodically_call(5, _refresh, app))
 
 

batch/batch/front_end/validate.py

@@ -82,6 +82,8 @@
                 }),
                 required('command'): listof(str_type),
                 'profile': bool_type,
+                # the worker uses this as a path component, hence the strict pattern
+                'spark_version': regex(r'\d+\.\d+\.\d+'),
             },
         },
     ),

batch/batch/globals.py

@@ -23,7 +23,10 @@
 
 BATCH_FORMAT_VERSION = 7
 STATUS_FORMAT_VERSION = 5
-INSTANCE_VERSION = 35
+INSTANCE_VERSION = 36
+
+# jvm jobs from hail versions that predate multi-spark support were built against spark 3.5
+DEFAULT_SPARK_VERSION = '3.5.3'
 
 MAX_PERSISTENT_SSD_SIZE_GIB = 64 * 1024
 RESERVED_STORAGE_GB_PER_CORE = 5