chore: import upstream snapshot with attribution

tests/projects/test_databricks.py

@@ -0,0 +1,498 @@
+import filecmp
+import json
+import os
+import shutil
+from unittest import mock
+
+import pytest
+
+import mlflow
+from mlflow import MlflowClient, cli
+from mlflow.entities import RunStatus
+from mlflow.environment_variables import MLFLOW_TRACKING_URI
+from mlflow.exceptions import MlflowException
+from mlflow.legacy_databricks_cli.configure.provider import DatabricksConfig
+from mlflow.projects import ExecutionException, databricks
+from mlflow.projects.databricks import DatabricksJobRunner, _get_cluster_mlflow_run_cmd
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, ErrorCode
+from mlflow.store.tracking.file_store import FileStore
+from mlflow.tracking.request_header.default_request_header_provider import (
+    DefaultRequestHeaderProvider,
+)
+from mlflow.utils import file_utils
+from mlflow.utils.mlflow_tags import (
+    MLFLOW_DATABRICKS_RUN_URL,
+    MLFLOW_DATABRICKS_SHELL_JOB_RUN_ID,
+    MLFLOW_DATABRICKS_WEBAPP_URL,
+)
+from mlflow.utils.rest_utils import MlflowHostCreds
+from mlflow.utils.uri import construct_db_uri_from_profile
+
+from tests import helper_functions
+from tests.integration.utils import invoke_cli_runner
+from tests.projects.utils import TEST_PROJECT_DIR, validate_exit_status
+
+
+@pytest.fixture
+def runs_cancel_mock():
+    """Mocks the Jobs Runs Cancel API request"""
+    with mock.patch(
+        "mlflow.projects.databricks.DatabricksJobRunner.jobs_runs_cancel"
+    ) as runs_cancel_mock:
+        runs_cancel_mock.return_value = None
+        yield runs_cancel_mock
+
+
+@pytest.fixture
+def runs_submit_mock():
+    """Mocks the Jobs Runs Submit API request"""
+    with mock.patch(
+        "mlflow.projects.databricks.DatabricksJobRunner._jobs_runs_submit",
+        return_value={"run_id": "-1"},
+    ) as runs_submit_mock:
+        yield runs_submit_mock
+
+
+@pytest.fixture
+def runs_get_mock():
+    """Mocks the Jobs Runs Get API request"""
+    with mock.patch(
+        "mlflow.projects.databricks.DatabricksJobRunner.jobs_runs_get"
+    ) as runs_get_mock:
+        yield runs_get_mock
+
+
+@pytest.fixture
+def databricks_cluster_mlflow_run_cmd_mock():
+    """Mocks the Jobs Runs Get API request"""
+    with mock.patch(
+        "mlflow.projects.databricks._get_cluster_mlflow_run_cmd"
+    ) as mlflow_run_cmd_mock:
+        yield mlflow_run_cmd_mock
+
+
+@pytest.fixture
+def cluster_spec_mock(tmp_path):
+    cluster_spec_handle = tmp_path.joinpath("cluster_spec.json")
+    cluster_spec_handle.write_text("{}")
+    return str(cluster_spec_handle)
+
+
+@pytest.fixture
+def dbfs_root_mock(tmp_path):
+    return str(tmp_path.joinpath("dbfs-root"))
+
+
+@pytest.fixture
+def upload_to_dbfs_mock(dbfs_root_mock):
+    def upload_mock_fn(_, src_path, dbfs_uri):
+        mock_dbfs_dst = os.path.join(dbfs_root_mock, dbfs_uri.split("/dbfs/")[1])
+        os.makedirs(os.path.dirname(mock_dbfs_dst))
+        shutil.copy(src_path, mock_dbfs_dst)
+
+    with mock.patch.object(
+        mlflow.projects.databricks.DatabricksJobRunner, "_upload_to_dbfs", new=upload_mock_fn
+    ) as upload_mock:
+        yield upload_mock
+
+
+@pytest.fixture
+def dbfs_path_exists_mock(dbfs_root_mock):
+    with mock.patch(
+        "mlflow.projects.databricks.DatabricksJobRunner._dbfs_path_exists"
+    ) as path_exists_mock:
+        yield path_exists_mock
+
+
+@pytest.fixture
+def dbfs_mocks(dbfs_path_exists_mock, upload_to_dbfs_mock):
+    return
+
+
+@pytest.fixture
+def before_run_validations_mock():
+    with mock.patch("mlflow.projects.databricks.before_run_validations"):
+        yield
+
+
+@pytest.fixture
+def set_tag_mock():
+    with mock.patch("mlflow.projects.databricks.tracking.MlflowClient") as m:
+        mlflow_service_mock = mock.Mock(wraps=MlflowClient())
+        m.return_value = mlflow_service_mock
+        yield mlflow_service_mock.set_tag
+
+
+def _get_mock_run_state(succeeded):
+    if succeeded is None:
+        return {"life_cycle_state": "RUNNING", "state_message": ""}
+    run_result_state = "SUCCESS" if succeeded else "FAILED"
+    return {"life_cycle_state": "TERMINATED", "state_message": "", "result_state": run_result_state}
+
+
+def mock_runs_get_result(succeeded):
+    run_state = _get_mock_run_state(succeeded)
+    return {"state": run_state, "run_page_url": "test_url"}
+
+
+def run_databricks_project(cluster_spec, **kwargs):
+    return mlflow.projects.run(
+        uri=TEST_PROJECT_DIR,
+        backend="databricks",
+        backend_config=cluster_spec,
+        parameters={"alpha": "0.4"},
+        **kwargs,
+    )
+
+
+def test_upload_project_to_dbfs(
+    dbfs_root_mock, tmp_path, dbfs_path_exists_mock, upload_to_dbfs_mock
+):
+    # Upload project to a mock directory
+    dbfs_path_exists_mock.return_value = False
+    runner = DatabricksJobRunner(databricks_profile_uri=construct_db_uri_from_profile("DEFAULT"))
+    dbfs_uri = runner._upload_project_to_dbfs(
+        project_dir=TEST_PROJECT_DIR, experiment_id=FileStore.DEFAULT_EXPERIMENT_ID
+    )
+    # Get expected tar
+    local_tar_path = os.path.join(dbfs_root_mock, dbfs_uri.split("/dbfs/")[1])
+    expected_tar_path = str(tmp_path.joinpath("expected.tar.gz"))
+    file_utils.make_tarfile(
+        output_filename=expected_tar_path,
+        source_dir=TEST_PROJECT_DIR,
+        archive_name=databricks.DB_TARFILE_ARCHIVE_NAME,
+    )
+    # Extract the tarred project, verify its contents
+    assert filecmp.cmp(local_tar_path, expected_tar_path, shallow=False)
+
+
+def test_upload_existing_project_to_dbfs(dbfs_path_exists_mock):
+    # Check that we don't upload the project if it already exists on DBFS
+    with mock.patch(
+        "mlflow.projects.databricks.DatabricksJobRunner._upload_to_dbfs"
+    ) as upload_to_dbfs_mock:
+        dbfs_path_exists_mock.return_value = True
+        runner = DatabricksJobRunner(
+            databricks_profile_uri=construct_db_uri_from_profile("DEFAULT")
+        )
+        runner._upload_project_to_dbfs(
+            project_dir=TEST_PROJECT_DIR, experiment_id=FileStore.DEFAULT_EXPERIMENT_ID
+        )
+        assert upload_to_dbfs_mock.call_count == 0
+
+
+@pytest.mark.parametrize(
+    "response_mock",
+    [
+        helper_functions.create_mock_response(400, "Error message but not a JSON string"),
+        helper_functions.create_mock_response(400, ""),
+        helper_functions.create_mock_response(400, None),
+    ],
+)
+def test_dbfs_path_exists_error_response_handling(response_mock):
+    with (
+        mock.patch(
+            "mlflow.utils.databricks_utils.get_databricks_host_creds"
+        ) as get_databricks_host_creds_mock,
+        mock.patch("mlflow.utils.rest_utils.http_request") as http_request_mock,
+    ):
+        # given a well formed DatabricksJobRunner
+        # note: databricks_profile is None needed because clients using profile are mocked
+        job_runner = DatabricksJobRunner(databricks_profile_uri=None)
+
+        # when the http request to validate the dbfs path returns a 400 response with an
+        # error message that is either well-formed JSON or not
+        get_databricks_host_creds_mock.return_value = None
+        http_request_mock.return_value = response_mock
+
+        # then _dbfs_path_exists should return a MlflowException
+        with pytest.raises(MlflowException, match="API request to check existence of file at DBFS"):
+            job_runner._dbfs_path_exists("some/path")
+
+
+def test_run_databricks_validations(
+    tmp_path,
+    monkeypatch,
+    cluster_spec_mock,
+    dbfs_mocks,
+    set_tag_mock,
+):
+    """
+    Tests that running on Databricks fails before making any API requests if validations fail.
+    """
+    monkeypatch.setenv("DATABRICKS_HOST", "test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    with mock.patch(
+        "mlflow.projects.databricks.DatabricksJobRunner._databricks_api_request"
+    ) as db_api_req_mock:
+        # Test bad tracking URI
+        mlflow.set_tracking_uri(f"sqlite:///{tmp_path / 'mlflow.db'}")
+        with pytest.raises(ExecutionException, match="MLflow tracking URI must be of"):
+            run_databricks_project(cluster_spec_mock, synchronous=True)
+        assert db_api_req_mock.call_count == 0
+        db_api_req_mock.reset_mock()
+        mlflow_service = MlflowClient()
+        assert len(mlflow_service.search_runs([FileStore.DEFAULT_EXPERIMENT_ID])) == 0
+        mlflow.set_tracking_uri("databricks")
+        # Test misspecified parameters
+        with pytest.raises(
+            ExecutionException, match="No value given for missing parameters: 'name'"
+        ):
+            mlflow.projects.run(
+                TEST_PROJECT_DIR,
+                backend="databricks",
+                entry_point="greeter",
+                backend_config=cluster_spec_mock,
+            )
+        assert db_api_req_mock.call_count == 0
+        db_api_req_mock.reset_mock()
+        # Test bad cluster spec
+        with pytest.raises(ExecutionException, match="Backend spec must be provided"):
+            mlflow.projects.run(
+                TEST_PROJECT_DIR, backend="databricks", synchronous=True, backend_config=None
+            )
+        assert db_api_req_mock.call_count == 0
+        db_api_req_mock.reset_mock()
+        # Test that validations pass with good tracking URIs
+        databricks.before_run_validations("http://", cluster_spec_mock)
+        databricks.before_run_validations("databricks", cluster_spec_mock)
+
+
+@pytest.mark.usefixtures(
+    "before_run_validations_mock",
+    "runs_cancel_mock",
+    "dbfs_mocks",
+    "databricks_cluster_mlflow_run_cmd_mock",
+)
+def test_run_databricks(
+    runs_submit_mock,
+    runs_get_mock,
+    cluster_spec_mock,
+    set_tag_mock,
+    databricks_cluster_mlflow_run_cmd_mock,
+    monkeypatch,
+):
+    monkeypatch.setenv("DATABRICKS_HOST", "https://test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    mlflow.set_tracking_uri("databricks")
+    # Test that MLflow gets the correct run status when performing a Databricks run
+    for run_succeeded, expect_status in [(True, RunStatus.FINISHED), (False, RunStatus.FAILED)]:
+        runs_get_mock.return_value = mock_runs_get_result(succeeded=run_succeeded)
+        submitted_run = run_databricks_project(cluster_spec_mock, synchronous=False)
+        assert submitted_run.wait() == run_succeeded
+        assert submitted_run.run_id is not None
+        assert runs_submit_mock.call_count == 1
+        assert databricks_cluster_mlflow_run_cmd_mock.call_count == 1
+        tags = {}
+        for call_args, _ in set_tag_mock.call_args_list:
+            tags[call_args[1]] = call_args[2]
+        assert tags[MLFLOW_DATABRICKS_RUN_URL] == "test_url"
+        assert tags[MLFLOW_DATABRICKS_SHELL_JOB_RUN_ID] == "-1"
+        assert tags[MLFLOW_DATABRICKS_WEBAPP_URL] == "https://test-host"
+        set_tag_mock.reset_mock()
+        runs_submit_mock.reset_mock()
+        databricks_cluster_mlflow_run_cmd_mock.reset_mock()
+        validate_exit_status(submitted_run.get_status(), expect_status)
+
+
+@pytest.mark.usefixtures(
+    "before_run_validations_mock",
+    "runs_cancel_mock",
+    "dbfs_mocks",
+    "cluster_spec_mock",
+    "set_tag_mock",
+)
+def test_run_databricks_cluster_spec_json(runs_submit_mock, runs_get_mock, monkeypatch):
+    monkeypatch.setenv("DATABRICKS_HOST", "test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    runs_get_mock.return_value = mock_runs_get_result(succeeded=True)
+    cluster_spec = {
+        "spark_version": "5.0.x-scala2.11",
+        "num_workers": 2,
+        "node_type_id": "i3.xlarge",
+    }
+    # Run project synchronously, verify that it succeeds (doesn't throw)
+    run_databricks_project(cluster_spec=cluster_spec, synchronous=True)
+    assert runs_submit_mock.call_count == 1
+    runs_submit_args, _ = runs_submit_mock.call_args_list[0]
+    req_body = runs_submit_args[0]
+    assert req_body["new_cluster"] == cluster_spec
+
+
+@pytest.mark.usefixtures(
+    "before_run_validations_mock",
+    "runs_cancel_mock",
+    "dbfs_mocks",
+    "cluster_spec_mock",
+    "set_tag_mock",
+)
+def test_run_databricks_extended_cluster_spec_json(runs_submit_mock, runs_get_mock, monkeypatch):
+    monkeypatch.setenv("DATABRICKS_HOST", "test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    runs_get_mock.return_value = mock_runs_get_result(succeeded=True)
+    new_cluster_spec = {
+        "spark_version": "6.5.x-scala2.11",
+        "num_workers": 2,
+        "node_type_id": "i3.xlarge",
+    }
+    extra_library = {"pypi": {"package": "tensorflow"}}
+
+    cluster_spec = {"new_cluster": new_cluster_spec, "libraries": [extra_library]}
+
+    # Run project synchronously, verify that it succeeds (doesn't throw)
+    run_databricks_project(cluster_spec=cluster_spec, synchronous=True)
+    assert runs_submit_mock.call_count == 1
+    runs_submit_args, _ = runs_submit_mock.call_args_list[0]
+    req_body = runs_submit_args[0]
+    assert req_body["new_cluster"] == new_cluster_spec
+    # This does test deep object equivalence
+    assert extra_library in req_body["libraries"]
+
+
+@pytest.mark.usefixtures(
+    "before_run_validations_mock",
+    "runs_cancel_mock",
+    "dbfs_mocks",
+    "cluster_spec_mock",
+    "set_tag_mock",
+)
+def test_run_databricks_extended_cluster_spec_json_without_libraries(
+    runs_submit_mock, runs_get_mock, monkeypatch
+):
+    monkeypatch.setenv("DATABRICKS_HOST", "test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    runs_get_mock.return_value = mock_runs_get_result(succeeded=True)
+    new_cluster_spec = {
+        "spark_version": "6.5.x-scala2.11",
+        "num_workers": 2,
+        "node_type_id": "i3.xlarge",
+    }
+
+    cluster_spec = {
+        "new_cluster": new_cluster_spec,
+    }
+
+    # Run project synchronously, verify that it succeeds (doesn't throw)
+    run_databricks_project(cluster_spec=cluster_spec, synchronous=True)
+    assert runs_submit_mock.call_count == 1
+    runs_submit_args, _ = runs_submit_mock.call_args_list[0]
+    req_body = runs_submit_args[0]
+    assert req_body["new_cluster"] == new_cluster_spec
+
+
+def test_run_databricks_throws_exception_when_spec_uses_existing_cluster(monkeypatch):
+    monkeypatch.setenv("DATABRICKS_HOST", "test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    existing_cluster_spec = {
+        "existing_cluster_id": "1000-123456-clust1",
+    }
+    with pytest.raises(
+        MlflowException, match="execution against existing clusters is not currently supported"
+    ) as exc:
+        run_databricks_project(cluster_spec=existing_cluster_spec)
+    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
+
+
+def test_run_databricks_cancel(
+    before_run_validations_mock,
+    runs_submit_mock,
+    dbfs_mocks,
+    set_tag_mock,
+    runs_cancel_mock,
+    runs_get_mock,
+    cluster_spec_mock,
+    monkeypatch,
+):
+    # Test that MLflow properly handles Databricks run cancellation. We mock the result of
+    # the runs-get API to indicate run failure so that cancel() exits instead of blocking while
+    # waiting for run status.
+    monkeypatch.setenv("DATABRICKS_HOST", "test-host")
+    monkeypatch.setenv("DATABRICKS_TOKEN", "foo")
+    runs_get_mock.return_value = mock_runs_get_result(succeeded=False)
+    submitted_run = run_databricks_project(cluster_spec_mock, synchronous=False)
+    submitted_run.cancel()
+    validate_exit_status(submitted_run.get_status(), RunStatus.FAILED)
+    assert runs_cancel_mock.call_count == 1
+    # Test that we raise an exception when a blocking Databricks run fails
+    runs_get_mock.return_value = mock_runs_get_result(succeeded=False)
+    with pytest.raises(mlflow.projects.ExecutionException, match=r"Run \(ID '.+'\) failed"):
+        run_databricks_project(cluster_spec_mock, synchronous=True)
+
+
+def test_get_tracking_uri_for_run(monkeypatch):
+    mlflow.set_tracking_uri("http://some-uri")
+    assert databricks._get_tracking_uri_for_run() == "http://some-uri"
+    mlflow.set_tracking_uri("databricks://profile")
+    assert databricks._get_tracking_uri_for_run() == "databricks"
+    mlflow.set_tracking_uri(None)
+    monkeypatch.setenv(MLFLOW_TRACKING_URI.name, "http://some-uri")
+    assert mlflow.tracking._tracking_service.utils.get_tracking_uri() == "http://some-uri"
+
+
+class MockProfileConfigProvider:
+    def __init__(self, profile):
+        assert profile == "my-profile"
+
+    def get_config(self):
+        return DatabricksConfig.from_password("host", "user", "pass", insecure=False)
+
+
+def test_databricks_http_request_integration():
+    def confirm_request_params(*args, **kwargs):
+        headers = DefaultRequestHeaderProvider().request_headers()
+        headers["Authorization"] = "Basic dXNlcjpwYXNz"
+        assert args == ("PUT", "host/clusters/list")
+        assert kwargs == {
+            "allow_redirects": True,
+            "headers": headers,
+            "verify": True,
+            "json": {"a": "b"},
+            "timeout": 120,
+        }
+        http_response = mock.MagicMock()
+        http_response.status_code = 200
+        http_response.text = '{"OK": "woo"}'
+        return http_response
+
+    with (
+        mock.patch("requests.Session.request", side_effect=confirm_request_params),
+        mock.patch(
+            "mlflow.utils.databricks_utils.get_databricks_host_creds",
+            return_value=MlflowHostCreds(
+                host="host", username="user", password="pass", ignore_tls_verification=False
+            ),
+        ),
+    ):
+        response = DatabricksJobRunner(databricks_profile_uri=None)._databricks_api_request(
+            "/clusters/list", "PUT", json={"a": "b"}
+        )
+        assert json.loads(response.text) == {"OK": "woo"}
+
+
+def test_run_databricks_failed():
+    text = '{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "Node type not supported"}'
+    with (
+        mock.patch("mlflow.utils.databricks_utils.get_databricks_host_creds"),
+        mock.patch(
+            "mlflow.utils.rest_utils.http_request",
+            return_value=mock.Mock(text=text, status_code=400),
+        ),
+    ):
+        runner = DatabricksJobRunner(construct_db_uri_from_profile("profile"))
+        with pytest.raises(
+            MlflowException, match="RESOURCE_DOES_NOT_EXIST: Node type not supported"
+        ):
+            runner._run_shell_command_job("/project", "command", {}, {})
+
+
+def test_run_databricks_generates_valid_mlflow_run_cmd():
+    cmd = _get_cluster_mlflow_run_cmd(
+        project_dir="my_project_dir",
+        run_id="hi",
+        entry_point="main",
+        parameters={"a": "b"},
+        env_manager="conda",
+    )
+    assert cmd[0] == "mlflow"
+    with mock.patch("mlflow.projects.run"):
+        invoke_cli_runner(cli.cli, cmd[1:])