[cherry-pick] KeyError: 'sparkConf' occurs when running a Databricks task without spark_conf (#3282)

Signed-off-by: Kevin Su <[email protected]>

Author: pingsutw

plugins/flytekit-spark/flytekitplugins/spark/agent.py

@@ -38,7 +38,7 @@ def _get_databricks_job_spec(task_template: TaskTemplate) -> dict:
         if not new_cluster.get("docker_image"):
             new_cluster["docker_image"] = {"url": container.image}
         if not new_cluster.get("spark_conf"):
-            new_cluster["spark_conf"] = custom["sparkConf"]
+            new_cluster["spark_conf"] = custom.get("sparkConf", {})
         if not new_cluster.get("spark_env_vars"):
             new_cluster["spark_env_vars"] = {k: v for k, v in envs.items()}
         else:

plugins/flytekit-spark/tests/test_spark_task.py

@@ -8,7 +8,7 @@
 
 from flytekit.core import context_manager
 from flytekitplugins.spark import Spark
-from flytekitplugins.spark.task import Databricks, new_spark_session
+from flytekitplugins.spark.task import Databricks, DatabricksV2, new_spark_session
 from pyspark.sql import SparkSession
 
 import flytekit
@@ -105,6 +105,45 @@ def my_databricks(a: int) -> int:
     assert my_databricks(a=3) == 3
 
 
+@pytest.mark.parametrize("spark_conf", [None, {"spark": "2"}])
+def test_databricks_v2(reset_spark_session, spark_conf):
+    databricks_conf = {
+        "name": "flytekit databricks plugin example",
+        "new_cluster": {
+            "spark_version": "11.0.x-scala2.12",
+            "node_type_id": "r3.xlarge",
+            "aws_attributes": {"availability": "ON_DEMAND"},
+            "num_workers": 4,
+            "docker_image": {"url": "pingsutw/databricks:latest"},
+        },
+        "timeout_seconds": 3600,
+        "max_retries": 1,
+        "spark_python_task": {
+            "python_file": "dbfs:///FileStore/tables/entrypoint-1.py",
+            "parameters": "ls",
+        },
+    }
+
+    databricks_instance = "account.cloud.databricks.com"
+
+    @task(
+        task_config=DatabricksV2(
+            databricks_conf=databricks_conf,
+            databricks_instance=databricks_instance,
+            spark_conf=spark_conf,
+        )
+    )
+    def my_databricks(a: int) -> int:
+        session = flytekit.current_context().spark_session
+        assert session.sparkContext.appName == "FlyteSpark: ex:local:local:local"
+        return a
+
+    assert my_databricks.task_config is not None
+    assert my_databricks.task_config.databricks_conf == databricks_conf
+    assert my_databricks.task_config.databricks_instance == databricks_instance
+    assert my_databricks.task_config.spark_conf == (spark_conf or {})
+    assert my_databricks(a=3) == 3
+
 def test_new_spark_session():
     name = "SessionName"
     spark_conf = {"spark1": "1", "spark2": "2"}

tests/flytekit/unit/experimental/test_eager_workflows.py

@@ -287,9 +287,15 @@ async def eager_wf_flyte_directory() -> str:
 @mock.patch("flytekit.core.data_persistence.FileAccessProvider.get_data")
 @mock.patch("flytekit.core.data_persistence.FileAccessProvider.put_data")
 @mock.patch("flytekit.core.utils.write_proto_to_file")
-def test_eager_workflow_dispatch(mock_write_to_file, mock_put_data, mock_get_data, mock_load_proto, event_loop):
+def test_eager_workflow_dispatch(mock_write_to_file, mock_put_data, mock_get_data, mock_load_proto):
     """Test that event loop is preserved after executing eager workflow via dispatch."""
 
+    try:
+        event_loop = asyncio.get_running_loop()
+    except RuntimeError:
+        event_loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(event_loop)
+
     @eager
     async def eager_wf():
         await asyncio.sleep(0.1)