Repository: mlflow/mlflow

Issue:
[BUG] mlflow projects run ignores `docker_args` for docker container environments if `synchronous` is disabled
### Willingness to contribute
The MLflow Community encourages bug fix contributions. Would you or another member of your organization be willing to contribute a fix for this bug to the MLflow code base?

- [x] Yes. I can contribute a fix for this bug independently.
- [ ] Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.
- [ ] No. I cannot contribute a bug fix at this time.

PR: https://github.com/mlflow/mlflow/pull/3563/

### System information
- **Have I written custom code (as opposed to using a stock example script provided in MLflow)**: No
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Ubuntu 18.04.5 LTS
- **MLflow installed from (source or binary)**: binary (`pip install mlflow`)
- **MLflow version (run ``mlflow --version``)**: 1.11.0
- **Python version**: 3.6.9
- **npm version, if running the dev UI**:
- **Exact command to reproduce**:

### Describe the problem
MLFlow projects run ignores `docker_args` for docker container environments if `synchronous` is disabled. The expectation is the behavior for docker container environment to be the same weather the run is synchronous or not.

Also, see https://github.com/mlflow/mlflow/issues/3372

### Code to reproduce issue
Provide a reproducible test case that is the bare minimum necessary to generate the problem.
```bash
git clone https://github.com/mlflow/mlflow.git
cd ~/mlflow

# build docker image:
docker build -t mlflow-docker-example -f Dockerfile examples/docker
```

Minimal python app:

```python
import mlflow

mlflow.projects.run(
    uri='examples/docker',
    parameters={"alpha": 0.5},
    synchronous=False,
    docker_args={"name": "mlflow-container-run"},
)
```

Current behavior:
```
$ docker container ls

CONTAINER ID        IMAGE                   COMMAND                  CREATED             STATUS                  PORTS                    NAMES
14f5012f2a53        docker-example:latest   "/usr/bin/tini -- py…"   2 seconds ago       Up Less than a second                            serene_gould
```
Container name is random (not set)


Expected behavior:
```
$ docker container ls

CONTAINER ID        IMAGE                   COMMAND                  CREATED             STATUS              PORTS                    NAMES
d4ccfb3ec5dc        docker-example:latest   "/usr/bin/tini -- py…"   2 seconds ago       Up 1 second                                  mlflow-container-run
```
Container name is set as passed in `docker_args` 

### What component(s), interfaces, languages, and integrations does this bug affect?
Components 
- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [x] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: Local serving, model deployment tools, spark UDFs
- [ ] `area/server-infra`: MLflow server, JavaScript dev server
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

Interface 
- [ ] `area/uiux`: Front-end, user experience, JavaScript, plotting
- [x] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

Language 
- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

Integrations
- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations



Discussion:

Resolving patch:
diff --git a/mlflow/projects/backend/local.py b/mlflow/projects/backend/local.py
--- a/mlflow/projects/backend/local.py
+++ b/mlflow/projects/backend/local.py
@@ -104,13 +104,14 @@ def run(
             parameters=params,
             experiment_id=experiment_id,
             use_conda=use_conda,
+            docker_args=docker_args,
             storage_dir=storage_dir,
             run_id=active_run.info.run_id,
         )
 
 
 def _invoke_mlflow_run_subprocess(
-    work_dir, entry_point, parameters, experiment_id, use_conda, storage_dir, run_id
+    work_dir, entry_point, parameters, experiment_id, use_conda, docker_args, storage_dir, run_id
 ):
     """
     Run an MLflow project asynchronously by invoking ``mlflow run`` in a subprocess, returning
@@ -120,6 +121,7 @@ def _invoke_mlflow_run_subprocess(
     mlflow_run_arr = _build_mlflow_run_cmd(
         uri=work_dir,
         entry_point=entry_point,
+        docker_args=docker_args,
         storage_dir=storage_dir,
         use_conda=use_conda,
         run_id=run_id,
@@ -131,12 +133,18 @@ def _invoke_mlflow_run_subprocess(
     return LocalSubmittedRun(run_id, mlflow_run_subprocess)
 
 
-def _build_mlflow_run_cmd(uri, entry_point, storage_dir, use_conda, run_id, parameters):
+def _build_mlflow_run_cmd(
+    uri, entry_point, docker_args, storage_dir, use_conda, run_id, parameters
+):
     """
     Build and return an array containing an ``mlflow run`` command that can be invoked to locally
     run the project at the specified URI.
     """
     mlflow_run_arr = ["mlflow", "run", uri, "-e", entry_point, "--run-id", run_id]
+    if docker_args is not None:
+        for key, value in docker_args.items():
+            args = key if isinstance(value, bool) else "%s=%s" % (key, value)
+            mlflow_run_arr.extend(["--docker-args", args])
     if storage_dir is not None:
         mlflow_run_arr.extend(["--storage-dir", storage_dir])
     if not use_conda:


Test patch:
diff --git a/tests/projects/test_docker_projects.py b/tests/projects/test_docker_projects.py
--- a/tests/projects/test_docker_projects.py
+++ b/tests/projects/test_docker_projects.py
@@ -41,6 +41,7 @@ def test_docker_project_execution(
         experiment_id=file_store.FileStore.DEFAULT_EXPERIMENT_ID,
         parameters=expected_params,
         entry_point="test_tracking",
+        docker_args={"memory": "1g", "privileged": True},
     )
     # Validate run contents in the FileStore
     run_id = submitted_run.run_id
@@ -69,6 +70,31 @@ def test_docker_project_execution(
         assert run_tags[k].startswith(v)
     artifacts = mlflow_service.list_artifacts(run_id=run_id)
     assert len(artifacts) == 1
+    docker_cmd = submitted_run.command_proc.args[2]
+    assert "--memory 1g" in docker_cmd
+    assert "--privileged" in docker_cmd
+
+
+@pytest.mark.large
+def test_docker_project_execution_async_docker_args(
+    tmpdir, docker_example_base_image
+):  # pylint: disable=unused-argument
+    submitted_run = mlflow.projects.run(
+        TEST_DOCKER_PROJECT_DIR,
+        experiment_id=file_store.FileStore.DEFAULT_EXPERIMENT_ID,
+        parameters={"use_start_run": "0"},
+        entry_point="test_tracking",
+        docker_args={"memory": "1g", "privileged": True},
+        synchronous=False,
+    )
+    submitted_run.wait()
+
+    args = submitted_run.command_proc.args
+    assert len([a for a in args if a == "--docker-args"]) == 2
+    first_idx = args.index("--docker-args")
+    second_idx = args.index("--docker-args", first_idx + 1)
+    assert args[first_idx + 1] == "memory=1g"
+    assert args[second_idx + 1] == "privileged"
 
 
 @pytest.mark.parametrize(