[DOP-26758] Include .jar files into worker image

Author: dolfinus

.pre-commit-config.yaml

@@ -39,13 +39,13 @@ repos:
       - id: chmod
         args: ['644']
         exclude_types: [shell]
-        exclude: ^(.*__main__\.py|syncmaster/server/scripts/.*\.py)$
+        exclude: ^(.*__main__\.py|syncmaster/server/scripts/.*\.py|docker/.*\.py)$
       - id: chmod
         args: ['755']
         types: [shell]
       - id: chmod
         args: ['755']
-        files: ^(.*__main__\.py|syncmaster/server/scripts/.*\.py)$
+        files: ^(.*__main__\.py|syncmaster/server/scripts/.*\.py|docker/.*\.py)$
       - id: insert-license
         types: [python]
         exclude: ^(syncmaster/server/dependencies/stub.py|docs/.*\.py|tests/.*\.py)$

docker-compose.test.yml

@@ -114,7 +114,6 @@ services:
       - SYNCMASTER__SERVER__STATIC_FILES__ENABLED=false
     volumes:
       - ./syncmaster:/app/syncmaster
-      - ./cached_jars:/root/.ivy2
       - ./reports:/app/reports
       - ./tests:/app/tests
       - ./pyproject.toml:/app/pyproject.toml

docker/Dockerfile.worker

@@ -40,19 +40,31 @@ RUN --mount=type=cache,target=/root/.cache/pypoetry \
         --without test,docs,dev \
     && python -m compileall -j 4 .venv
 
+FROM builder AS maven_packages
 
-FROM base AS prod
+ARG SPARK_CONNECTORS_TO_DOWNLOAD="all"
+ENV SPARK_CONNECTORS_TO_DOWNLOAD=$SPARK_CONNECTORS_TO_DOWNLOAD
+RUN --mount=type=bind,source=./syncmaster,target=/app/syncmaster \
+    --mount=type=bind,source=./docker/download_maven_packages.py,target=/app/docker/download_maven_packages.py \
+    mkdir /root/.ivy2 && \
+    python /app/docker/download_maven_packages.py
 
-# We don't need poetry and compilers in final image
-COPY --from=builder /app/.venv/ /app/.venv/
-COPY ./syncmaster/ /app/syncmaster/
-RUN python -m compileall syncmaster
+
+FROM base AS prod
 
 # Do not run production as root, to improve security.
 # Also user does not own anything inside the image, including venv and source code.
 RUN useradd syncmaster && \
-    mkdir -p /home/syncmaster && \
+    mkdir -p /home/syncmaster /home/syncmaster/.ivy2 && \
     chown -R syncmaster:syncmaster /home/syncmaster
+
+# We don't need poetry and compilers in final image
+COPY --from=builder /app/.venv/ /app/.venv/
+# custom Spark session function may download different jars, so syncmaster have to own them
+COPY --from=maven_packages --chown=syncmaster:syncmaster /root/.ivy2/ /home/syncmaster/.ivy2/
+
+COPY ./syncmaster/ /app/syncmaster/
+RUN python -m compileall syncmaster
 USER syncmaster
 
 
@@ -68,6 +80,7 @@ RUN --mount=type=cache,target=/root/.cache/pypoetry \
         --without docs,dev \
     && python -m compileall -j 4 .venv
 
+COPY --from=maven_packages /root/.ivy2/ /root/.ivy2/
 ENV SYNCMASTER__WORKER__CREATE_SPARK_SESSION_FUNCTION=tests.spark.get_worker_spark_session
 
 # Collect coverage from worker

docker/download_maven_packages.py

@@ -0,0 +1,59 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: 2024-2025 MTS PJSC
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+
+from syncmaster.worker.spark import get_excluded_packages, get_packages
+
+if TYPE_CHECKING:
+    from pyspark.sql import SparkSession
+
+log = logging.getLogger(__name__)
+
+
+def get_spark_session_conf_for_docker_image() -> dict:
+    types = os.getenv("SPARK_CONNECTORS_TO_DOWNLOAD", "all").split(",")
+    maven_packages: list[str] = get_packages(connection_types=set(types))
+    excluded_packages: list[str] = get_excluded_packages()
+
+    return {
+        "spark.jars.packages": ",".join(maven_packages),
+        "spark.jars.excludes": ",".join(excluded_packages),
+        "spark.sql.pyspark.jvmStacktrace.enabled": "true",
+        # use only minimal available resoures
+        "spark.driver.cores": "1",
+        "spark.driver.memory": "512M",
+        "spark.executor.cores": "1",
+        "spark.executor.memory": "512M",
+        "spark.executor.instances": "1",
+    }
+
+
+def get_worker_spark_session_for_docker() -> SparkSession:
+    """
+    Construct dummy Spark session with all .jars included.
+    Designed to be used in Dockerfile.worker to populate the image.
+    """
+    from pyspark.sql import SparkSession
+
+    spark_builder = SparkSession.builder.appName("syncmaster_jar_downloader")
+
+    for k, v in get_spark_session_conf_for_docker_image().items():
+        spark_builder = spark_builder.config(k, v)
+
+    return spark_builder.getOrCreate()
+
+
+def download_maven_packages():
+    log.info("Downloading Maven packages...")
+    with get_worker_spark_session_for_docker():
+        log.info("Done!")
+
+
+if __name__ == "__main__":
+    download_maven_packages()

docs/changelog/0.2.4.rst

@@ -0,0 +1,8 @@
+0.2.4 (2025-07-07)
+==================
+
+Improvements
+------------
+
+- Include all required jars from Maven to worker image. This increases image size, but drastically reduces time of Spark session startup.
+

docs/changelog/index.rst

@@ -3,6 +3,7 @@
     :caption: Changelog
 
     DRAFT
+    0.2.4
     0.2.3
     0.2.2
     0.2.1

syncmaster/worker/spark.py

@@ -47,44 +47,42 @@ def get_worker_spark_session(
     return spark_builder.getOrCreate()
 
 
-def get_packages(connection_type: str) -> list[str]:  # noqa: WPS212
+def get_packages(connection_types: set[str]) -> list[str]:  # noqa: WPS212
     import pyspark
     from onetl.connection import MSSQL, Clickhouse, MySQL, Oracle, Postgres, SparkS3
     from onetl.file.format import XML, Excel
 
+    spark_version = pyspark.__version__
     # excel version is hardcoded due to https://github.com/nightscape/spark-excel/issues/902
     file_formats_spark_packages: list[str] = [
-        *XML.get_packages(spark_version=pyspark.__version__),
+        *XML.get_packages(spark_version=spark_version),
         *Excel.get_packages(spark_version="3.5.1"),
     ]
 
-    if connection_type == "postgres":
-        return Postgres.get_packages()
-    if connection_type == "oracle":
-        return Oracle.get_packages()
-    if connection_type == "clickhouse":
-        return [
-            "io.github.mtsongithub.doetl:spark-dialect-extension_2.12:0.0.2",
-            *Clickhouse.get_packages(),
-        ]
-    if connection_type == "mssql":
-        return MSSQL.get_packages()
-    if connection_type == "mysql":
-        return MySQL.get_packages()
-    if connection_type == "s3":
-        import pyspark
-
-        spark_version = pyspark.__version__
-        return SparkS3.get_packages(spark_version=spark_version) + file_formats_spark_packages
-
-    if connection_type in ("hdfs", "sftp", "ftp", "ftps", "samba", "webdav"):
-        return file_formats_spark_packages
-
-    # If the database type does not require downloading .jar packages
-    return []
-
-
-def get_excluded_packages(db_type: str) -> list[str]:
+    result = []
+
+    if connection_types & {"postgres", "all"}:
+        result.extend(Postgres.get_packages())
+    if connection_types & {"oracle", "all"}:
+        result.extend(Oracle.get_packages())
+    if connection_types & {"clickhouse", "all"}:
+        result.append("io.github.mtsongithub.doetl:spark-dialect-extension_2.12:0.0.2")
+        result.extend(Clickhouse.get_packages())
+    if connection_types & {"mssql", "all"}:
+        result.extend(MSSQL.get_packages())
+    if connection_types & {"mysql", "all"}:
+        result.extend(MySQL.get_packages())
+
+    if connection_types & {"s3", "all"}:
+        result.extend(SparkS3.get_packages(spark_version=spark_version))
+
+    if connection_types & {"s3", "hdfs", "sftp", "ftp", "ftps", "samba", "webdav", "all"}:
+        result.extend(file_formats_spark_packages)
+
+    return result
+
+
+def get_excluded_packages() -> list[str]:
     from onetl.connection import SparkS3
 
     return SparkS3.get_exclude_packages()
@@ -95,16 +93,11 @@ def get_spark_session_conf(
     target: ConnectionDTO,
     resources: dict,
 ) -> dict:
-    maven_packages: list[str] = []
-    excluded_packages: list[str] = []
-
-    for db_type in source, target:
-        maven_packages.extend(get_packages(connection_type=db_type.type))  # type: ignore
-        excluded_packages.extend(get_excluded_packages(db_type=db_type.type))  # type: ignore
+    maven_packages: list[str] = get_packages(connection_types={source.type, target.type})
+    excluded_packages: list[str] = get_excluded_packages()
 
     memory_mb = math.ceil(resources["ram_bytes_per_task"] / 1024 / 1024)
     config = {
-        "spark.jars.packages": ",".join(maven_packages),
         "spark.sql.pyspark.jvmStacktrace.enabled": "true",
         "spark.hadoop.mapreduce.fileoutputcommitter.marksuccessfuljobs": "false",
         "spark.executor.cores": resources["cpu_cores_per_task"],