feat: add classic cluster option for Switch job deployment (#2222)

## Changes

### What does this PR do?

Add support for classic job cluster via environment variable
`LAKEBRIDGE_CLUSTER_TYPE=CLASSIC` for Switch jobs. This provides a
workaround for workspaces where serverless is not available.

### Relevant implementation details

- Added `switch_use_serverless` flag to `LakebridgeConfiguration` (not
persisted, used only during install)
- Environment variable `LAKEBRIDGE_CLUSTER_TYPE=CLASSIC` triggers
classic cluster mode
- `SwitchDeployment.install()` now accepts `use_serverless` parameter
- When `use_serverless=False`, creates a job cluster with:
  - LTS Spark version
  - Single worker node (16GB+ memory)
  - `USER_ISOLATION` data security mode

### Usage

```bash
# Default: Serverless (no change needed)
databricks labs lakebridge install-transpile --include-llm-transpiler true

# Classic cluster
LAKEBRIDGE_CLUSTER_TYPE=CLASSIC databricks labs lakebridge install-transpile --include-llm-transpiler true
```

### Caveats/things to watch out for when reviewing:

- Default behavior remains serverless (backwards compatible)
- Classic cluster sizing (16GB+ memory, 1 worker) is fixed for now

### Linked issues

Resolves #2219

### Functionality

- [ ] added relevant user documentation
- [ ] added new CLI command
- [x] modified existing command: `databricks labs lakebridge
install-transpile`

### Tests

- [x] manually tested
- [x] added unit tests
- [ ] added integration tests

Author: hiroyukinakazato-db

src/databricks/labs/lakebridge/cli.py

@@ -761,16 +761,23 @@ def install_transpile(
     ctx.add_user_agent_extra("cmd", "install-transpile")
     if artifact:
         ctx.add_user_agent_extra("artifact-overload", Path(artifact).name)
+    # Internal: use LAKEBRIDGE_CLUSTER_TYPE=CLASSIC env var to use classic job cluster
+    switch_use_serverless = os.environ.get("LAKEBRIDGE_CLUSTER_TYPE", "").upper() != "CLASSIC"
     if include_llm_transpiler:
         ctx.add_user_agent_extra("include-llm-transpiler", "true")
         # Decision was made not to prompt when include_llm_transpiler is set, and we expect users to use llm-transpile
         # and pass all the arguments.
         logger.info("Including LLM transpiler as part of install, interactive mode disabled: will skip questionnaire.")
         is_interactive = False
+
     user = w.current_user
     logger.debug(f"User: {user}")
     transpile_installer = installer(
-        w, transpiler_repository, is_interactive=is_interactive, include_llm=include_llm_transpiler
+        w,
+        transpiler_repository,
+        is_interactive=is_interactive,
+        include_llm=include_llm_transpiler,
+        switch_use_serverless=switch_use_serverless,
     )
     transpile_installer.run(module="transpile", artifact=artifact)
 

src/databricks/labs/lakebridge/config.py

@@ -276,3 +276,5 @@ class LakebridgeConfiguration:
     reconcile: ReconcileConfig | None
     # Temporary flag, indicating whether to include the LLM-based Switch transpiler.
     include_switch: bool = False
+    # Internal: Use serverless compute for Switch job. Set via LAKEBRIDGE_CLUSTER_TYPE env var.
+    switch_use_serverless: bool = True

src/databricks/labs/lakebridge/deployment/installation.py

@@ -97,7 +97,7 @@ def install(self, config: LakebridgeConfiguration):
             self._recon_deployment.install(config.reconcile, wheel_path)
         if config.include_switch:
             logger.info("Installing Switch transpiler to workspace.")
-            self._switch_deployment.install()
+            self._switch_deployment.install(use_serverless=config.switch_use_serverless)
 
     def uninstall(self, config: LakebridgeConfiguration):
         # This will remove all the Lakebridge modules

src/databricks/labs/lakebridge/deployment/switch.py

@@ -12,7 +12,8 @@
 from databricks.labs.blueprint.paths import WorkspacePath
 from databricks.sdk import WorkspaceClient
 from databricks.sdk.errors import InvalidParameterValue, NotFound
-from databricks.sdk.service.jobs import JobParameterDefinition, JobSettings, NotebookTask, Source, Task
+from databricks.sdk.service import compute
+from databricks.sdk.service.jobs import JobCluster, JobParameterDefinition, JobSettings, NotebookTask, Source, Task
 
 
 logger = logging.getLogger(__name__)
@@ -32,12 +33,12 @@ def __init__(
         self._installation = installation
         self._install_state = install_state
 
-    def install(self) -> None:
+    def install(self, use_serverless: bool = True) -> None:
         """Deploy Switch to workspace and configure resources."""
         logger.debug("Deploying Switch resources to workspace...")
         try:
             self._deploy_resources_to_workspace()
-            self._setup_job()
+            self._setup_job(use_serverless)
             logger.debug("Switch deployment completed")
         except (RuntimeError, ValueError, InvalidParameterValue) as e:
             msg = f"Failed to setup required resources for Switch llm transpiler: {e}"
@@ -108,11 +109,11 @@ def _enumerate_resources(
 
         yield from _enumerate_resources(root)
 
-    def _setup_job(self) -> None:
+    def _setup_job(self, use_serverless: bool = True) -> None:
         """Create or update Switch job."""
         existing_job_id = self._get_existing_job_id()
         logger.debug("Setting up Switch job in workspace...")
-        job_id = self._create_or_update_switch_job(existing_job_id)
+        job_id = self._create_or_update_switch_job(existing_job_id, use_serverless)
         self._install_state.jobs[self._INSTALL_STATE_KEY] = job_id
         self._install_state.save()
         job_url = f"{self._ws.config.host}/jobs/{job_id}"
@@ -129,9 +130,9 @@ def _get_existing_job_id(self) -> str | None:
         except (InvalidParameterValue, NotFound, ValueError):
             return None
 
-    def _create_or_update_switch_job(self, job_id: str | None) -> str:
+    def _create_or_update_switch_job(self, job_id: str | None, use_serverless: bool = True) -> str:
         """Create or update Switch job, returning job ID."""
-        job_settings = self._get_switch_job_settings()
+        job_settings = self._get_switch_job_settings(use_serverless)
 
         # Try to update existing job
         if job_id:
@@ -149,7 +150,7 @@ def _create_or_update_switch_job(self, job_id: str | None) -> str:
         assert new_job_id is not None
         return new_job_id
 
-    def _get_switch_job_settings(self) -> dict[str, Any]:
+    def _get_switch_job_settings(self, use_serverless: bool = True) -> dict[str, Any]:
         """Build job settings for Switch transpiler."""
         job_name = "Lakebridge_Switch"
         notebook_path = self._get_switch_workspace_path() / "notebooks" / "00_main"
@@ -160,14 +161,31 @@ def _get_switch_job_settings(self) -> dict[str, Any]:
             disable_auto_optimization=True,  # To disable retries on failure
         )
 
-        return {
+        settings: dict[str, Any] = {
             "name": job_name,
             "tags": {"created_by": self._ws.current_user.me().user_name, "switch_version": f"v{switch_version}"},
             "tasks": [task],
             "parameters": self._generate_switch_job_parameters(),
             "max_concurrent_runs": 100,  # Allow simultaneous transpilations
         }
 
+        if not use_serverless:
+            job_cluster_key = "Switch_Cluster"
+            settings["job_clusters"] = [
+                JobCluster(
+                    job_cluster_key=job_cluster_key,
+                    new_cluster=compute.ClusterSpec(
+                        spark_version=self._ws.clusters.select_spark_version(latest=True, long_term_support=True),
+                        node_type_id=self._ws.clusters.select_node_type(local_disk=True, min_memory_gb=16),
+                        num_workers=1,
+                        data_security_mode=compute.DataSecurityMode.USER_ISOLATION,
+                    ),
+                )
+            ]
+            task.job_cluster_key = job_cluster_key
+
+        return settings
+
     @staticmethod
     def _generate_switch_job_parameters() -> Sequence[JobParameterDefinition]:
         # Add required runtime parameters, static for now.

src/databricks/labs/lakebridge/install.py

@@ -57,6 +57,7 @@ def __init__(  # pylint: disable=too-many-arguments
         *,
         is_interactive: bool = True,
         include_llm: bool = False,
+        switch_use_serverless: bool = True,
         transpiler_repository: TranspilerRepository = TranspilerRepository.user_home(),
         transpiler_installers: Sequence[Callable[[TranspilerRepository], TranspilerInstaller]] = (
             BladebridgeInstaller,
@@ -73,6 +74,7 @@ def __init__(  # pylint: disable=too-many-arguments
         # TODO: Refactor the 'prompts' property in preference to using this flag, which should be redundant.
         self._is_interactive = is_interactive
         self._include_llm = include_llm
+        self._switch_use_serverless = switch_use_serverless
         self._transpiler_repository = transpiler_repository
         self._transpiler_installer_factories = transpiler_installers
 
@@ -142,7 +144,10 @@ def configure(self, module: str) -> LakebridgeConfiguration:
             case "transpile":
                 logger.info("Configuring lakebridge `transpile`.")
                 return LakebridgeConfiguration(
-                    self._configure_transpile(), reconcile=None, include_switch=self._include_llm
+                    self._configure_transpile(),
+                    reconcile=None,
+                    include_switch=self._include_llm,
+                    switch_use_serverless=self._switch_use_serverless,
                 )
             case "reconcile":
                 logger.info("Configuring lakebridge `reconcile`.")
@@ -153,6 +158,7 @@ def configure(self, module: str) -> LakebridgeConfiguration:
                     self._configure_transpile(),
                     self._configure_reconcile(),
                     include_switch=self._include_llm,
+                    switch_use_serverless=self._switch_use_serverless,
                 )
             case _:
                 raise ValueError(f"Invalid input: {module}")
@@ -410,6 +416,7 @@ def installer(
     *,
     is_interactive: bool,
     include_llm: bool = False,
+    switch_use_serverless: bool = True,
 ) -> WorkspaceInstaller:
     app_context = ApplicationContext(_verify_workspace_client(ws))
     return WorkspaceInstaller(
@@ -423,6 +430,7 @@ def installer(
         transpiler_repository=transpiler_repository,
         is_interactive=is_interactive,
         include_llm=include_llm,
+        switch_use_serverless=switch_use_serverless,
     )
 
 

tests/unit/deployment/test_switch.py

@@ -10,6 +10,7 @@
 from databricks.labs.lakebridge.deployment.switch import SwitchDeployment
 from databricks.sdk import WorkspaceClient, JobsExt
 from databricks.sdk.errors import NotFound, InvalidParameterValue
+from databricks.sdk.service import compute
 from databricks.sdk.service.jobs import CreateResponse
 from databricks.sdk.service.iam import User
 
@@ -319,3 +320,62 @@ def test_install_creates_new_job_on_update_failure(
     assert install_state.jobs["Switch"] == str(expected_job_id)
     mock_workspace_client.jobs.reset.assert_called_once()
     mock_workspace_client.jobs.create.assert_called_once()
+
+
+# Tests for serverless vs classic cluster selection
+
+
+def test_install_with_serverless_creates_job_without_cluster_key(
+    switch_deployment: SwitchDeployment, mock_workspace_client: Any
+) -> None:
+    """Test serverless mode does not include job_clusters in settings."""
+    new_job = CreateResponse(job_id=123)
+    mock_workspace_client.jobs.create.return_value = new_job
+
+    switch_deployment.install(use_serverless=True)
+
+    call_kwargs = mock_workspace_client.jobs.create.call_args.kwargs
+    assert "job_clusters" not in call_kwargs
+    assert call_kwargs["tasks"][0].job_cluster_key is None
+
+
+def test_install_with_classic_cluster_creates_job_with_cluster_key(
+    switch_deployment: SwitchDeployment, mock_workspace_client: Any
+) -> None:
+    """Test classic cluster mode includes job_clusters and job_cluster_key."""
+    new_job = CreateResponse(job_id=123)
+    mock_workspace_client.jobs.create.return_value = new_job
+    mock_workspace_client.clusters.select_spark_version.return_value = "15.4.x-scala2.12"
+    mock_workspace_client.clusters.select_node_type.return_value = "m5.xlarge"
+
+    switch_deployment.install(use_serverless=False)
+
+    call_kwargs = mock_workspace_client.jobs.create.call_args.kwargs
+    assert "job_clusters" in call_kwargs
+    assert len(call_kwargs["job_clusters"]) == 1
+    assert call_kwargs["job_clusters"][0].job_cluster_key == "Switch_Cluster"
+    assert call_kwargs["tasks"][0].job_cluster_key == "Switch_Cluster"
+
+
+def test_install_with_classic_cluster_configures_correct_cluster_spec(
+    switch_deployment: SwitchDeployment, mock_workspace_client: Any
+) -> None:
+    """Test classic cluster mode configures ClusterSpec correctly."""
+    new_job = CreateResponse(job_id=123)
+    mock_workspace_client.jobs.create.return_value = new_job
+    mock_workspace_client.clusters.select_spark_version.return_value = "15.4.x-scala2.12"
+    mock_workspace_client.clusters.select_node_type.return_value = "m5.xlarge"
+
+    switch_deployment.install(use_serverless=False)
+
+    call_kwargs = mock_workspace_client.jobs.create.call_args.kwargs
+    cluster_spec = call_kwargs["job_clusters"][0].new_cluster
+
+    assert cluster_spec.spark_version == "15.4.x-scala2.12"
+    assert cluster_spec.node_type_id == "m5.xlarge"
+    assert cluster_spec.num_workers == 1
+    assert cluster_spec.data_security_mode == compute.DataSecurityMode.USER_ISOLATION
+
+    # Verify cluster selection methods were called with correct parameters
+    mock_workspace_client.clusters.select_spark_version.assert_called_once_with(latest=True, long_term_support=True)
+    mock_workspace_client.clusters.select_node_type.assert_called_once_with(local_disk=True, min_memory_gb=16)

tests/unit/test_cli_other.py

@@ -1,4 +1,5 @@
 import io
+import os
 from unittest.mock import Mock, patch, MagicMock, create_autospec, PropertyMock
 
 import pytest
@@ -102,3 +103,32 @@ def test_prompts_question():
     prompts = MockPrompts({"Some question": "something"})
     response = option.prompt_for_value(prompts)
     assert response == "something"
+
+
+@pytest.mark.parametrize(
+    ("env_value", "expected_serverless"),
+    (
+        (None, True),  # Default: serverless (no env var)
+        ("", True),  # Empty string: serverless
+        ("CLASSIC", False),  # CLASSIC: use classic cluster
+        ("classic", False),  # lowercase: use classic cluster
+        ("Classic", False),  # Mixed case: use classic cluster
+        ("SERVERLESS", True),  # Any other value: serverless
+        ("other", True),  # Any other value: serverless
+    ),
+)
+def test_install_transpile_cluster_type_env_var(
+    env_value: str | None,
+    expected_serverless: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Test LAKEBRIDGE_CLUSTER_TYPE env var determines serverless vs classic cluster."""
+    if env_value is None:
+        monkeypatch.delenv("LAKEBRIDGE_CLUSTER_TYPE", raising=False)
+    else:
+        monkeypatch.setenv("LAKEBRIDGE_CLUSTER_TYPE", env_value)
+
+    # The logic from cli.py install_transpile:
+    # switch_use_serverless = os.environ.get("LAKEBRIDGE_CLUSTER_TYPE", "").upper() != "CLASSIC"
+    actual = os.environ.get("LAKEBRIDGE_CLUSTER_TYPE", "").upper() != "CLASSIC"
+    assert actual == expected_serverless

tests/unit/test_install.py

@@ -1726,3 +1726,38 @@ def test_transpiler_installers_llm_flag(
     )
     assert installer.configure("transpile").include_switch == should_include_switch
     assert installer.configure("all").include_switch == should_include_switch
+
+
+@pytest.mark.parametrize(
+    ("switch_use_serverless", "expected_serverless"),
+    (
+        (True, True),  # Default: use serverless
+        (False, False),  # Flag disabled: use classic cluster
+    ),
+)
+def test_configure_switch_use_serverless_flag(
+    ws_installer: Callable[..., WorkspaceInstaller],
+    ws: WorkspaceClient,
+    switch_use_serverless: bool,
+    expected_serverless: bool,
+) -> None:
+    """Test switch_use_serverless configuration flag is passed through correctly."""
+    ctx = ApplicationContext(ws).replace(
+        product_info=ProductInfo.for_testing(LakebridgeConfiguration),
+        prompts=MockPrompts({}),
+        installation=MockInstallation({}),
+    )
+    installer = ws_installer(
+        ctx.workspace_client,
+        ctx.prompts,
+        ctx.installation,
+        ctx.install_state,
+        ctx.product_info,
+        ctx.resource_configurator,
+        ctx.workspace_installation,
+        is_interactive=False,
+        include_llm=True,
+        switch_use_serverless=switch_use_serverless,
+    )
+    config = installer.configure("transpile")
+    assert config.switch_use_serverless == expected_serverless