Replace super(AsyncAgentBase) with super(DatabricksAgent) (#2590)

Signed-off-by: Kevin Su <pingsutw@apache.org>
flyteorg · Jul 19, 2024 · f04f206 · f04f206
1 parent 5a4ba2f
commit f04f206
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/plugins/flytekit-spark/flytekitplugins/spark/agent.py b/plugins/flytekit-spark/flytekitplugins/spark/agent.py
@@ -126,7 +126,7 @@ class DatabricksAgentV2(DatabricksAgent):
     """
 
     def __init__(self):
-        super(AsyncAgentBase, self).__init__(task_type_name="databricks", metadata_type=DatabricksJobMetadata)
+        super(DatabricksAgent, self).__init__(task_type_name="databricks", metadata_type=DatabricksJobMetadata)
 
 
 def get_header() -> typing.Dict[str, str]:

diff --git a/plugins/flytekit-spark/flytekitplugins/spark/task.py b/plugins/flytekit-spark/flytekitplugins/spark/task.py
@@ -145,9 +145,14 @@ def __init__(
                     self._default_applications_path or "local:///usr/local/bin/entrypoint.py"
                 )
 
+        if isinstance(task_config, DatabricksV2):
+            task_type = "databricks"
+        else:
+            task_type = "spark"
+
         super(PysparkFunctionTask, self).__init__(
             task_config=task_config,
-            task_type=self._SPARK_TASK_TYPE,
+            task_type=task_type,
             task_function=task_function,
             container_image=container_image,
             **kwargs,