Repository: mlflow/mlflow

Issue:
Use a context manager when creating a spark session
### Summary

A context manager should be used when creating a spark session to kill it gracefully on error:

```diff
diff --git a/tests/models/test_signature.py b/tests/models/test_signature.py
index 78c1f3848..d8525d6bc 100644
--- a/tests/models/test_signature.py
+++ b/tests/models/test_signature.py
@@ -163,14 +163,14 @@ def test_signature_inference_infers_datime_types_as_expected():
     signature = infer_signature(test_df)
     assert signature.inputs == Schema([ColSpec(DataType.datetime, name=col_name)])
 
-    spark = pyspark.sql.SparkSession.builder.getOrCreate()
-    spark_df = spark.range(1).selectExpr(
-        "current_timestamp() as timestamp", "current_date() as date"
-    )
-    signature = infer_signature(spark_df)
-    assert signature.inputs == Schema(
-        [ColSpec(DataType.datetime, name="timestamp"), ColSpec(DataType.datetime, name="date")]
-    )
+    with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+        spark_df = spark.range(1).selectExpr(
+            "current_timestamp() as timestamp", "current_date() as date"
+        )
+        signature = infer_signature(spark_df)
+        assert signature.inputs == Schema(
+            [ColSpec(DataType.datetime, name="timestamp"), ColSpec(DataType.datetime, name="date")]
+        )
 
 
 def test_set_signature_to_logged_model():
diff --git a/tests/pyfunc/test_spark.py b/tests/pyfunc/test_spark.py
index 6d7cdf1bf..bc785469c 100644
--- a/tests/pyfunc/test_spark.py
+++ b/tests/pyfunc/test_spark.py
@@ -15,6 +15,7 @@ import pyspark
 import pytest
 from packaging.version import Version
 from pandas.testing import assert_frame_equal
+from pyspark.sql import SparkSession
 from pyspark.sql.functions import col, pandas_udf, struct
 from pyspark.sql.types import (
     ArrayType,
@@ -58,13 +59,13 @@ types = [np.int32, int, str, np.float32, np.double, bool]
 
 
 def score_model_as_udf(model_uri, pandas_df, result_type="double"):
-    spark = get_spark_session(pyspark.SparkConf())
-    spark_df = spark.createDataFrame(pandas_df).coalesce(1)
-    pyfunc_udf = spark_udf(
-        spark=spark, model_uri=model_uri, result_type=result_type, env_manager="local"
-    )
-    new_df = spark_df.withColumn("prediction", pyfunc_udf(*pandas_df.columns))
-    return [x["prediction"] for x in new_df.collect()]
+    with get_spark_session(pyspark.SparkConf()) as spark:
+        spark_df = spark.createDataFrame(pandas_df).coalesce(1)
+        pyfunc_udf = spark_udf(
+            spark=spark, model_uri=model_uri, result_type=result_type, env_manager="local"
+        )
+        new_df = spark_df.withColumn("prediction", pyfunc_udf(*pandas_df.columns))
+        return [x["prediction"] for x in new_df.collect()]
 
 
 class ConstantPyfuncWrapper:
diff --git a/tests/spark/test_spark_model_export.py b/tests/spark/test_spark_model_export.py
index ada83712c..fe6d135fa 100644
--- a/tests/spark/test_spark_model_export.py
+++ b/tests/spark/test_spark_model_export.py
@@ -122,9 +122,9 @@ def iris_pandas_df():
 def iris_df(spark_context):
     pdf = iris_pandas_df()
     feature_names = list(pdf.drop("label", axis=1).columns)
-    spark_session = pyspark.sql.SparkSession(spark_context)
-    iris_spark_df = spark_session.createDataFrame(pdf)
-    return feature_names, pdf, iris_spark_df
+    with pyspark.sql.SparkSession(spark_context) as spark:
+        iris_spark_df = spark.createDataFrame(pdf)
+        yield feature_names, pdf, iris_spark_df
 
 
 @pytest.fixture(scope="module")
```

### Notes

- Make sure to open a PR from a **non-master** branch.
- Sign off the commit using the `-s` flag when making a commit:

  ```sh
  git commit -s -m "..."
           # ^^ make sure to use this
  ```

- Include `#{issue_number}` (e.g. `#123`) in the PR description when opening a PR.



Discussion:
Hi, Can I give it a try?
Hi @Dennis40816 it's all yours! Thank you! :) 

Resolving patch:
diff --git a/examples/multistep_workflow/als.py b/examples/multistep_workflow/als.py
--- a/examples/multistep_workflow/als.py
+++ b/examples/multistep_workflow/als.py
@@ -23,45 +23,46 @@
 def train_als(ratings_data, split_prop, max_iter, reg_param, rank, cold_start_strategy):
     seed = 42
 
-    spark = pyspark.sql.SparkSession.builder.getOrCreate()
+    with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+        ratings_df = spark.read.parquet(ratings_data)
+        (training_df, test_df) = ratings_df.randomSplit([split_prop, 1 - split_prop], seed=seed)
+        training_df.cache()
+        test_df.cache()
 
-    ratings_df = spark.read.parquet(ratings_data)
-    (training_df, test_df) = ratings_df.randomSplit([split_prop, 1 - split_prop], seed=seed)
-    training_df.cache()
-    test_df.cache()
+        mlflow.log_metric("training_nrows", training_df.count())
+        mlflow.log_metric("test_nrows", test_df.count())
 
-    mlflow.log_metric("training_nrows", training_df.count())
-    mlflow.log_metric("test_nrows", test_df.count())
+        print(f"Training: {training_df.count()}, test: {test_df.count()}")
 
-    print(f"Training: {training_df.count()}, test: {test_df.count()}")
+        als = (
+            ALS()
+            .setUserCol("userId")
+            .setItemCol("movieId")
+            .setRatingCol("rating")
+            .setPredictionCol("predictions")
+            .setMaxIter(max_iter)
+            .setSeed(seed)
+            .setRegParam(reg_param)
+            .setColdStartStrategy(cold_start_strategy)
+            .setRank(rank)
+        )
 
-    als = (
-        ALS()
-        .setUserCol("userId")
-        .setItemCol("movieId")
-        .setRatingCol("rating")
-        .setPredictionCol("predictions")
-        .setMaxIter(max_iter)
-        .setSeed(seed)
-        .setRegParam(reg_param)
-        .setColdStartStrategy(cold_start_strategy)
-        .setRank(rank)
-    )
+        als_model = Pipeline(stages=[als]).fit(training_df)
 
-    als_model = Pipeline(stages=[als]).fit(training_df)
+        reg_eval = RegressionEvaluator(
+            predictionCol="predictions", labelCol="rating", metricName="mse"
+        )
 
-    reg_eval = RegressionEvaluator(predictionCol="predictions", labelCol="rating", metricName="mse")
+        predicted_test_dF = als_model.transform(test_df)
 
-    predicted_test_dF = als_model.transform(test_df)
+        test_mse = reg_eval.evaluate(predicted_test_dF)
+        train_mse = reg_eval.evaluate(als_model.transform(training_df))
 
-    test_mse = reg_eval.evaluate(predicted_test_dF)
-    train_mse = reg_eval.evaluate(als_model.transform(training_df))
-
-    print(f"The model had a MSE on the test set of {test_mse}")
-    print(f"The model had a MSE on the (train) set of {train_mse}")
-    mlflow.log_metric("test_mse", test_mse)
-    mlflow.log_metric("train_mse", train_mse)
-    mlflow.spark.log_model(als_model, "als-model")
+        print(f"The model had a MSE on the test set of {test_mse}")
+        print(f"The model had a MSE on the (train) set of {train_mse}")
+        mlflow.log_metric("test_mse", test_mse)
+        mlflow.log_metric("train_mse", train_mse)
+        mlflow.spark.log_model(als_model, "als-model")
 
 
 if __name__ == "__main__":
diff --git a/examples/multistep_workflow/etl_data.py b/examples/multistep_workflow/etl_data.py
--- a/examples/multistep_workflow/etl_data.py
+++ b/examples/multistep_workflow/etl_data.py
@@ -22,20 +22,20 @@ def etl_data(ratings_csv, max_row_limit):
     with mlflow.start_run():
         tmpdir = tempfile.mkdtemp()
         ratings_parquet_dir = os.path.join(tmpdir, "ratings-parquet")
-        spark = pyspark.sql.SparkSession.builder.getOrCreate()
         print(f"Converting ratings CSV {ratings_csv} to Parquet {ratings_parquet_dir}")
-        ratings_df = (
-            spark.read.option("header", "true")
-            .option("inferSchema", "true")
-            .csv(ratings_csv)
-            .drop("timestamp")
-        )  # Drop unused column
-        ratings_df.show()
-        if max_row_limit != -1:
-            ratings_df = ratings_df.limit(max_row_limit)
-        ratings_df.write.parquet(ratings_parquet_dir)
-        print(f"Uploading Parquet ratings: {ratings_parquet_dir}")
-        mlflow.log_artifacts(ratings_parquet_dir, "ratings-parquet-dir")
+        with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+            ratings_df = (
+                spark.read.option("header", "true")
+                .option("inferSchema", "true")
+                .csv(ratings_csv)
+                .drop("timestamp")
+            )  # Drop unused column
+            ratings_df.show()
+            if max_row_limit != -1:
+                ratings_df = ratings_df.limit(max_row_limit)
+            ratings_df.write.parquet(ratings_parquet_dir)
+            print(f"Uploading Parquet ratings: {ratings_parquet_dir}")
+            mlflow.log_artifacts(ratings_parquet_dir, "ratings-parquet-dir")
 
 
 if __name__ == "__main__":
diff --git a/examples/multistep_workflow/train_keras.py b/examples/multistep_workflow/train_keras.py
--- a/examples/multistep_workflow/train_keras.py
+++ b/examples/multistep_workflow/train_keras.py
@@ -29,85 +29,87 @@ def train_keras(ratings_data, als_model_uri, hidden_units):
     np.random.seed(0)
     tf.set_random_seed(42)  # For reproducibility
 
-    spark = pyspark.sql.SparkSession.builder.getOrCreate()
-    als_model = mlflow.spark.load_model(als_model_uri).stages[0]
-
-    ratings_df = spark.read.parquet(ratings_data)
-
-    (training_df, test_df) = ratings_df.randomSplit([0.8, 0.2], seed=42)
-    training_df.cache()
-    test_df.cache()
-
-    mlflow.log_metric("training_nrows", training_df.count())
-    mlflow.log_metric("test_nrows", test_df.count())
-
-    print(f"Training: {training_df.count()}, test: {test_df.count()}")
-
-    user_factors = als_model.userFactors.selectExpr("id as userId", "features as uFeatures")
-    item_factors = als_model.itemFactors.selectExpr("id as movieId", "features as iFeatures")
-    joined_train_df = training_df.join(item_factors, on="movieId").join(user_factors, on="userId")
-    joined_test_df = test_df.join(item_factors, on="movieId").join(user_factors, on="userId")
-
-    # We'll combine the movies and ratings vectors into a single vector of length 24.
-    # We will then explode this features vector into a set of columns.
-    def concat_arrays(*args):
-        return list(chain(*args))
-
-    concat_arrays_udf = udf(concat_arrays, ArrayType(FloatType()))
-
-    concat_train_df = joined_train_df.select(
-        "userId",
-        "movieId",
-        concat_arrays_udf(col("iFeatures"), col("uFeatures")).alias("features"),
-        col("rating").cast("float"),
-    )
-    concat_test_df = joined_test_df.select(
-        "userId",
-        "movieId",
-        concat_arrays_udf(col("iFeatures"), col("uFeatures")).alias("features"),
-        col("rating").cast("float"),
-    )
-
-    pandas_df = concat_train_df.toPandas()
-    pandas_test_df = concat_test_df.toPandas()
-
-    # This syntax will create a new DataFrame where elements of the 'features' vector
-    # are each in their own column. This is what we'll train our neural network on.
-    x_test = pd.DataFrame(pandas_test_df.features.values.tolist(), index=pandas_test_df.index)
-    x_train = pd.DataFrame(pandas_df.features.values.tolist(), index=pandas_df.index)
-
-    # Show matrix for example.
-    print("Training matrix:")
-    print(x_train)
-
-    # Create our Keras model with two fully connected hidden layers.
-    model = Sequential()
-    model.add(Dense(30, input_dim=24, activation="relu"))
-    model.add(Dense(hidden_units, activation="relu"))
-    model.add(Dense(1, activation="linear"))
-
-    model.compile(loss="mse", optimizer=keras.optimizers.Adam(lr=0.0001))
-
-    early_stopping = EarlyStopping(monitor="val_loss", min_delta=0.0001, patience=2, mode="auto")
-
-    model.fit(
-        x_train,
-        pandas_df["rating"],
-        validation_split=0.2,
-        verbose=2,
-        epochs=3,
-        batch_size=128,
-        shuffle=False,
-        callbacks=[early_stopping],
-    )
-
-    train_mse = model.evaluate(x_train, pandas_df["rating"], verbose=2)
-    test_mse = model.evaluate(x_test, pandas_test_df["rating"], verbose=2)
-    mlflow.log_metric("test_mse", test_mse)
-    mlflow.log_metric("train_mse", train_mse)
-
-    print(f"The model had a MSE on the test set of {test_mse}")
-    mlflow.tensorflow.log_model(model, "keras-model")
+    with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+        als_model = mlflow.spark.load_model(als_model_uri).stages[0]
+        ratings_df = spark.read.parquet(ratings_data)
+        (training_df, test_df) = ratings_df.randomSplit([0.8, 0.2], seed=42)
+        training_df.cache()
+        test_df.cache()
+
+        mlflow.log_metric("training_nrows", training_df.count())
+        mlflow.log_metric("test_nrows", test_df.count())
+
+        print(f"Training: {training_df.count()}, test: {test_df.count()}")
+
+        user_factors = als_model.userFactors.selectExpr("id as userId", "features as uFeatures")
+        item_factors = als_model.itemFactors.selectExpr("id as movieId", "features as iFeatures")
+        joined_train_df = training_df.join(item_factors, on="movieId").join(
+            user_factors, on="userId"
+        )
+        joined_test_df = test_df.join(item_factors, on="movieId").join(user_factors, on="userId")
+
+        # We'll combine the movies and ratings vectors into a single vector of length 24.
+        # We will then explode this features vector into a set of columns.
+        def concat_arrays(*args):
+            return list(chain(*args))
+
+        concat_arrays_udf = udf(concat_arrays, ArrayType(FloatType()))
+
+        concat_train_df = joined_train_df.select(
+            "userId",
+            "movieId",
+            concat_arrays_udf(col("iFeatures"), col("uFeatures")).alias("features"),
+            col("rating").cast("float"),
+        )
+        concat_test_df = joined_test_df.select(
+            "userId",
+            "movieId",
+            concat_arrays_udf(col("iFeatures"), col("uFeatures")).alias("features"),
+            col("rating").cast("float"),
+        )
+
+        pandas_df = concat_train_df.toPandas()
+        pandas_test_df = concat_test_df.toPandas()
+
+        # This syntax will create a new DataFrame where elements of the 'features' vector
+        # are each in their own column. This is what we'll train our neural network on.
+        x_test = pd.DataFrame(pandas_test_df.features.values.tolist(), index=pandas_test_df.index)
+        x_train = pd.DataFrame(pandas_df.features.values.tolist(), index=pandas_df.index)
+
+        # Show matrix for example.
+        print("Training matrix:")
+        print(x_train)
+
+        # Create our Keras model with two fully connected hidden layers.
+        model = Sequential()
+        model.add(Dense(30, input_dim=24, activation="relu"))
+        model.add(Dense(hidden_units, activation="relu"))
+        model.add(Dense(1, activation="linear"))
+
+        model.compile(loss="mse", optimizer=keras.optimizers.Adam(lr=0.0001))
+
+        early_stopping = EarlyStopping(
+            monitor="val_loss", min_delta=0.0001, patience=2, mode="auto"
+        )
+
+        model.fit(
+            x_train,
+            pandas_df["rating"],
+            validation_split=0.2,
+            verbose=2,
+            epochs=3,
+            batch_size=128,
+            shuffle=False,
+            callbacks=[early_stopping],
+        )
+
+        train_mse = model.evaluate(x_train, pandas_df["rating"], verbose=2)
+        test_mse = model.evaluate(x_test, pandas_test_df["rating"], verbose=2)
+        mlflow.log_metric("test_mse", test_mse)
+        mlflow.log_metric("train_mse", train_mse)
+
+        print(f"The model had a MSE on the test set of {test_mse}")
+        mlflow.tensorflow.log_model(model, "keras-model")
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/tests/models/test_signature.py b/tests/models/test_signature.py
--- a/tests/models/test_signature.py
+++ b/tests/models/test_signature.py
@@ -163,14 +163,14 @@ def test_signature_inference_infers_datime_types_as_expected():
     signature = infer_signature(test_df)
     assert signature.inputs == Schema([ColSpec(DataType.datetime, name=col_name)])
 
-    spark = pyspark.sql.SparkSession.builder.getOrCreate()
-    spark_df = spark.range(1).selectExpr(
-        "current_timestamp() as timestamp", "current_date() as date"
-    )
-    signature = infer_signature(spark_df)
-    assert signature.inputs == Schema(
-        [ColSpec(DataType.datetime, name="timestamp"), ColSpec(DataType.datetime, name="date")]
-    )
+    with pyspark.sql.SparkSession.builder.getOrCreate() as spark:
+        spark_df = spark.range(1).selectExpr(
+            "current_timestamp() as timestamp", "current_date() as date"
+        )
+        signature = infer_signature(spark_df)
+        assert signature.inputs == Schema(
+            [ColSpec(DataType.datetime, name="timestamp"), ColSpec(DataType.datetime, name="date")]
+        )
 
 
 def test_set_signature_to_logged_model():
diff --git a/tests/pyfunc/test_spark.py b/tests/pyfunc/test_spark.py
--- a/tests/pyfunc/test_spark.py
+++ b/tests/pyfunc/test_spark.py
@@ -58,7 +58,7 @@
 
 
 def score_model_as_udf(model_uri, pandas_df, result_type="double"):
-    spark = get_spark_session(pyspark.SparkConf())
+    spark = pyspark.sql.SparkSession.getActiveSession()
     spark_df = spark.createDataFrame(pandas_df).coalesce(1)
     pyfunc_udf = spark_udf(
         spark=spark, model_uri=model_uri, result_type=result_type, env_manager="local"
diff --git a/tests/spark/test_spark_model_export.py b/tests/spark/test_spark_model_export.py
--- a/tests/spark/test_spark_model_export.py
+++ b/tests/spark/test_spark_model_export.py
@@ -72,21 +72,23 @@ def spark_custom_env(tmp_path):
 
 def _get_spark_session_with_retry(max_tries=3):
     conf = pyspark.SparkConf()
-    for num_tries in range(max_tries):
+    for attempt in range(max_tries):
         try:
             return get_spark_session(conf)
-        except Exception:
-            if num_tries >= max_tries - 1:
+        except Exception as e:
+            if attempt >= max_tries - 1:
                 raise
-            _logger.exception(f"Attempt {num_tries} to create a SparkSession failed, retrying...")
+            _logger.exception(
+                f"Attempt {attempt} to create a SparkSession failed ({e!r}), retrying..."
+            )
 
 
 # Specify `autouse=True` to ensure that a context is created
 # before any tests are executed. This ensures that the Hadoop filesystem
 # does not create its own SparkContext without the MLeap libraries required by
 # other tests.
-@pytest.fixture(scope="module", autouse=True)
-def spark_context():
+@pytest.fixture(scope="module")
+def spark():
     if Version(pyspark.__version__) < Version("3.1"):
         # A workaround for this issue:
         # https://stackoverflow.com/questions/62109276/errorjava-lang-unsupportedoperationexception-for-pyspark-pandas-udf-documenta
@@ -103,9 +105,9 @@ def spark_context():
 spark.executor.extraJavaOptions="-Dio.netty.tryReflectionSetAccessible=true"
 """
             f.write(conf)
-    spark = _get_spark_session_with_retry()
-    yield spark.sparkContext
-    spark.stop()
+
+    with _get_spark_session_with_retry() as spark:
+        yield spark
 
 
 def iris_pandas_df():
@@ -119,11 +121,10 @@ def iris_pandas_df():
 
 
 @pytest.fixture(scope="module")
-def iris_df(spark_context):
+def iris_df(spark):
     pdf = iris_pandas_df()
     feature_names = list(pdf.drop("label", axis=1).columns)
-    spark_session = pyspark.sql.SparkSession(spark_context)
-    iris_spark_df = spark_session.createDataFrame(pdf)
+    iris_spark_df = spark.createDataFrame(pdf)
     return feature_names, pdf, iris_spark_df
 
 
@@ -170,8 +171,7 @@ def spark_model_transformer(iris_df):
 
 
 @pytest.fixture(scope="module")
-def spark_model_estimator(iris_df, spark_context):
-    # pylint: disable=unused-argument
+def spark_model_estimator(iris_df):
     feature_names, iris_pandas_df, iris_spark_df = iris_df
     assembler = VectorAssembler(inputCols=feature_names, outputCol="features")
     features_df = assembler.transform(iris_spark_df)
@@ -190,6 +190,7 @@ def model_path(tmp_path):
     return os.path.join(tmp_path, "model")
 
 
+@pytest.mark.usefixtures("spark")
 def test_hadoop_filesystem(tmp_path):
     # copy local dir to and back from HadoopFS and make sure the results match
     from mlflow.spark import _HadoopFileSystem as FS
diff --git a/tests/types/test_schema.py b/tests/types/test_schema.py
--- a/tests/types/test_schema.py
+++ b/tests/types/test_schema.py
@@ -632,10 +632,12 @@ def test_spark_type_mapping(pandas_df_with_all_types):
     )
     actual_spark_schema = schema.as_spark_schema()
     assert expected_spark_schema.jsonValue() == actual_spark_schema.jsonValue()
-    spark_session = pyspark.sql.SparkSession(pyspark.SparkContext.getOrCreate())
-    sparkdf = spark_session.createDataFrame(pandas_df_with_all_types, schema=actual_spark_schema)
-    schema2 = _infer_schema(sparkdf)
-    assert schema == schema2
+    with pyspark.sql.SparkSession(pyspark.SparkContext.getOrCreate()) as spark_session:
+        sparkdf = spark_session.createDataFrame(
+            pandas_df_with_all_types, schema=actual_spark_schema
+        )
+        schema2 = _infer_schema(sparkdf)
+        assert schema == schema2
 
     # test unnamed columns
     schema = Schema([ColSpec(col.type) for col in schema.inputs])