Add support for providing a default instance to the databricks agent (#3190)

Author: andrei-trandafir

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

@@ -1,5 +1,6 @@
 import http
 import json
+import os
 import typing
 from dataclasses import dataclass
 from typing import Optional
@@ -17,6 +18,7 @@
 aiohttp = lazy_module("aiohttp")
 
 DATABRICKS_API_ENDPOINT = "/api/2.1/jobs"
+DEFAULT_DATABRICKS_INSTANCE_ENV_KEY = "FLYTE_DATABRICKS_INSTANCE"
 
 
 @dataclass
@@ -69,7 +71,15 @@ async def create(
         self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs
     ) -> DatabricksJobMetadata:
         data = json.dumps(_get_databricks_job_spec(task_template))
-        databricks_instance = task_template.custom["databricksInstance"]
+        databricks_instance = task_template.custom.get(
+            "databricksInstance", os.getenv(DEFAULT_DATABRICKS_INSTANCE_ENV_KEY)
+        )
+
+        if not databricks_instance:
+            raise ValueError(
+                f"Missing databricks instance. Please set the value through the task config or set the {DEFAULT_DATABRICKS_INSTANCE_ENV_KEY} environment variable in the agent."
+            )
+
         databricks_url = f"https://{databricks_instance}{DATABRICKS_API_ENDPOINT}/runs/submit"
 
         async with aiohttp.ClientSession() as session:

plugins/flytekit-spark/tests/test_agent.py

@@ -9,19 +9,17 @@
 
 from flytekit.core.constants import FLYTE_FAIL_ON_ERROR
 from flytekitplugins.spark.agent import DATABRICKS_API_ENDPOINT, DatabricksJobMetadata, get_header, \
-    _get_databricks_job_spec
+    _get_databricks_job_spec, DEFAULT_DATABRICKS_INSTANCE_ENV_KEY
 
 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
+import os
 
-
-@pytest.mark.asyncio
-async def test_databricks_agent():
-    agent = AgentRegistry.get_agent("spark")
-
+@pytest.fixture(scope="function")
+def task_template() -> TaskTemplate:
     task_id = Identifier(
         resource_type=ResourceType.TASK, project="project", domain="domain", name="name", version="version"
     )
@@ -55,8 +53,7 @@ async def test_databricks_agent():
             },
             "timeout_seconds": 3600,
             "max_retries": 1,
-        },
-        "databricksInstance": "test-account.cloud.databricks.com",
+        }
     }
     container = Container(
         image="flyteorg/flytekit:databricks-0.18.0-py3.7",
@@ -103,6 +100,16 @@ async def test_databricks_agent():
         interface=None,
         type="spark",
     )
+
+    return dummy_template
+
+
+@pytest.mark.asyncio
+async def test_databricks_agent(task_template: TaskTemplate):
+    agent = AgentRegistry.get_agent("spark")
+
+    task_template.custom["databricksInstance"] = "test-account.cloud.databricks.com"
+
     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
@@ -124,8 +131,8 @@ async def test_databricks_agent():
     delete_url = f"https://test-account.cloud.databricks.com{DATABRICKS_API_ENDPOINT}/runs/cancel"
     with aioresponses() as mocked:
         mocked.post(create_url, status=http.HTTPStatus.OK, payload=mock_create_response)
-        res = await agent.create(dummy_template, None)
-        spec = _get_databricks_job_spec(dummy_template)
+        res = await agent.create(task_template, None)
+        spec = _get_databricks_job_spec(task_template)
         data = json.dumps(spec)
         mocked.assert_called_with(create_url, method="POST", data=data, headers=get_header())
         spark_envs = spec["new_cluster"]["spark_env_vars"]
@@ -147,3 +154,39 @@ async def test_databricks_agent():
     assert get_header() == {"Authorization": f"Bearer {mocked_token}", "content-type": "application/json"}
 
     mock.patch.stopall()
+
+
+@pytest.mark.asyncio
+async def test_agent_create_with_default_instance(task_template: TaskTemplate):
+    agent = AgentRegistry.get_agent("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
+
+    databricks_metadata = DatabricksJobMetadata(
+        databricks_instance="test-account.cloud.databricks.com",
+        run_id="123",
+    )
+
+    mock_create_response = {"run_id": "123"}
+
+    os.environ[DEFAULT_DATABRICKS_INSTANCE_ENV_KEY] = "test-account.cloud.databricks.com"
+
+    create_url = f"https://test-account.cloud.databricks.com{DATABRICKS_API_ENDPOINT}/runs/submit"
+    with aioresponses() as mocked:
+        mocked.post(create_url, status=http.HTTPStatus.OK, payload=mock_create_response)
+        res = await agent.create(task_template, None)
+        spec = _get_databricks_job_spec(task_template)
+        data = json.dumps(spec)
+        mocked.assert_called_with(create_url, method="POST", data=data, headers=get_header())
+        assert res == databricks_metadata
+
+    mock.patch.stopall()
+
+@pytest.mark.asyncio
+async def test_agent_create_with_no_instance(task_template: TaskTemplate):
+    agent = AgentRegistry.get_agent("spark")
+
+    with pytest.raises(ValueError) as e:
+        await agent.create(task_template, None)