|
| 1 | +"""Single source of truth for setting up the Spark session shared by FLAML's |
| 2 | +Spark-dependent test modules. |
| 3 | +
|
| 4 | +Several test files (``test/spark/test_0sparkml.py``, |
| 5 | +``test/spark/test_internal_mlflow.py``, ``test/automl/test_extra_models.py``) |
| 6 | +historically duplicated the same large ``SparkSession.builder`` configuration |
| 7 | +(synapseml + hadoop-azure + mlflow-spark JARs, MMLSpark Maven repo, |
| 8 | +log_model_allowlistFile, ANSI-mode handling). Keeping the duplicated copies |
| 9 | +in sync turned out to be brittle — this helper centralises it. |
| 10 | +
|
| 11 | +The module name starts with an underscore so pytest will not try to collect |
| 12 | +it as a test module. |
| 13 | +""" |
| 14 | + |
| 15 | +from __future__ import annotations |
| 16 | + |
| 17 | +import atexit |
| 18 | +import os |
| 19 | +import sys |
| 20 | +import warnings |
| 21 | + |
| 22 | +import mlflow |
| 23 | +from packaging.version import Version |
| 24 | + |
| 25 | +from flaml.automl.spark import disable_spark_ansi_mode, restore_spark_ansi_mode |
| 26 | +from flaml.tune.spark.utils import check_spark |
| 27 | + |
| 28 | + |
| 29 | +def _spark_jars_packages() -> str: |
| 30 | + """Maven coordinates for the JARs added to the Spark session.""" |
| 31 | + return ( |
| 32 | + "com.microsoft.azure:synapseml_2.12:1.0.14," |
| 33 | + "org.apache.hadoop:hadoop-azure:3.3.5," |
| 34 | + "com.microsoft.azure:azure-storage:8.6.6," |
| 35 | + f"org.mlflow:mlflow-spark_2.12:{mlflow.__version__}" |
| 36 | + if Version(mlflow.__version__) >= Version("2.9.0") |
| 37 | + else f"org.mlflow:mlflow-spark:{mlflow.__version__}" |
| 38 | + ) |
| 39 | + |
| 40 | + |
| 41 | +def init_spark_session(app_name: str = "MyApp", master: str = "local[2]"): |
| 42 | + """Build (or fetch) the SparkSession used by FLAML's Spark tests. |
| 43 | +
|
| 44 | + Imports ``pyspark`` lazily so this module remains importable in |
| 45 | + environments where pyspark is not installed (collection-time safety). |
| 46 | + """ |
| 47 | + import pyspark |
| 48 | + |
| 49 | + spark = ( |
| 50 | + pyspark.sql.SparkSession.builder.appName(app_name) |
| 51 | + .master(master) |
| 52 | + .config("spark.jars.packages", _spark_jars_packages()) |
| 53 | + .config("spark.jars.repositories", "https://mmlspark.azureedge.net/maven") |
| 54 | + .config("spark.sql.debug.maxToStringFields", "100") |
| 55 | + .config("spark.driver.extraJavaOptions", "-Xss1m") |
| 56 | + .config("spark.executor.extraJavaOptions", "-Xss1m") |
| 57 | + .getOrCreate() |
| 58 | + ) |
| 59 | + spark.sparkContext._conf.set( |
| 60 | + "spark.mlflow.pysparkml.autolog.logModelAllowlistFile", |
| 61 | + "https://mmlspark.blob.core.windows.net/publicwasb/log_model_allowlist.txt", |
| 62 | + ) |
| 63 | + return spark |
| 64 | + |
| 65 | + |
| 66 | +def setup_spark_for_tests( |
| 67 | + app_name: str = "MyApp", |
| 68 | + master: str = "local[2]", |
| 69 | +) -> tuple[object | None, bool]: |
| 70 | + """Initialise Spark for the importing test module. |
| 71 | +
|
| 72 | + Returns ``(spark, skip_spark)``: |
| 73 | +
|
| 74 | + * ``spark`` -- the configured ``SparkSession`` (or ``None`` when |
| 75 | + pyspark is not available, e.g. on macOS / Windows |
| 76 | + runners or when the package is not installed). |
| 77 | + * ``skip_spark`` -- ``True`` when the Spark-dependent tests in the caller |
| 78 | + should be skipped. Suitable for use with |
| 79 | + ``pytest.mark.skipif(skip_spark, ...)``. |
| 80 | +
|
| 81 | + ANSI-mode handling matches the prior duplicated setup: it is invoked |
| 82 | + unconditionally and ``restore_spark_ansi_mode`` is registered with |
| 83 | + ``atexit`` (this is a no-op when no active Spark session exists). |
| 84 | + """ |
| 85 | + warnings.simplefilter(action="ignore") |
| 86 | + |
| 87 | + spark: object | None = None |
| 88 | + skip_spark = False |
| 89 | + |
| 90 | + if sys.platform == "darwin" or "nt" in os.name: |
| 91 | + skip_spark = True |
| 92 | + else: |
| 93 | + try: |
| 94 | + spark = init_spark_session(app_name=app_name, master=master) |
| 95 | + spark_available, _ = check_spark() |
| 96 | + skip_spark = not spark_available |
| 97 | + except ImportError: |
| 98 | + skip_spark = True |
| 99 | + |
| 100 | + active_spark, ansi_conf, adjusted = disable_spark_ansi_mode() |
| 101 | + atexit.register(restore_spark_ansi_mode, active_spark, ansi_conf, adjusted) |
| 102 | + |
| 103 | + return spark, skip_spark |
0 commit comments