Repository: mlflow/mlflow

Issue:
[BUG] Cannot configure number of workers of MLServer generated with `mlflow models build-docker` command
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.1.1


### System information

- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Debian 11 16cpus/RAM 32Go **and** Ubuntu 20.04.5 128cpus/RAM 540Go
- **Python version**: 3.9.2 **and** 3.8.10
- **Docker version**: 23.0.1 **and** 20.10.21


### Describe the problem

I am a MLFlow enthusiastic (really great job !) and I am validating MLFlow integration in the company I work for.
We have several big models which work together with around 1.5Go of memory foot print per model.
We opted for MLServer as model server and we use `mlflow models build-docker` to generate dockerized API of models.

But my issue is that we cannot control the number of worker/process MLServer spawns through environment variables (or any other way) because `MLSERVER_INFER_WORKERS` variable environment (see https://mlserver.readthedocs.io/en/latest/reference/settings.html) seems to be preempted by `cpu_count` passed as command argument of Python `subprocess.Popen` that will start MLServer.
see https://github.com/mlflow/mlflow/blob/dcdc0db38f5cebaa6aa3b2474c669b81d1d9f197/mlflow/models/container/__init__.py#L189-L198

So on a server with 128 cpus / 512Go RAM we can serve only 2 models of 2Go of memory footprint without worrying about RAM exhaustion.

### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```bash
MLFLOW_TRACKING_URI=http://xxxxxxxxxxxxx/ \
mlflow models build-docker -m runs:/xxxxxxxxxxxxxxxxxxxxxxxxxxx/model \
--enable-mlserver -n mlserver_model_x:latest
docker run --rm -it -p 8080:8080 -p 8082:8082 -e MLSERVER_INFER_WORKERS=2 mlserver_model_x:latest
```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
On my 16cpus machine:
```
➜  ~ docker run --rm -it -p 8080:8080 -p 8082:8082 -e MLSERVER_INFER_WORKERS=2 mlserver_model_x:latest
2023-03-14 17:42:14,028 [mlserver] INFO - Using asyncio event-loop policy: uvloop
2023-03-14 17:42:15,579 [mlserver.parallel] DEBUG - Starting response processing loop...
2023-03-14 17:42:15,594 [mlserver.rest] INFO - HTTP server running on http://0.0.0.0:8080
INFO:     Started server process [39]
INFO:     Waiting for application startup.
2023-03-14 17:42:15,646 [mlserver.metrics] INFO - Metrics server running on http://0.0.0.0:8082
2023-03-14 17:42:15,646 [mlserver.metrics] INFO - Prometheus scraping endpoint can be accessed on http://0.0.0.0:8082/metrics
INFO:     Started server process [39]
INFO:     Waiting for application startup.
INFO:     Application startup complete.
2023-03-14 17:42:15,679 [mlserver.grpc] INFO - gRPC server running on http://0.0.0.0:8081
INFO:     Application startup complete.
INFO:     Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)
INFO:     Uvicorn running on http://0.0.0.0:8082 (Press CTRL+C to quit)
[2023-03-14 17:43:53,686] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:54,655] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:54,935] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:55,282] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:55,485] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:55,662] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:56,652] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:56,844] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:57,063] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:57,301] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:57,350] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,106] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,107] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,229] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:58,264] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:58,373] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:58,380] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:58,597] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:58,688] DEBUG /opt/ml/model/code/config.py:31 - LOG_LEVEL setted to DEBUG
[2023-03-14 17:43:58,927] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:59,128] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:59,136] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:59,290] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
[2023-03-14 17:43:59,379] DEBUG /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/tensorflow/python/distribute/cluster_resolver/tpu/tpu_cluster_resolver.py:32 - Falling back to TensorFlow client; we recommended you install the Cloud TPU client directly with pip install cloud-tpu-client.
2023-03-14 17:45:11,935 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:11,935] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:12,195 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:12,195] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:12,254 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:12,254] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:12,527 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:12,527] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:12,648 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:12,648] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:12,707 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:12,707] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:12,920 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:12,920] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,186 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,186] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,298 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,298] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,307 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,307] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,499 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,499] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,568 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,568] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,728 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,728] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,731 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,731] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:13,884 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:13,884] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:14,071 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
[2023-03-14 17:45:14,071] INFO /root/.mlflow/envs/mlflow-dd05953b9216fb8a26c0bd2830b6dcfb487bec8e/lib/python3.8/site-packages/mlserver/registry.py:160 - Loaded model 'mlflow-model' succesfully.
2023-03-14 17:45:14,152 [mlserver] INFO - Loaded model 'mlflow-model' succesfully.
INFO:     172.17.0.1:46884 - "GET /health HTTP/1.1" 200 OK
```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [X] `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
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [X] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [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

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:

Resolving patch:
diff --git a/mlflow/models/container/__init__.py b/mlflow/models/container/__init__.py
--- a/mlflow/models/container/__init__.py
+++ b/mlflow/models/container/__init__.py
@@ -187,12 +187,21 @@ def _serve_pyfunc(model, env_manager):
         procs.append(nginx)
 
     cpu_count = multiprocessing.cpu_count()
+    if enable_mlserver:
+        inference_server = mlserver
+        # Allows users to choose the number of workers using MLServer var env settings.
+        # Default to cpu count
+        nworkers = int(os.getenv("MLSERVER_INFER_WORKERS", cpu_count))
+        # Since MLServer will run without NGINX, expose the server in the `8080`
+        # port, which is the assumed "public" port.
+        port = DEFAULT_MLSERVER_PORT
+    else:
+        inference_server = scoring_server
+        # users can use GUNICORN_CMD_ARGS="--workers=3" var env to override the number of workers
+        nworkers = cpu_count
+        port = DEFAULT_INFERENCE_SERVER_PORT
 
-    inference_server = mlserver if enable_mlserver else scoring_server
-    # Since MLServer will run without NGINX, expose the server in the `8080`
-    # port, which is the assumed "public" port.
-    port = DEFAULT_MLSERVER_PORT if enable_mlserver else DEFAULT_INFERENCE_SERVER_PORT
-    cmd, cmd_env = inference_server.get_cmd(model_uri=MODEL_PATH, nworkers=cpu_count, port=port)
+    cmd, cmd_env = inference_server.get_cmd(model_uri=MODEL_PATH, nworkers=nworkers, port=port)
 
     bash_cmds.append(cmd)
     inference_server_process = Popen(["/bin/bash", "-c", " && ".join(bash_cmds)], env=cmd_env)
diff --git a/mlflow/utils/cli_args.py b/mlflow/utils/cli_args.py
--- a/mlflow/utils/cli_args.py
+++ b/mlflow/utils/cli_args.py
@@ -172,7 +172,11 @@ def _create_env_manager_option(help_string):
     "--enable-mlserver",
     is_flag=True,
     default=False,
-    help="Enable serving with MLServer through the v2 inference protocol.",
+    help=(
+        "Enable serving with MLServer through the v2 inference protocol. "
+        "You can use environment variables to configure MLServer. "
+        "(See https://mlserver.readthedocs.io/en/latest/reference/settings.html)"
+    ),
 )
 
 ARTIFACTS_DESTINATION = click.option(


Test patch: