Databricks Agent

flytekit/extend/backend/agent_service.py

@@ -40,7 +40,6 @@ async def CreateTask(self, request: CreateTaskRequest, context: grpc.ServicerCon
                 logger.error(f"failed to run sync create with error {e}")
                 raise
         except Exception as e:
-            logger.error(f"failed to create task with error {e}")
             context.set_code(grpc.StatusCode.INTERNAL)
             context.set_details(f"failed to create task with error {e}")
 
@@ -60,7 +59,6 @@ async def GetTask(self, request: GetTaskRequest, context: grpc.ServicerContext)
                 logger.error(f"failed to run sync get with error {e}")
                 raise
         except Exception as e:
-            logger.error(f"failed to get task with error {e}")
             context.set_code(grpc.StatusCode.INTERNAL)
             context.set_details(f"failed to get task with error {e}")
 
@@ -80,6 +78,5 @@ async def DeleteTask(self, request: DeleteTaskRequest, context: grpc.ServicerCon
                 logger.error(f"failed to run sync delete with error {e}")
                 raise
         except Exception as e:
-            logger.error(f"failed to delete task with error {e}")
             context.set_code(grpc.StatusCode.INTERNAL)
             context.set_details(f"failed to delete task with error {e}")

plugins/flytekit-spark/dev-requirements.in

@@ -0,0 +1,2 @@
+aioresponses
+pytest-asyncio

plugins/flytekit-spark/dev-requirements.txt

@@ -0,0 +1,44 @@
+#
+# This file is autogenerated by pip-compile with Python 3.9
+# by the following command:
+#
+#    pip-compile dev-requirements.in
+#
+aiohttp==3.8.5
+    # via aioresponses
+aioresponses==0.7.4
+    # via -r dev-requirements.in
+aiosignal==1.3.1
+    # via aiohttp
+async-timeout==4.0.3
+    # via aiohttp
+attrs==23.1.0
+    # via aiohttp
+charset-normalizer==3.2.0
+    # via aiohttp
+exceptiongroup==1.1.3
+    # via pytest
+frozenlist==1.4.0
+    # via
+    #   aiohttp
+    #   aiosignal
+idna==3.4
+    # via yarl
+iniconfig==2.0.0
+    # via pytest
+multidict==6.0.4
+    # via
+    #   aiohttp
+    #   yarl
+packaging==23.1
+    # via pytest
+pluggy==1.3.0
+    # via pytest
+pytest==7.4.0
+    # via pytest-asyncio
+pytest-asyncio==0.21.1
+    # via -r dev-requirements.in
+tomli==2.0.1
+    # via pytest
+yarl==1.9.2
+    # via aiohttp

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

@@ -17,6 +17,7 @@
 
 from flytekit.configuration import internal as _internal
 
+from .agent import DatabricksAgent
 from .pyspark_transformers import PySparkPipelineModelTransformer
 from .schema import SparkDataFrameSchemaReader, SparkDataFrameSchemaWriter, SparkDataFrameTransformer  # noqa
 from .sd_transformers import ParquetToSparkDecodingHandler, SparkToParquetEncodingHandler

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

@@ -0,0 +1,99 @@
+import json
+import pickle
+import typing
+from dataclasses import dataclass
+from typing import Optional
+
+import aiohttp
+import grpc
+from flyteidl.admin.agent_pb2 import PENDING, CreateTaskResponse, DeleteTaskResponse, GetTaskResponse, Resource
+
+import flytekit
+from flytekit.extend.backend.base_agent import AgentBase, AgentRegistry, convert_to_flyte_state
+from flytekit.models.literals import LiteralMap
+from flytekit.models.task import TaskTemplate
+
+
+@dataclass
+class Metadata:
+    databricks_instance: str
+    run_id: str
+
+
+class DatabricksAgent(AgentBase):
+    def __init__(self):
+        super().__init__(task_type="spark")
+
+    async def async_create(
+        self,
+        context: grpc.ServicerContext,
+        output_prefix: str,
+        task_template: TaskTemplate,
+        inputs: Optional[LiteralMap] = None,
+    ) -> CreateTaskResponse:
+        custom = task_template.custom
+        container = task_template.container
+        databricks_job = custom["databricksConf"]
+        if not databricks_job["new_cluster"].get("docker_image"):
+            databricks_job["new_cluster"]["docker_image"] = {"url": container.image}
+        if not databricks_job["new_cluster"].get("spark_conf"):
+            databricks_job["new_cluster"]["spark_conf"] = custom["sparkConf"]
+        databricks_job["spark_python_task"] = {
+            "python_file": custom["mainApplicationFile"],
+            "parameters": tuple(container.args),
+        }
+
+        databricks_instance = custom["databricksInstance"]
+        databricks_url = f"https://{databricks_instance}/api/2.0/jobs/runs/submit"
+        data = json.dumps(databricks_job)
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(databricks_url, headers=get_header(), data=data) as resp:
+                if resp.status != 200:
+                    raise Exception(f"Failed to create databricks job with error: {resp.reason}")
+                response = await resp.json()
+
+        metadata = Metadata(
+            databricks_instance=databricks_instance,
+            run_id=str(response["run_id"]),
+        )
+        return CreateTaskResponse(resource_meta=pickle.dumps(metadata))
+
+    async def async_get(self, context: grpc.ServicerContext, resource_meta: bytes) -> GetTaskResponse:
+        metadata = pickle.loads(resource_meta)
+        databricks_instance = metadata.databricks_instance
+        databricks_url = f"https://{databricks_instance}/api/2.0/jobs/runs/get?run_id={metadata.run_id}"
+
+        async with aiohttp.ClientSession() as session:
+            async with session.get(databricks_url, headers=get_header()) as resp:
+                if resp.status != 200:
+                    raise Exception(f"Failed to get databricks job {metadata.run_id} with error: {resp.reason}")
+                response = await resp.json()
+
+        cur_state = PENDING
+        if response.get("state") and response["state"].get("result_state"):
+            cur_state = convert_to_flyte_state(response["state"]["result_state"])
+
+        return GetTaskResponse(resource=Resource(state=cur_state))
+
+    async def async_delete(self, context: grpc.ServicerContext, resource_meta: bytes) -> DeleteTaskResponse:
+        metadata = pickle.loads(resource_meta)
+
+        databricks_url = f"https://{metadata.databricks_instance}/api/2.0/jobs/runs/cancel"
+        data = json.dumps({"run_id": metadata.run_id})
+
+        async with aiohttp.ClientSession() as session:
+            async with session.post(databricks_url, headers=get_header(), data=data) as resp:
+                if resp.status != 200:
+                    raise Exception(f"Failed to cancel databricks job {metadata.run_id} with error: {resp.reason}")
+                await resp.json()
+
+        return DeleteTaskResponse()
+
+
+def get_header() -> typing.Dict[str, str]:
+    token = flytekit.current_context().secrets.get("databricks", "access_token")
+    return {"Authorization": f"Bearer {token}", "content-type": "application/json"}
+
+
+AgentRegistry.register(DatabricksAgent())

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

@@ -111,16 +111,18 @@ def __init__(
         **kwargs,
     ):
         self.sess: Optional[SparkSession] = None
-        self._default_executor_path: Optional[str] = task_config.executor_path
-        self._default_applications_path: Optional[str] = task_config.applications_path
+        self._default_executor_path: str = task_config.executor_path
+        self._default_applications_path: str = task_config.applications_path
 
         if isinstance(container_image, ImageSpec):
             if container_image.base_image is None:
                 img = f"cr.flyte.org/flyteorg/flytekit:spark-{DefaultImages.get_version_suffix()}"
                 container_image.base_image = img
                 # default executor path and applications path in apache/spark-py:3.3.1
-                self._default_executor_path = "/usr/bin/python3"
-                self._default_applications_path = "local:///usr/local/bin/entrypoint.py"
+                self._default_executor_path = self._default_executor_path or "/usr/bin/python3"
+                self._default_applications_path = (
+                    self._default_applications_path or "local:///usr/local/bin/entrypoint.py"
+                )
         super(PysparkFunctionTask, self).__init__(
             task_config=task_config,
             task_type=self._SPARK_TASK_TYPE,

plugins/flytekit-spark/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "pyspark>=3.0.0"]
+plugin_requires = ["flytekit>=1.3.0b2,<2.0.0", "pyspark>=3.0.0", "aiohttp"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-spark/tests/test_agent.py

@@ -0,0 +1,136 @@
+import pickle
+from datetime import timedelta
+from unittest import mock
+from unittest.mock import MagicMock
+
+import grpc
+import pytest
+from aioresponses import aioresponses
+from flyteidl.admin.agent_pb2 import SUCCEEDED
+from flytekitplugins.spark.agent import Metadata, get_header
+
+from flytekit.extend.backend.base_agent import AgentRegistry
+from flytekit.interfaces.cli_identifiers import Identifier
+from flytekit.models import literals, task
+from flytekit.models.core.identifier import ResourceType
+from flytekit.models.task import Container, Resources, TaskTemplate
+
+
+@pytest.mark.asyncio
+async def test_databricks_agent():
+    ctx = MagicMock(spec=grpc.ServicerContext)
+    agent = AgentRegistry.get_agent("spark")
+
+    task_id = Identifier(
+        resource_type=ResourceType.TASK, project="project", domain="domain", name="name", version="version"
+    )
+    task_metadata = task.TaskMetadata(
+        True,
+        task.RuntimeMetadata(task.RuntimeMetadata.RuntimeType.FLYTE_SDK, "1.0.0", "python"),
+        timedelta(days=1),
+        literals.RetryStrategy(3),
+        True,
+        "0.1.1b0",
+        "This is deprecated!",
+        True,
+        "A",
+    )
+    task_config = {
+        "sparkConf": {
+            "spark.driver.memory": "1000M",
+            "spark.executor.memory": "1000M",
+            "spark.executor.cores": "1",
+            "spark.executor.instances": "2",
+            "spark.driver.cores": "1",
+        },
+        "mainApplicationFile": "dbfs:/entrypoint.py",
+        "databricksConf": {
+            "run_name": "flytekit databricks plugin example",
+            "new_cluster": {
+                "spark_version": "12.2.x-scala2.12",
+                "node_type_id": "n2-highmem-4",
+                "num_workers": 1,
+            },
+            "timeout_seconds": 3600,
+            "max_retries": 1,
+        },
+        "databricksInstance": "test-account.cloud.databricks.com",
+    }
+    container = Container(
+        image="flyteorg/flytekit:databricks-0.18.0-py3.7",
+        command=[],
+        args=[
+            "pyflyte-fast-execute",
+            "--additional-distribution",
+            "s3://my-s3-bucket/flytesnacks/development/24UYJEF2HDZQN3SG4VAZSM4PLI======/script_mode.tar.gz",
+            "--dest-dir",
+            "/root",
+            "--",
+            "pyflyte-execute",
+            "--inputs",
+            "s3://my-s3-bucket",
+            "--output-prefix",
+            "s3://my-s3-bucket",
+            "--raw-output-data-prefix",
+            "s3://my-s3-bucket",
+            "--checkpoint-path",
+            "s3://my-s3-bucket",
+            "--prev-checkpoint",
+            "s3://my-s3-bucket",
+            "--resolver",
+            "flytekit.core.python_auto_container.default_task_resolver",
+            "--",
+            "task-module",
+            "spark_local_example",
+            "task-name",
+            "hello_spark",
+        ],
+        resources=Resources(
+            requests=[],
+            limits=[],
+        ),
+        env={},
+        config={},
+    )
+
+    dummy_template = TaskTemplate(
+        id=task_id,
+        custom=task_config,
+        metadata=task_metadata,
+        container=container,
+        interface=None,
+        type="spark",
+    )
+    mocked_token = "mocked_databricks_token"
+    mocked_context = mock.patch("flytekit.current_context", autospec=True).start()
+    mocked_context.return_value.secrets.get.return_value = mocked_token
+
+    metadata_bytes = pickle.dumps(
+        Metadata(
+            databricks_instance="test-account.cloud.databricks.com",
+            run_id="123",
+        )
+    )
+
+    mock_create_response = {"run_id": "123"}
+    mock_get_response = {"run_id": "123", "state": {"result_state": "SUCCESS"}}
+    mock_delete_response = {}
+    create_url = "https://test-account.cloud.databricks.com/api/2.0/jobs/runs/submit"
+    get_url = "https://test-account.cloud.databricks.com/api/2.0/jobs/runs/get?run_id=123"
+    delete_url = "https://test-account.cloud.databricks.com/api/2.0/jobs/runs/cancel"
+    with aioresponses() as mocked:
+        mocked.post(create_url, status=200, payload=mock_create_response)
+        res = await agent.async_create(ctx, "/tmp", dummy_template, None)
+        assert res.resource_meta == metadata_bytes
+
+        mocked.get(get_url, status=200, payload=mock_get_response)
+        res = await agent.async_get(ctx, metadata_bytes)
+        assert res.resource.state == SUCCEEDED
+        assert res.resource.outputs == literals.LiteralMap({}).to_flyte_idl()
+
+        mocked.post(delete_url, status=200, payload=mock_delete_response)
+        await agent.async_delete(ctx, metadata_bytes)
+
+    assert get_header() == {"Authorization": f"Bearer {mocked_token}", "content-type": "application/json"}
+
+    mock.patch.stopall()