Repository: mlflow/mlflow

Issue:
[FR] Add support for cluster library definition in MLflow project execution
Hi!

MLFlow redefines library dependencies on databricks cluster creation. In my project, extra libraries are needed on cluster nodes, but when I'm specifying them in cluster json file, they are ignored and replaced by the single mlflow dependency. 

I suppose this happens in this line of code: https://github.com/mlflow/mlflow/blob/58a6053cad8d6c01f095a4e3dac8a1b560dc4a61/mlflow/projects/databricks.py#L182

This is not very flexible and doesn't allow me to install extra libraries on mlflow cluster nodes.
So, better behaviour would be to add mlflow dependency to dependency given in cluster specification file.

Below is the example of my cluster.json:
```
{
  "spark_conf": {
    "spark.databricks.hive.metastore.glueCatalog.enabled": "true"
  },
  "libraries": [
    {
      "whl": "dbfs:/path_to_my_lib.whl"
    }
  ],
  "num_workers": 20
}
```

But my whl dependency is ignored on cluster creation. Below is the resulting cluster properties:
```
  "cluster_spec": {
    "new_cluster": {
    "libraries": [
      {
        "pypi": {
          "package": "'mlflow<=1.7.2'"
        }
      }
    ]
  },
```
Thanks!


Discussion:
@Shmuma thank you for raising this issue and finding relevant parts of our implementation!

The `conda-env` parameter in MLproject is intended to capture the dependencies required to run code in MLflow in a reproducible way. Needing a wheel in dbfs breaks the reproducibility across different environments, but perhaps that's okay in your context and could be managed outside of the MLproject format.

With regards to the behavior of the running MLprojects on Databricks, there are two relevant limitations here:
- the conda environment is only installed on the driver
- there's no easy way to reference a wheel from a conda environment

In your example, `libraries` is not a recognized field in the field that is accepted in the command (https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobsclusterspecnewcluster). In this case, your example wouldn't be expected to work. 

This would need some additional design on how to pass additional libraries to the cluster specification that gets constructed, as it would have to be defined outside the `cluster.json`. We'd also want to think through its impact on reproducibility. 

One thing that would be helpful for us is to understand the use case a bit better - what aspects of the MLproject formats are you particularly looking to leverage that go beyond what the Databrick's job APIs provide? 



Hi!

Thanks for your reply!

My requirements is quite simple. The code I'm writing requires some third-party libraries to be installed on computing nodes, not only on the driver's side. 

For example, we're using maxmind GeoLocation libraries (maxminddb) to perform large-scale preprocessing, so, library maxminddb need to be present on workers. That's not the only example, but rather quite generic requirements to be able to specify cluster's workers configuration.

When I'm using mlflow via databricks-connect, I'm specifying already configured cluster, so, it's not a problem to install any set of packages. I want to be able to do the same when I'm running mlflow job with `-b databricks -c cluster.json` mode.

In theory, it is possible to solve in a hacky way by running `pip install` or doing `SparkContext.addFile`, but that becomes a mess in a large project with several dependencies.

So, ideal solution will be to have a way to specify set of packages to be installed on a cluster before mlflow job run.
Thanks for the extra context @Shmuma! This may take a bit of design on our side and prioritization to surface this all the way up through to the MLflow command line for project usage. I'll try to get some clarity with the team over the next couple days when we could address it.

One potential workaround would be to use the Databricks jobs rest API. `runs-submit` (https://docs.databricks.com/dev-tools/api/latest/jobs.html#runs-submit) allows you to create a cluster and run a job on the cluster with whatever libraries you need installed on the workers. You should be able to log to MLflow from these jobs, you'll just need to be sure to include MLflow in the libraries if you are using the Databricks Runtime. MLflow is installed on the Databricks Runtime for Machine Learning. Would this unblock you while we work on a more integrated solution?


Thanks for suggesting!

I've quickly skimmed through jobs api documentation, but found it quite limited in comparison to mlflow functionality. In particular:

- haven't found a way to setup conda env on driver
- jobs api supports "spark_python_task", but it is limited to a single file, not a whole dir or git repo
- spark_python_task documentation (https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobssparkpythontask) says libraries are not supported in this mode, only pyFiles, which is, again won't work for worker's libraries.

I don't have much experience with jobs api and I might be missing something. But from what I've seen, it's not a solution of the issue. 


Hi @Shmuma - 

There are a few approaches we could use here. Let assuming that you don't want to wait for (3), let me know which approach you would like to pursue among (1) and (2).

**Approach 1.** With the `spark_python_task`, library management for the workers and drivers would all be through the `libraries` block of the cluster specification, and would have a conda environment on the driver. The warning you see refers to the `spark_submit_task` task type, which is a different task type than the `spark_python_task`. However, your understanding is correct that the job is limited to referencing a single python file. You'd have to bundle up other code you depend on into a wheel or library and install that on the cluster.

**Approach 2.** We're also considering on a patch to address this in the MLflow client to unblock this use case relatively quickly. It would have a quirk. You would need to specify the libraries you depend on in both the conda environment and the `libraries` property for cluster configuration, in slightly different ways. As you have a wheel in your previous example, I'm assuming you have a wheel uploaded to 

To install wheels on the driver, in your conda yaml file:
```
dependencies:
  - pip:
    - /dbfs/path_to_my_lib.whl
```
Notice that the path to the wheel is not of the form `dbfs:/`, but instead uses the FUSE mount in `/dbfs` to make the file appear local. For other libraries to install in the driver, you would specify them using standard conda conventions.

To install wheels on the workers, you would need to specify them in the cluster json, in a slightly different format than what you've been using:
```
{
  "new_cluster": {
    "spark_version": "6.5.x-scala2.11",
    "num_workers": 1,
    "node_type_id": "i3.xlarge"
  },
  "libraries": [
    {
      "whl": "dbfs:/path_to_my_lib.whl"
    }
  ]
}
```
To add other additional libraries you would use the [standard Databricks JSON format](https://docs.databricks.com/dev-tools/api/latest/libraries.html#managedlibrarieslibrary) for libraries in APIs.

It's a bit ugly to have to specify the libraries for workers and the driver separately in two different formats, but we could get a patch into MLflow in the next few days. To use it you'd have to wait until the next release (which may be up to a month from now) or build MLflow yourself from master (as described [here](https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.rst#prerequisites), except you could probably get away without installing the test requirements).

**Approach 3.** In the long run we would like the workers to respect the conda environment too, but that would be a much longer project to get released and available to you, likely on the order of several months from now.

Hi Paul!

Thanks for your reply. From my perspective Approach 2 is the best solution to my problem, I'm glad it won't take too much time to implement it. 

Meanwhile, I found a workaround my issue by running `pip install package` on workers first time I'm importing required packages. That's a bit hacky and could not be seen as a long-term production solution, but it works.

In terms of Approach 3, my suggestion will be to allow separate conda configuration file for workers, as, it is quite likely that driver and workes might have a different set of packages needed. 

In any case, thanks for your help! Hope for the proper solution to be implemented soon!

Wbr, Max

Resolving patch:
diff --git a/mlflow/projects/databricks.py b/mlflow/projects/databricks.py
--- a/mlflow/projects/databricks.py
+++ b/mlflow/projects/databricks.py
@@ -162,12 +162,28 @@ def _run_shell_command_job(self, project_uri, command, env_vars, cluster_spec):
         :param command: Shell command to run.
         :param env_vars: Environment variables to set in the process running ``command``.
         :param cluster_spec: Dictionary containing a `Databricks cluster specification
-                             <https://docs.databricks.com/api/latest/jobs.html#clusterspec>`_
-                             to use when launching a run.
+                             <https://docs.databricks.com/dev-tools/api/latest/jobs.html#clusterspec>`_
+                             or a `Databricks new cluster specification
+                             <https://docs.databricks.com/dev-tools/api/latest/jobs.html#jobsclusterspecnewcluster>`_
+                             to use when launching a run. If you specify libraries, this function
+                             will add MLflow to the library list. This function does not support
+                             installation of conda environment libraries on the workers.
         :return: ID of the Databricks job run. Can be used to query the run's status via the
                  Databricks
                  `Runs Get <https://docs.databricks.com/api/latest/jobs.html#runs-get>`_ API.
         """
+        # NB: We use <= on the version specifier to allow running projects on pre-release
+        # versions, where we will select the most up-to-date mlflow version available.
+        # Also note, that we escape this so '<' is not treated as a shell pipe.
+        libraries = [{"pypi": {"package": "'mlflow<=%s'" % VERSION}}]
+
+        # Check syntax of JSON - if it contains libraries and new_cluster, pull those out
+        if 'new_cluster' in cluster_spec:
+            # Libraries are optional, so we don't require that this be specified
+            if 'libraries' in cluster_spec:
+                libraries.extend(cluster_spec['libraries'])
+            cluster_spec = cluster_spec['new_cluster']
+
         # Make jobs API request to launch run.
         req_body_json = {
             'run_name': 'MLflow Run for %s' % project_uri,
@@ -176,10 +192,7 @@ def _run_shell_command_job(self, project_uri, command, env_vars, cluster_spec):
                 'command': command,
                 "env_vars": env_vars
             },
-            # NB: We use <= on the version specifier to allow running projects on pre-release
-            # versions, where we will select the most up-to-date mlflow version available.
-            # Also note, that we escape this so '<' is not treated as a shell pipe.
-            "libraries": [{"pypi": {"package": "'mlflow<=%s'" % VERSION}}],
+            "libraries": libraries,
         }
         run_submit_res = self._jobs_runs_submit(req_body_json)
         databricks_run_id = run_submit_res["run_id"]


Test patch:
diff --git a/tests/projects/test_databricks.py b/tests/projects/test_databricks.py
--- a/tests/projects/test_databricks.py
+++ b/tests/projects/test_databricks.py
@@ -268,6 +268,58 @@ def test_run_databricks_cluster_spec_json(
         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):
+    with mock.patch.dict(os.environ, {'DATABRICKS_HOST': 'test-host', '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):
+    with mock.patch.dict(os.environ, {'DATABRICKS_HOST': 'test-host', '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():
     with mock.patch.dict(os.environ, {'DATABRICKS_HOST': 'test-host', 'DATABRICKS_TOKEN': 'foo'}):
         existing_cluster_spec = {