Repository: mlflow/mlflow
Pull request: Optimize Docker image for model serving
<details><summary>&#x1F6E0 DevTools &#x1F6E0</summary>
<p>

[![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/B-Step62/mlflow/pull/10954?quickstart=1)

#### Install mlflow from this PR

```
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/10954/merge
```

#### Checkout with GitHub CLI

```
gh pr checkout 10954
```

</p>
</details>

### Related Issues/PRs

Resolve #2426, #5927

### What changes are proposed in this pull request?
#### Background
Docker image for model serving was not optimized and its large size has been pain for users (#2426, #5927). For example, few MB of scikit-learn model becomes > 3.7GB image after `build-image` command. Also the long building time (> 5 minutes) has been another pain.

After internal design discussion, we decided to introduce three optimizations:
1. Remove Java and dependencies if the model flavor doesn't require them.
2. Use Python base image instead of Ubuntu if possible.
3. Remove unnecessary conda/virtualenv layer if using Python image.
This PR implements those changes, as well as comprehensive integration tests with (almost) all flavors.

#### Implementation Notes
Some cases we cannot apply optimization, e.g. model_uri is not specified. To simplify the logic, I implemented those conditions as "all-or-nothing", namely, apply 1+2+3 all or do nothing. 
The optimizations apply when all the following conditions are met:
1. `model_uri` is specified.
2. User doesn't enable `--install-java` flag. (new CLI param)
3. The model flavor is not one of those require Java (e.g. `spark`, `mleap`)
4. Python version can be determined from the model metadata.
If any of this isn't met, we fallback to the original Ubuntu image with virtualenv and java. Tho majority of usage should meet these conditions so can benefit from the optimization.

#### Testing
For better safety, I've added integration tests covers (almost) all flavors in `tests/pyfunc/test_docker_flavors.py`. These tests are a bit time consuming (~20 mins) so I mark them to be skipped in the CI.

#### Impact
The impact depends on the actual size of necessary dependencies and model files, but significant for small models. For example, an image for small scikit-learn model becomes 0.99 GB (vs 3.7GB original; 73% reduction). Build time also decreases to 67 secs (vs 430 secs original; 84% reduction). The margin could be bigger for thin/small models like OpenAI, Langchain, while smaller for large models like Pytorch, Transformers.

#### For Reviewers
While the number of changed file is a bit large (28), most of them doesn't require deep review, such as test Dockerfiles, small tweak to fixture naming (to avoid conflict). Basically core logic changes reside following three iles:
* `mlflow/models/container/__init__.py`
* `mlflow/models/docker_utils.py`
* `mlflow/pyfunc/backend.py`

#### Can we reduce the image size further?
Speaking of the scikit-learn image, the majority of the size is used by MLflow itself and its dependencies (0.79 GB out of 0.99GB total), while many modules and dependencies are not required for model serving.


### How is this PR tested?

- [x] Existing unit/integration tests
- [x] New unit/integration tests
- [x] Manual tests

![Screenshot 2024-02-01 at 13 40 42](https://github.com/mlflow/mlflow/assets/31463517/be827d0d-0f42-49c5-b555-fb2df0573b28)

### Does this PR require documentation update?

- [x] No. You can skip the rest of this section.
- [ ] Yes. I've updated:
  - [ ] Examples
  - [ ] API references
  - [ ] Instructions

I will update deployment doc with the change accordingly (user-facing change shouldn't be too big).

### Release Notes

#### Is this a user-facing change?

- [ ] No. You can skip the rest of this section.
- [x] Yes. Give a description of this change to be included in the release notes for MLflow users.

Optimize Docker image for model serving. This reduces image size ~80% and build time ~85% depending on the model size and flavor. Note that this includes default behavior change for Java installation; the Docker image doesn't include Java by default, for all model flavors except Spark, Mleap, JohnSnowLabs, H2O. Please specify `--install-java` flag to enforce installation.

#### What component(s), interfaces, languages, and integrations does this PR affect?

Components

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `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

Interface

- [ ] `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

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
- [x] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations

<!--
Insert an empty named anchor here to allow jumping to this section with a fragment URL
(e.g. https://github.com/mlflow/mlflow/pull/123#user-content-release-note-category).
Note that GitHub prefixes anchor names in markdown with "user-content-".
-->

<a name="release-note-category"></a>

#### How should the PR be classified in the release notes? Choose one:

- [ ] `rn/none` - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
- [ ] `rn/breaking-change` - The PR will be mentioned in the "Breaking Changes" section
- [x] `rn/feature` - A new user-facing feature worth mentioning in the release notes
- [ ] `rn/bug-fix` - A user-facing bug fix worth mentioning in the release notes
- [ ] `rn/documentation` - A user-facing documentation change worth mentioning in the release notes


Problem statement:
[FR] Reduce Docker Image Size for Predict
### Motivation
Docker Image produced by MLFlow for predicting the models is supposed to be the most portable format but unfortunately, the docker image produced by MLFlow is of huge size. We tried with a simple model file with 1kb size (a simple regression model) but the resultant docker image is 3.3 GB.  We are trying to deploy these docker images (about 60) in Kubernetes but since K8s nodes will have limited resources, docker images of this size make mlflow a non-viable option. 
### Proposed Changes
Please let us know how the docker image size for predicting the model can be reduced.

Hints from issue discussion:
We should avoid packaging the jar files if it's a python served model, and remove condas for java (MLeap) serving. 
It'll also reduce the build time

I suggest we add the libraries (python or java) depending on the serving flavor of the model we download in the the `copy_model_into_container` function instead of hardcoding all the dependancies in `docker_utils.py`

For java models, maven shouldn't be packaged in the image. We should use a first image to run the build, and then copy the resulting jars in the final image.

If it's the best way to improve the image size & the build time I can implement it.
Hey Quentin,

Just wondering if there was any sense of prioritisation of this feature, as I think it would be a great enhancement?

I've taken a look at `docker_utils.py` and can have a go at an initial patch which strips out the java dependencies from the `_DOCKERFILE_TEMPLATE` & `_get_mlflow_install_step`, given an optional `python_only` arg to `_build_image`, because I think it may be marginally preferable image size/caching wise than installing the dependencies slightly later on as part of the `custom_setup_steps_hook`
Yes I've started a PR and removing the java part when only python is needed, but also removing the conda env when the serving layer is in java.
Will update the ticket soon with the PR content
That's awesome, thanks!

As a tangential idea, it would also be great to add a `cached_image` arg to the `_build_image` func

i.e. `["docker", "build", "-t", image_name, "-f", "Dockerfile", ".", "--cache-from", cached_image]`

(Happy to raise a PR if it's outside the scope of the changes you're planning to make)
I have worked on optimising the docker image size. I am able to reduce the size from 3.45GB to 1.41GB. I have done the following changes. I am planning to generalize the changes and raise a PR.
1. As we were using only python model serve, removed all java dependencies like maven downloads and maven builds.
2. Used the base miniconda environment instead of creating a new virtual env "custom_env". 
3. Make --no-conda as True by default. 
4. Also from conda.yaml install only the required dependencies to serve the model (instead of creating a new environment with all the dependencies)
I have created a pull request #2496 and all default test cases succeeded. I have added a new flag "--python-only" and "--no-conda" to build a docker image without java jar files and uses the base conda environment (doesn't create new virtual environment custom_env). 

The model that I have tested created a docker image of size 2.11GB and after optimization docker image size is 886MB.

Please review the code changes.

With my new changes in pull request #2496, docker image is created either for python flavor or java flavor based on the model for which docker image is being created for.

With the new argument added to build-docker  ``--no-conda`` that doesn't create a new conda environment "custom_env" instead uses the default miniconda base environment.

The current docker image size for sklearn iris example is 3.21GB and with the new changes docker image size is reduced to 2.14GB and by using --no-conda option (ie., using base miniconda environment), docker image size can be further reduced to 887MB.
Apologies for the delay, thanks all for chiming in! @ambiSanthu the approach of adding dependencies based on the model's flavor makes sense. I wonder if instead of adding a `--no-conda` option, it would make sense to simply remove the base miniconda image to save space (not sure how many models can be served solely using the base miniconda image). Thoughts?
As per the docker image creation flow, a new conda env called "custom_env" is being created and dependencies from conda.yml is installed in this custom_env. With --no-conda option, I am not creating the custom_env but parsing the conda.yml and installing the dependencies directly into base env. Hence I make sure required dependencies are installed for the model to run.
Hi, 

I'm very interested in this issue! We want to deploy tens/hundreds of models using mlflow build-docker, and the image size is clearly a blocker!

In general, I have 2 comments:
1. The docker image already provides a decent form of virtualization. Creating another conda environment inside it can be redundant.  
2. I would imagine an even more customized feature (please point me to an issue if created before), which allows a user to reference his pre-written Docker template file to create the model image. Currently there are some dependencies hard-coded [here](https://github.com/mlflow/mlflow/blob/master/mlflow/models/docker_utils.py#L14). I see the PR #2496  split it into 2 for Java and Python separately, but that still looks a little rigid to me

Commit message: more and more fix

Signed-off-by: B-Step62 <yuki.watanabe@databricks.com>

Patch to review:
diff --git a/.github/workflows/master.yml b/.github/workflows/master.yml
index b824a2c8f5415..2af15fb42b0ad 100644
--- a/.github/workflows/master.yml
+++ b/.github/workflows/master.yml
@@ -276,7 +276,8 @@ jobs:
       - name: Run tests
         run: |
           pytest --splits=${{ matrix.splits }} --group=${{ matrix.group }} --durations=30 \
-            tests/pyfunc --ignore tests/pyfunc/test_spark_connect.py
+            tests/pyfunc --ignore tests/pyfunc/test_spark_connect.py \
+            --ignore tests/pyfunc/docker/test_docker_flavors.py
 
           # test_spark_connect.py fails if it's run with ohter tests, so run it separately.
           pytest tests/pyfunc/test_spark_connect.py
diff --git a/mlflow/models/cli.py b/mlflow/models/cli.py
index 60c0bf1cc614c..98072483566f5 100644
--- a/mlflow/models/cli.py
+++ b/mlflow/models/cli.py
@@ -170,10 +170,17 @@ def prepare_env(
 )
 @cli_args.ENV_MANAGER
 @cli_args.MLFLOW_HOME
+@cli_args.INSTALL_JAVA
 @cli_args.INSTALL_MLFLOW
 @cli_args.ENABLE_MLSERVER
 def generate_dockerfile(
-    model_uri, output_directory, env_manager, mlflow_home, install_mlflow, enable_mlserver
+    model_uri,
+    output_directory,
+    env_manager,
+    mlflow_home,
+    install_java,
+    install_mlflow,
+    enable_mlserver,
 ):
     """
     Generates a directory with Dockerfile whose default entrypoint serves an MLflow model at port
@@ -191,6 +198,7 @@ def generate_dockerfile(
             model_uri,
             output_directory,
             mlflow_home=mlflow_home,
+            install_java=install_java,
             install_mlflow=install_mlflow,
             enable_mlserver=enable_mlserver,
         )
@@ -208,6 +216,7 @@ def generate_dockerfile(
 @click.option("--name", "-n", default="mlflow-pyfunc-servable", help="Name to use for built image")
 @cli_args.ENV_MANAGER
 @cli_args.MLFLOW_HOME
+@cli_args.INSTALL_JAVA
 @cli_args.INSTALL_MLFLOW
 @cli_args.ENABLE_MLSERVER
 def build_docker(**kwargs):
diff --git a/mlflow/models/container/__init__.py b/mlflow/models/container/__init__.py
index 988190ab51a54..b77c876de17a9 100644
--- a/mlflow/models/container/__init__.py
+++ b/mlflow/models/container/__init__.py
@@ -12,6 +12,7 @@
 import sys
 from pathlib import Path
 from subprocess import Popen, check_call
+from typing import List
 
 import mlflow
 import mlflow.version
@@ -22,6 +23,7 @@
 from mlflow.pyfunc import _extract_conda_env, mlserver, scoring_server
 from mlflow.store.artifact.models_artifact_repo import REGISTERED_MODEL_META_FILE_NAME
 from mlflow.utils import env_manager as em
+from mlflow.utils.environment import _PythonEnv
 from mlflow.utils.file_utils import read_yaml
 from mlflow.utils.virtualenv import _get_or_create_virtualenv
 from mlflow.version import VERSION as MLFLOW_VERSION
@@ -86,34 +88,7 @@ def _install_pyfunc_deps(
     Creates a conda env for serving the model at the specified path and installs almost all serving
     dependencies into the environment - MLflow is not installed as it's not available via conda.
     """
-    # If model is a pyfunc model, create its conda env (even if it also has mleap flavor)
-    activate_cmd = []
-    if model_path:
-        model_config_path = os.path.join(model_path, MLMODEL_FILE_NAME)
-        model = Model.load(model_config_path)
-        # NOTE: this differs from _serve cause we always activate the env even if you're serving
-        # an mleap model
-        if pyfunc.FLAVOR_NAME not in model.flavors:
-            return
-        conf = model.flavors[pyfunc.FLAVOR_NAME]
-        if pyfunc.ENV in conf:
-            _logger.info("creating and activating custom environment")
-            env = _extract_conda_env(conf[pyfunc.ENV])
-            env_path_dst = os.path.join("/opt/mlflow/", env)
-            env_path_dst_dir = os.path.dirname(env_path_dst)
-            if not os.path.exists(env_path_dst_dir):
-                os.makedirs(env_path_dst_dir)
-            shutil.copy2(os.path.join(MODEL_PATH, env), env_path_dst)
-            if env_manager == em.CONDA:
-                conda_create_model_env = f"conda env create -n custom_env -f {env_path_dst}"
-                if Popen(["bash", "-c", conda_create_model_env]).wait() != 0:
-                    raise Exception("Failed to create model environment.")
-                activate_cmd = ["source /miniconda/bin/activate custom_env"]
-            elif env_manager == em.VIRTUALENV:
-                env_activate_cmd = _get_or_create_virtualenv(model_path)
-                path = env_activate_cmd.split(" ")[-1]
-                os.symlink(path, "/opt/activate")
-                activate_cmd = [env_activate_cmd]
+    activate_cmd = _install_model_dependencies_to_env(model_path, env_manager) if model_path else []
 
     # NB: install gunicorn[gevent] from pip rather than from conda because gunicorn is already
     # dependency of mlflow on pip and we expect mlflow to be part of the environment.
@@ -125,6 +100,8 @@ def _install_pyfunc_deps(
     if Popen(["bash", "-c", " && ".join(activate_cmd + install_server_deps)]).wait() != 0:
         raise Exception("Failed to install serving dependencies into the model environment.")
 
+    # NB: If we don't use virtualenv or conda env, we don't need to install mlflow here as
+    # it's already installed in the container.
     if len(activate_cmd) and install_mlflow:
         install_mlflow_cmd = [
             "pip install /opt/mlflow/."
@@ -136,6 +113,57 @@ def _install_pyfunc_deps(
     return activate_cmd
 
 
+def _install_model_dependencies_to_env(model_path, env_manager) -> List[str]:
+    """:
+    Installs model dependencies to the specified environment, which can be either a local
+    environment, a conda environment, or a virtualenv.
+
+    Returns:
+        Empty list if local environment, otherwise a list of bash commands to activate the
+        virtualenv or conda environment.
+    """
+    model_config_path = os.path.join(model_path, MLMODEL_FILE_NAME)
+    model = Model.load(model_config_path)
+
+    conf = model.flavors.get(pyfunc.FLAVOR_NAME, {})
+    if pyfunc.ENV not in conf:
+        return []
+    env_conf = conf[mlflow.pyfunc.ENV]
+
+    if env_manager == em.LOCAL:
+        # Install pip dependencies directly into the local environment
+        python_env_config_path = os.path.join(model_path, env_conf[em.VIRTUALENV])
+        python_env = _PythonEnv.from_yaml(python_env_config_path)
+        deps = " ".join(python_env.build_dependencies + python_env.dependencies)
+        deps = deps.replace("requirements.txt", os.path.join(model_path, "requirements.txt"))
+        if Popen(["bash", "-c", f"python -m pip install {deps}"]).wait() != 0:
+            raise Exception("Failed to install model dependencies.")
+        return []
+
+    _logger.info("creating and activating custom environment")
+
+    env = _extract_conda_env(env_conf)
+    env_path_dst = os.path.join("/opt/mlflow/", env)
+    env_path_dst_dir = os.path.dirname(env_path_dst)
+    if not os.path.exists(env_path_dst_dir):
+        os.makedirs(env_path_dst_dir)
+    shutil.copy2(os.path.join(MODEL_PATH, env), env_path_dst)
+
+    if env_manager == em.CONDA:
+        conda_create_model_env = f"conda env create -n custom_env -f {env_path_dst}"
+        if Popen(["bash", "-c", conda_create_model_env]).wait() != 0:
+            raise Exception("Failed to create model environment.")
+        activate_cmd = ["source /miniconda/bin/activate custom_env"]
+
+    elif env_manager == em.VIRTUALENV:
+        env_activate_cmd = _get_or_create_virtualenv(model_path)
+        path = env_activate_cmd.split(" ")[-1]
+        os.symlink(path, "/opt/activate")
+        activate_cmd = [env_activate_cmd]
+
+    return activate_cmd
+
+
 def _serve_pyfunc(model, env_manager):
     # option to disable manually nginx. The default behavior is to enable nginx.
     disable_nginx = os.getenv(DISABLE_NGINX, "false").lower() == "true"
@@ -145,6 +173,10 @@ def _serve_pyfunc(model, env_manager):
     conf = model.flavors[pyfunc.FLAVOR_NAME]
     bash_cmds = []
     if pyfunc.ENV in conf:
+        # NB: MLFLOW_DISABLE_ENV_CREATION is False only for SageMaker deployment, where the model
+        # files are loaded into the container at runtime rather than build time. In this case,
+        # we need to create a virtual environment and install the model dependencies into it when
+        # starting the container.
         if not disable_env_creation:
             _install_pyfunc_deps(
                 MODEL_PATH,
diff --git a/mlflow/models/docker_utils.py b/mlflow/models/docker_utils.py
index 183c9265e1898..484cf2e64a955 100644
--- a/mlflow/models/docker_utils.py
+++ b/mlflow/models/docker_utils.py
@@ -1,6 +1,6 @@
-import logging
 import os
 from subprocess import PIPE, STDOUT, Popen
+from typing import Optional, Union
 from urllib.parse import urlparse
 
 from mlflow.utils import env_manager as em
@@ -8,20 +8,12 @@
 from mlflow.utils.logging_utils import eprint
 from mlflow.version import VERSION
 
-_logger = logging.getLogger(__name__)
+UBUNTU_BASE_IMAGE = "ubuntu:20.04"
+PYTHON_SLIM_BASE_IMAGE = "python:{version}-slim"
 
-SETUP_MINICONDA = """
-# Setup miniconda
-RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh >> miniconda.sh
-RUN bash ./miniconda.sh -b -p /miniconda && rm ./miniconda.sh
-ENV PATH="/miniconda/bin:$PATH"
-"""
 
-SETUP_PYENV_AND_VIRTUALENV = r"""
-# Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
+SETUP_PYENV_AND_VIRTUALENV = r"""# Setup pyenv
+RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata \
     libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
     libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
 RUN git clone \
@@ -30,106 +22,112 @@
     https://github.com/pyenv/pyenv.git /root/.pyenv
 ENV PYENV_ROOT="/root/.pyenv"
 ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
+RUN apt install -y python3.8 python3.8-distutils \
+    && ln -s -f $(which python3.8) /usr/bin/python \
+    && wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py \
+    && python /tmp/get-pip.py
 RUN pip install virtualenv
 """
 
+_DOCKERFILE_TEMPLATE = """# Build an image that can serve mlflow models.
+FROM {base_image}
 
-def _get_maven_proxy():
-    http_proxy = os.getenv("http_proxy")
-    https_proxy = os.getenv("https_proxy")
-    if not http_proxy or not https_proxy:
-        return ""
-
-    # Expects proxies as either PROTOCOL://{USER}:{PASSWORD}@HOSTNAME:PORT
-    # or PROTOCOL://HOSTNAME:PORT
-    parsed_http_proxy = urlparse(http_proxy)
-    assert parsed_http_proxy.hostname is not None, "Invalid `http_proxy` hostname."
-    assert parsed_http_proxy.port is not None, f"Invalid proxy port: {parsed_http_proxy.port}"
-
-    parsed_https_proxy = urlparse(https_proxy)
-    assert parsed_https_proxy.hostname is not None, "Invalid `https_proxy` hostname."
-    assert parsed_https_proxy.port is not None, f"Invalid proxy port: {parsed_https_proxy.port}"
-
-    maven_proxy_options = (
-        "-DproxySet=true",
-        f"-Dhttp.proxyHost={parsed_http_proxy.hostname}",
-        f"-Dhttp.proxyPort={parsed_http_proxy.port}",
-        f"-Dhttps.proxyHost={parsed_https_proxy.hostname}",
-        f"-Dhttps.proxyPort={parsed_https_proxy.port}",
-        "-Dhttps.nonProxyHosts=repo.maven.apache.org",
-    )
+{setup_python_venv}
 
-    if parsed_http_proxy.username is None or parsed_http_proxy.password is None:
-        return " ".join(maven_proxy_options)
+{setup_java}
 
-    return " ".join(
-        (
-            *maven_proxy_options,
-            f"-Dhttp.proxyUser={parsed_http_proxy.username}",
-            f"-Dhttp.proxyPassword={parsed_http_proxy.password}",
-        )
-    )
-
-
-_DOCKERFILE_TEMPLATE = """
-# Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends \
-         wget \
-         curl \
-         nginx \
-         ca-certificates \
-         bzip2 \
-         build-essential \
-         cmake \
-         openjdk-8-jdk \
-         git-core \
-         maven \
-    && rm -rf /var/lib/apt/lists/*
-
-{setup_python_env}
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
 WORKDIR /opt/mlflow
 
 {install_mlflow}
 
-{custom_setup_steps}
+{install_model_and_deps}
 
 ENV MLFLOW_DISABLE_ENV_CREATION={disable_env_creation}
 ENV ENABLE_MLSERVER={enable_mlserver}
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-{entrypoint}
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "{entrypoint}"]
 """
 
 
-def _get_mlflow_install_step(dockerfile_context_dir, mlflow_home):
+SETUP_MINICONDA = """# Setup miniconda
+RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh >> miniconda.sh
+RUN bash ./miniconda.sh -b -p /miniconda && rm ./miniconda.sh
+ENV PATH="/miniconda/bin:$PATH"
+"""
+
+
+def generate_dockerfile(
+    output_dir: str,
+    base_image: str,
+    model_install_steps: Optional[str],
+    entrypoint: str,
+    env_manager: Union[em.CONDA, em.LOCAL, em.VIRTUALENV],
+    mlflow_home: Optional[str] = None,
+    enable_mlserver: bool = False,
+    disable_env_creation_at_runtime: bool = True,
+):
     """
-    Get docker build commands for installing MLflow given a Docker context dir and optional source
-    directory
+    Generates a Dockerfile that can be used to build a docker image, that serves ML model
+    stored and tracked in MLflow.
     """
-    mlflow_home = os.path.abspath(mlflow_home) if mlflow_home else None
+    if base_image.startswith("python:"):
+        setup_python_venv_steps = (
+            "RUN apt-get -y update && apt-get install -y --no-install-recommends nginx"
+        )
+        setup_java_steps = ""
+        install_mlflow_steps = _pip_mlflow_install_step(output_dir, mlflow_home)
+
+    elif base_image == "ubuntu:20.04":
+        setup_python_venv_steps = (
+            "RUN apt-get -y update && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y "
+            "--no-install-recommends wget curl nginx ca-certificates bzip2 build-essential cmake "
+            "git-core\n\n"
+        )
+        setup_python_venv_steps += (
+            SETUP_MINICONDA if env_manager == em.CONDA else SETUP_PYENV_AND_VIRTUALENV
+        )
+
+        setup_java_steps = (
+            "# Setup Java\n"
+            "RUN apt-get install -y --no-install-recommends openjdk-8-jdk maven\n"
+            "ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64"
+        )
 
+        install_mlflow_steps = _pip_mlflow_install_step(output_dir, mlflow_home)
+        install_mlflow_steps += "\n\n" + _java_mlflow_install_step(mlflow_home)
+
+    else:
+        raise ValueError(f"Unsupported base image: {base_image}")
+
+    with open(os.path.join(output_dir, "Dockerfile"), "w") as f:
+        f.write(
+            _DOCKERFILE_TEMPLATE.format(
+                base_image=base_image,
+                setup_python_venv=setup_python_venv_steps,
+                setup_java=setup_java_steps,
+                install_mlflow=install_mlflow_steps,
+                install_model_and_deps=model_install_steps,
+                entrypoint=entrypoint,
+                enable_mlserver=enable_mlserver,
+                disable_env_creation=disable_env_creation_at_runtime,
+            )
+        )
+
+
+def _java_mlflow_install_step(mlflow_home):
     maven_proxy = _get_maven_proxy()
     if mlflow_home:
-        mlflow_dir = _copy_project(src_path=mlflow_home, dst_path=dockerfile_context_dir)
         return (
-            f"COPY {mlflow_dir} /opt/mlflow\n"
-            "RUN pip install /opt/mlflow\n"
+            "# Install Java mlflow-scoring from local source\n"
             "RUN cd /opt/mlflow/mlflow/java/scoring && "
             f"mvn --batch-mode package -DskipTests {maven_proxy} && "
             "mkdir -p /opt/java/jars && "
@@ -138,7 +136,7 @@ def _get_mlflow_install_step(dockerfile_context_dir, mlflow_home):
         )
     else:
         return (
-            f"RUN pip install mlflow=={VERSION}\n"
+            "# Install Java mlflow-scoring from Maven Central\n"
             "RUN mvn"
             " --batch-mode dependency:copy"
             f" -Dartifact=org.mlflow:mlflow-scoring:{VERSION}:pom"
@@ -154,34 +152,59 @@ def _get_mlflow_install_step(dockerfile_context_dir, mlflow_home):
         )
 
 
-def generate_dockerfile(
-    output_dir,
-    entrypoint,
-    env_manager,
-    mlflow_home=None,
-    custom_setup_steps=None,
-    enable_mlserver=False,
-    disable_env_creation=False,
-):
+def _get_maven_proxy():
+    http_proxy = os.getenv("http_proxy")
+    https_proxy = os.getenv("https_proxy")
+    if not http_proxy or not https_proxy:
+        return ""
+
+    # Expects proxies as either PROTOCOL://{USER}:{PASSWORD}@HOSTNAME:PORT
+    # or PROTOCOL://HOSTNAME:PORT
+    parsed_http_proxy = urlparse(http_proxy)
+    assert parsed_http_proxy.hostname is not None, "Invalid `http_proxy` hostname."
+    assert parsed_http_proxy.port is not None, f"Invalid proxy port: {parsed_http_proxy.port}"
+
+    parsed_https_proxy = urlparse(https_proxy)
+    assert parsed_https_proxy.hostname is not None, "Invalid `https_proxy` hostname."
+    assert parsed_https_proxy.port is not None, f"Invalid proxy port: {parsed_https_proxy.port}"
+
+    maven_proxy_options = (
+        "-DproxySet=true",
+        f"-Dhttp.proxyHost={parsed_http_proxy.hostname}",
+        f"-Dhttp.proxyPort={parsed_http_proxy.port}",
+        f"-Dhttps.proxyHost={parsed_https_proxy.hostname}",
+        f"-Dhttps.proxyPort={parsed_https_proxy.port}",
+        "-Dhttps.nonProxyHosts=repo.maven.apache.org",
+    )
+
+    if parsed_http_proxy.username is None or parsed_http_proxy.password is None:
+        return " ".join(maven_proxy_options)
+
+    return " ".join(
+        (
+            *maven_proxy_options,
+            f"-Dhttp.proxyUser={parsed_http_proxy.username}",
+            f"-Dhttp.proxyPassword={parsed_http_proxy.password}",
+        )
+    )
+
+
+def _pip_mlflow_install_step(dockerfile_context_dir, mlflow_home):
     """
-    Generates a Dockerfile that can be used to build a docker image, that serves ML model
-    stored and tracked in MLflow.
+    Get docker build commands for installing MLflow given a Docker context dir and optional source
+    directory
     """
-    install_mlflow_steps = _get_mlflow_install_step(output_dir, mlflow_home)
-    setup_python_env_steps = (
-        SETUP_MINICONDA if env_manager == em.CONDA else SETUP_PYENV_AND_VIRTUALENV
-    )
-    with open(os.path.join(output_dir, "Dockerfile"), "w") as f:
-        f.write(
-            _DOCKERFILE_TEMPLATE.format(
-                setup_python_env=setup_python_env_steps,
-                install_mlflow=install_mlflow_steps,
-                custom_setup_steps=custom_setup_steps,
-                entrypoint=entrypoint,
-                enable_mlserver=enable_mlserver,
-                disable_env_creation=disable_env_creation,
-            )
+    if mlflow_home:
+        mlflow_dir = _copy_project(
+            src_path=os.path.abspath(mlflow_home), dst_path=dockerfile_context_dir
         )
+        return (
+            "# Install MLflow from local source\n"
+            f"COPY {mlflow_dir} /opt/mlflow\n"
+            "RUN pip install /opt/mlflow"
+        )
+    else:
+        return f"# Install MLflow\nRUN pip install mlflow=={VERSION}"
 
 
 def build_image_from_context(context_dir: str, image_name: str):
diff --git a/mlflow/models/python_api.py b/mlflow/models/python_api.py
index 71cc7956bfc3a..7c1cb4e8a5379 100644
--- a/mlflow/models/python_api.py
+++ b/mlflow/models/python_api.py
@@ -19,6 +19,7 @@ def build_docker(
     name="mlflow-pyfunc",
     env_manager=_EnvManager.VIRTUALENV,
     mlflow_home=None,
+    install_java=False,
     install_mlflow=False,
     enable_mlserver=False,
 ):
@@ -48,6 +49,7 @@ def build_docker(
         model_uri,
         name,
         mlflow_home=mlflow_home,
+        install_java=install_java,
         install_mlflow=install_mlflow,
         enable_mlserver=enable_mlserver,
     )
diff --git a/mlflow/pyfunc/backend.py b/mlflow/pyfunc/backend.py
index 672e158b3c5b9..a2cc6cb6cbb84 100644
--- a/mlflow/pyfunc/backend.py
+++ b/mlflow/pyfunc/backend.py
@@ -10,8 +10,11 @@
 import warnings
 from pathlib import Path
 
+from mlflow import pyfunc
 from mlflow.exceptions import MlflowException
 from mlflow.models import FlavorBackend, docker_utils
+from mlflow.models.docker_utils import PYTHON_SLIM_BASE_IMAGE, UBUNTU_BASE_IMAGE
+from mlflow.models.model import MLMODEL_FILE_NAME, Model
 from mlflow.pyfunc import (
     ENV,
     _extract_conda_env,
@@ -22,13 +25,14 @@
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
 from mlflow.utils import env_manager as _EnvManager
 from mlflow.utils.conda import get_conda_bin_executable, get_or_create_conda_env
-from mlflow.utils.environment import Environment
+from mlflow.utils.environment import Environment, _PythonEnv
 from mlflow.utils.file_utils import (
     TempDir,
     get_or_create_nfs_tmp_dir,
     get_or_create_tmp_dir,
     path_to_local_file_uri,
 )
+from mlflow.utils.model_utils import _get_all_flavor_configurations
 from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir
 from mlflow.utils.process import ShellCommandException, cache_return_value_per_process
 from mlflow.utils.virtualenv import (
@@ -42,6 +46,18 @@
 _IS_UNIX = os.name != "nt"
 _STDIN_SERVER_SCRIPT = Path(__file__).parent.joinpath("stdin_server.py")
 
+# Flavors that require Java to be installed in the environment
+JAVA_FLAVORS = {"johnsnowlabs", "h2o", "mleap", "spark"}
+
+# Some flavor requires additional packages to be installed in the environment
+FLAVOR_SPECIFIC_APT_PACKAGES = {
+    "lightgbm": ["libgomp1"],
+    "paddle": ["libgomp1"],
+}
+
+# Directory to store loaded model inside the Docker context directory
+_MODEL_DIR_NAME = "model_dir"
+
 
 class PyFuncBackend(FlavorBackend):
     """
@@ -318,69 +334,147 @@ def can_score_model(self):
             # Can not find conda
             return False
 
+    def build_image(
+        self,
+        model_uri,
+        image_name,
+        install_java=False,
+        install_mlflow=False,
+        mlflow_home=None,
+        enable_mlserver=False,
+    ):
+        with TempDir() as tmp:
+            cwd = tmp.path()
+            self.generate_dockerfile(
+                model_uri=model_uri,
+                output_dir=cwd,
+                install_java=install_java,
+                install_mlflow=install_mlflow,
+                mlflow_home=mlflow_home,
+                enable_mlserver=enable_mlserver,
+            )
+
+            _logger.info("Building docker image with name %s", image_name)
+            docker_utils.build_image_from_context(context_dir=cwd, image_name=image_name)
+
     def generate_dockerfile(
-        self, model_uri, output_dir, install_mlflow=False, mlflow_home=None, enable_mlserver=False
+        self,
+        model_uri,
+        output_dir,
+        install_java=False,
+        install_mlflow=False,
+        mlflow_home=None,
+        enable_mlserver=False,
     ):
         os.makedirs(output_dir, exist_ok=True)
         _logger.debug("Created all folders in path", extra={"output_directory": output_dir})
 
-        # Copy model to image if model_uri is specified
-        custom_setup_steps = (
-            self._get_copy_model_steps(output_dir, model_uri, install_mlflow, enable_mlserver)
-            if model_uri
-            else ""
-        )
+        if model_uri:
+            model_cwd = os.path.join(output_dir, _MODEL_DIR_NAME)
+            pathlib.Path(model_cwd).mkdir(parents=True, exist_ok=True)
+            model_path = _download_artifact_from_uri(model_uri, output_path=model_cwd)
+            base_image = self._get_base_image(model_path, install_java)
+
+            # We don't need virtualenv or conda if base image is python
+            env_manager = (
+                _EnvManager.LOCAL if base_image.startswith("python") else self._env_manager
+            )
 
-        pyfunc_entrypoint = self._pyfunc_entrypoint(model_uri, install_mlflow, enable_mlserver)
+            model_install_steps = self._model_installation_steps(
+                model_path, env_manager, install_mlflow, enable_mlserver
+            )
+            entrypoint = f"from mlflow.models import container as C; C._serve('{env_manager}')"
+
+        else:
+            base_image = UBUNTU_BASE_IMAGE
+            model_install_steps = ""
+            # If model_uri is not specified, dependencies are installed at runtime
+            entrypoint = (
+                self._get_install_pyfunc_deps_cmd(
+                    self._env_manager, install_mlflow, enable_mlserver
+                )
+                + f" C._serve('{self._env_manager}')"
+            )
 
         dockerfile_text = docker_utils.generate_dockerfile(
             output_dir=output_dir,
-            custom_setup_steps=custom_setup_steps,
-            entrypoint=pyfunc_entrypoint,
+            base_image=base_image,
+            model_install_steps=model_install_steps,
+            entrypoint=entrypoint,
             env_manager=self._env_manager,
             mlflow_home=mlflow_home,
             enable_mlserver=enable_mlserver,
-            disable_env_creation=True,  # Always disable env creation for pyfunc
+            # always disable env creation at runtime for pyfunc
+            disable_env_creation_at_runtime=True,
         )
         _logger.debug("generated dockerfile at {output_dir}", extra={"dockerfile": dockerfile_text})
 
-    def build_image(
-        self, model_uri, image_name, install_mlflow=False, mlflow_home=None, enable_mlserver=False
-    ):
-        with TempDir() as tmp:
-            cwd = tmp.path()
-            self.generate_dockerfile(model_uri, cwd, install_mlflow, mlflow_home, enable_mlserver)
+    def _get_base_image(self, model_path: str, install_java: bool) -> str:
+        """
+        Determine the base image to use for the Dockerfile.
 
-            _logger.info("Building docker image with name %s", image_name)
-            docker_utils.build_image_from_context(context_dir=cwd, image_name=image_name)
+        We use Python slim base image when all of the following conditions are met:
+          1. Model URI is specified by the user
+          2. Model flavor does not require Java
+          3. Python version is specified in the model
 
-    def _get_copy_model_steps(self, output_dir, model_uri, install_mlflow, enable_mlserver):
-        model_cwd = os.path.join(output_dir, "model_dir")
-        pathlib.Path(model_cwd).mkdir(parents=True, exist_ok=True)
+        Returns:
+            Either the Ubuntu base image or the Python slim base image.
+        """
+        # Check if the model requires Java
+        if not install_java:
+            flavors = _get_all_flavor_configurations(model_path).keys()
+            if java_flavors := JAVA_FLAVORS & flavors:
+                _logger.info(f"Detected java flavors {java_flavors}, installing Java in the image")
+                install_java = True
+
+        # Use ubuntu base image if Java is required
+        if install_java:
+            return UBUNTU_BASE_IMAGE
+
+        # Get Python version from MLmodel
+        try:
+            model_config_path = os.path.join(model_path, MLMODEL_FILE_NAME)
+            model = Model.load(model_config_path)
+
+            conf = model.flavors[pyfunc.FLAVOR_NAME]
+            env_conf = conf[pyfunc.ENV]
+            python_env_config_path = os.path.join(model_path, env_conf[_EnvManager.VIRTUALENV])
+
+            python_env = _PythonEnv.from_yaml(python_env_config_path)
+            return PYTHON_SLIM_BASE_IMAGE.format(version=python_env.python)
+        except Exception as e:
+            _logger.warning(
+                f"Failed to determine Python version from {model_config_path}. "
+                f"Defaulting to {UBUNTU_BASE_IMAGE}. Error: {e}"
+            )
+            return UBUNTU_BASE_IMAGE
 
-        # If model_uri is specified, copy the model to the image and install its dependencies
-        model_path = _download_artifact_from_uri(model_uri, output_path=model_cwd)
-        model_dir = str(posixpath.join("model_dir", os.path.basename(model_path)))
+    def _model_installation_steps(self, model_path, env_manager, install_mlflow, enable_mlserver):
+        model_dir = str(posixpath.join(_MODEL_DIR_NAME, os.path.basename(model_path)))
+        # Copy model to image if model_uri is specified
+        steps = (
+            "# Copy model to image and install dependencies\n"
+            f"COPY {model_dir} /opt/ml/model\nRUN python -c "
+        )
+        steps += (
+            f'"{self._get_install_pyfunc_deps_cmd(env_manager, install_mlflow, enable_mlserver)}"'
+        )
 
-        install_deps_cmd = self._get_install_pyfunc_deps_cmd(install_mlflow, enable_mlserver)
-        return f'COPY {model_dir} /opt/ml/model\nRUN python -c "{install_deps_cmd}"'
+        # Install flavor-specific dependencies if needed
+        flavors = _get_all_flavor_configurations(model_path).keys()
+        for flavor in flavors:
+            if flavor in FLAVOR_SPECIFIC_APT_PACKAGES:
+                packages = " ".join(FLAVOR_SPECIFIC_APT_PACKAGES[flavor])
+                steps += f"\nRUN apt-get install -y --no-install-recommends {packages}"
 
-    def _pyfunc_entrypoint(self, model_uri, install_mlflow, enable_mlserver):
-        if model_uri:
-            # If model_uri is specified, dependencies are installed at build time so we don't
-            # need to run the install command at runtime
-            install_deps_cmd = ""
-        else:
-            install_deps_cmd = self._get_install_pyfunc_deps_cmd(install_mlflow, enable_mlserver)
-        entrypoint = (
-            f"from mlflow.models import container as C;{install_deps_cmd} "
-            f"C._serve('{self._env_manager}')"
-        )
-        return f'ENTRYPOINT ["python", "-c", "{entrypoint}"]'
+        return steps
 
-    def _get_install_pyfunc_deps_cmd(self, install_mlflow, enable_mlserver):
+    def _get_install_pyfunc_deps_cmd(
+        self, env_manager: _EnvManager, install_mlflow: bool, enable_mlserver: bool
+    ):
         return (
-            "from mlflow.models.container import _install_pyfunc_deps; "
-            f"_install_pyfunc_deps('/opt/ml/model', install_mlflow={install_mlflow}, "
-            f"enable_mlserver={enable_mlserver}, env_manager='{self._env_manager}');"
+            "from mlflow.models import container as C; "
+            f"C._install_pyfunc_deps('/opt/ml/model', install_mlflow={install_mlflow}, "
+            f"enable_mlserver={enable_mlserver}, env_manager='{env_manager}');"
         )
diff --git a/mlflow/sagemaker/cli.py b/mlflow/sagemaker/cli.py
index 0d068ab2598ac..5fe573d862090 100644
--- a/mlflow/sagemaker/cli.py
+++ b/mlflow/sagemaker/cli.py
@@ -358,22 +358,26 @@ def build_and_push_container(build, push, container, env_manager, mlflow_home):
         click.echo("skipping both build and push, have nothing to do!")
     if build:
         sagemaker_image_entrypoint = (
-            'ENTRYPOINT ["python", "-c", "import sys; from mlflow.models import container as C; '
-            f"C._init(sys.argv[1], '{env_manager}')\"]"
+            "import sys; from mlflow.models import container as C; "
+            f"C._init(sys.argv[1], '{env_manager}')"
         )
 
         setup_container = (
+            "# Install minimal serving dependencies\n"
             'RUN python -c "from mlflow.models.container import _install_pyfunc_deps;'
             '_install_pyfunc_deps(None, False)"'
         )
 
         with tempfile.TemporaryDirectory() as tmp:
             docker_utils.generate_dockerfile(
+                base_image=mlflow.models.docker_utils.UBUNTU_BASE_IMAGE,
                 output_dir=tmp,
                 entrypoint=sagemaker_image_entrypoint,
                 env_manager=env_manager,
                 mlflow_home=os.path.abspath(mlflow_home) if mlflow_home else None,
-                custom_setup_steps=setup_container,
+                model_install_steps=setup_container,
+                # Create a conda env or virtualenv at runtime after the model is loaded
+                disable_env_creation_at_runtime=False,
             )
 
             docker_utils.build_image_from_context(tmp, image_name=container)
diff --git a/mlflow/utils/cli_args.py b/mlflow/utils/cli_args.py
index 247e3824db3dc..e697d17742ea4 100644
--- a/mlflow/utils/cli_args.py
+++ b/mlflow/utils/cli_args.py
@@ -212,3 +212,9 @@ def _create_env_manager_option(help_string, default=None):
     is_flag=True,
     help="If specified, use local environment.",
 )
+
+INSTALL_JAVA = click.option(
+    "--install-java",
+    is_flag=True,
+    help="Install Java in the image. Default is False.",
+)
diff --git a/mlflow/utils/model_utils.py b/mlflow/utils/model_utils.py
index 80d149518ee9c..79c84605ef46d 100644
--- a/mlflow/utils/model_utils.py
+++ b/mlflow/utils/model_utils.py
@@ -22,18 +22,15 @@
 FLAVOR_CONFIG_CODE = "code"
 
 
-def _get_flavor_configuration(model_path, flavor_name):
-    """Obtains the configuration for the specified flavor from the specified
-    MLflow model path. If the model does not contain the specified flavor,
-    an exception will be thrown.
+def _get_all_flavor_configurations(model_path):
+    """Obtains all the flavor configurations from the specified MLflow model path.
 
     Args:
         model_path: The path to the root directory of the MLflow model for which to load
             the specified flavor configuration.
-        flavor_name: The name of the flavor configuration to load.
 
     Returns:
-        The flavor configuration as a dictionary.
+        The dictionary contains all flavor configurations with flavor name as key.
 
     """
     model_configuration_path = os.path.join(model_path, MLMODEL_FILE_NAME)
@@ -44,12 +41,30 @@ def _get_flavor_configuration(model_path, flavor_name):
         )
 
     model_conf = Model.load(model_configuration_path)
-    if flavor_name not in model_conf.flavors:
+    return model_conf.flavors
+
+
+def _get_flavor_configuration(model_path, flavor_name):
+    """Obtains the configuration for the specified flavor from the specified
+    MLflow model path. If the model does not contain the specified flavor,
+    an exception will be thrown.
+
+    Args:
+        model_path: The path to the root directory of the MLflow model for which to load
+            the specified flavor configuration.
+        flavor_name: The name of the flavor configuration to load.
+
+    Returns:
+        The flavor configuration as a dictionary.
+
+    """
+    flavors = _get_all_flavor_configurations(model_path)
+    if flavor_name not in flavors:
         raise MlflowException(
             f'Model does not have the "{flavor_name}" flavor',
             RESOURCE_DOES_NOT_EXIST,
         )
-    return model_conf.flavors[flavor_name]
+    return flavors[flavor_name]
 
 
 def _get_flavor_configuration_from_uri(model_uri, flavor_name, logger):
diff --git a/pyproject.toml b/pyproject.toml
index 298f6e237a84a..708c25301ae55 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -89,6 +89,7 @@ extend-exclude = [
 "examples/*" = ["T20"]
 "docs/broken_links.py" = ["T20"]
 "mlflow/*" = ["PT018"]
+"tests/pyfunc/docker/test_docker_flavors.py" = ["F401"] # For pytest fixture imports
 
 [tool.ruff.flake8-pytest-style]
 mark-parentheses = false
diff --git a/tests/diviner/test_diviner_model_export.py b/tests/diviner/test_diviner_model_export.py
index 983a16b6dcb34..d4e8cc50d8142 100644
--- a/tests/diviner/test_diviner_model_export.py
+++ b/tests/diviner/test_diviner_model_export.py
@@ -79,6 +79,17 @@ def diviner_custom_env(tmp_path):
     return conda_env
 
 
+@pytest.fixture
+def diviner_groups(diviner_data):
+    groups = []
+    for i in [0, -1]:
+        key_entries = []
+        for value in diviner_data.df[diviner_data.key_columns].iloc[[i]].to_dict().values():
+            key_entries.append(list(value.values())[0])
+        groups.append(tuple(key_entries))
+    return groups
+
+
 def test_diviner_native_save_and_load(grouped_prophet, model_path):
     mlflow.diviner.save_model(diviner_model=grouped_prophet, path=model_path)
 
@@ -137,37 +148,27 @@ def test_diviner_pyfunc_invalid_config_raises(grouped_prophet, model_path):
         loaded_pyfunc_model.predict(bad_conf)
 
 
-def test_diviner_pyfunc_group_predict_prophet(grouped_prophet, model_path, diviner_data):
-    groups = []
-    for i in [0, -1]:
-        key_entries = []
-        for value in diviner_data.df[diviner_data.key_columns].iloc[[i]].to_dict().values():
-            key_entries.append(list(value.values())[0])
-        groups.append(tuple(key_entries))
-
+def test_diviner_pyfunc_group_predict_prophet(grouped_prophet, model_path, diviner_groups):
     mlflow.diviner.save_model(diviner_model=grouped_prophet, path=model_path)
     loaded_pyfunc_model = pyfunc.load_pyfunc(model_uri=model_path)
 
-    local_group_pred = grouped_prophet.predict_groups(groups=groups, horizon=10, frequency="D")
-    pyfunc_conf = pd.DataFrame({"groups": [groups], "horizon": 10, "frequency": "D"}, index=[0])
+    local_group_pred = grouped_prophet.predict_groups(
+        groups=diviner_groups, horizon=10, frequency="D"
+    )
+    pyfunc_conf = pd.DataFrame(
+        {"groups": [diviner_groups], "horizon": 10, "frequency": "D"}, index=[0]
+    )
     pyfunc_group_predict = loaded_pyfunc_model.predict(pyfunc_conf)
 
     pd.testing.assert_frame_equal(local_group_pred, pyfunc_group_predict)
 
 
-def test_diviner_pyfunc_group_predict_pmdarima(grouped_pmdarima, model_path, diviner_data):
-    groups = []
-    for i in [0, -1]:
-        key_entries = []
-        for value in diviner_data.df[diviner_data.key_columns].iloc[[i]].to_dict().values():
-            key_entries.append(list(value.values())[0])
-        groups.append(tuple(key_entries))
-
+def test_diviner_pyfunc_group_predict_pmdarima(grouped_pmdarima, model_path, diviner_groups):
     mlflow.diviner.save_model(diviner_model=grouped_pmdarima, path=model_path)
     loaded_pyfunc_model = pyfunc.load_pyfunc(model_uri=model_path)
 
     local_group_pred = grouped_pmdarima.predict_groups(
-        groups=groups,
+        groups=diviner_groups,
         n_periods=10,
         predict_col="prediction",
         alpha=0.1,
@@ -176,7 +177,7 @@ def test_diviner_pyfunc_group_predict_pmdarima(grouped_pmdarima, model_path, div
     )
     pyfunc_conf = pd.DataFrame(
         {
-            "groups": [groups],
+            "groups": [diviner_groups],
             "n_periods": 10,
             "predict_col": "prediction",
             "alpha": 0.1,
@@ -419,7 +420,7 @@ def test_pmdarima_pyfunc_serve_and_score(grouped_prophet):
     pd.testing.assert_frame_equal(local_predict, scores)
 
 
-def test_pmdarima_pyfunc_serve_and_score_groups(grouped_prophet, diviner_data):
+def test_pmdarima_pyfunc_serve_and_score_groups(grouped_prophet, diviner_groups):
     artifact_path = "model"
     with mlflow.start_run():
         mlflow.diviner.log_model(
@@ -428,16 +429,11 @@ def test_pmdarima_pyfunc_serve_and_score_groups(grouped_prophet, diviner_data):
         )
         model_uri = mlflow.get_artifact_uri(artifact_path)
 
-    groups = []
-    for i in [0, -1]:
-        key_entries = []
-        for value in diviner_data.df[diviner_data.key_columns].iloc[[i]].to_dict().values():
-            key_entries.append(list(value.values())[0])
-        groups.append(tuple(key_entries))
+    local_predict = grouped_prophet.predict_groups(groups=diviner_groups, horizon=10, frequency="W")
 
-    local_predict = grouped_prophet.predict_groups(groups=groups, horizon=10, frequency="W")
-
-    inference_data = pd.DataFrame({"groups": [groups], "horizon": 10, "frequency": "W"}, index=[0])
+    inference_data = pd.DataFrame(
+        {"groups": [diviner_groups], "horizon": 10, "frequency": "W"}, index=[0]
+    )
 
     from mlflow.deployments import PredictionsResponse
 
diff --git a/tests/pyfunc/docker/__init__.py b/tests/pyfunc/docker/__init__.py
new file mode 100644
index 0000000000000..e69de29bb2d1d
diff --git a/tests/pyfunc/docker/conftest.py b/tests/pyfunc/docker/conftest.py
new file mode 100644
index 0000000000000..602cedfb0ef5f
--- /dev/null
+++ b/tests/pyfunc/docker/conftest.py
@@ -0,0 +1,56 @@
+import os
+from functools import lru_cache
+
+import pytest
+import requests
+from packaging.version import Version
+
+import docker
+import mlflow
+
+TEST_IMAGE_NAME = "test_image"
+MLFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+RESOURCE_DIR = os.path.join(MLFLOW_ROOT, "tests", "resources", "dockerfile")
+
+docker_client = docker.from_env()
+
+
+@pytest.fixture(autouse=True)
+def clean_up_docker_image():
+    try:
+        # Get all containers using the test image
+        containers = docker_client.containers.list(filters={"ancestor": TEST_IMAGE_NAME})
+        for container in containers:
+            container.remove(force=True)
+
+        # Clean up the image
+        docker_client.images.remove(TEST_IMAGE_NAME, force=True)
+    except Exception:
+        pass
+
+
+@lru_cache(maxsize=1)
+def get_released_mlflow_version():
+    url = "https://pypi.org/pypi/mlflow/json"
+    response = requests.get(url)
+    response.raise_for_status()
+    data = response.json()
+    versions = [Version(v) for v in data["releases"].keys() if not v.endswith(("a", "b", "rc"))]
+    return str(sorted(versions, reverse=True)[0])
+
+
+def save_model_with_latest_mlflow_version(flavor, **kwargs):
+    """
+    Save a model with overriding MLflow version from dev version to the latest released version.
+    By default a model is saved with the dev version of MLflow, which is not available on PyPI.
+    Usually we can be workaround this by adding --serve-wheel flag that starts local PyPI server,
+    however, this doesn't work when installing dependencies inside Docker container. Hence, this
+    function uses `extra_pip_requirements` to save the model with the latest released MLflow.
+    """
+    latest_mlflow_version = get_released_mlflow_version()
+    if flavor == "langchain":
+        kwargs["pip_requirements"] = [f"mlflow[gateway]=={latest_mlflow_version}", "langchain"]
+    else:
+        kwargs["extra_pip_requirements"] = [f"mlflow=={latest_mlflow_version}"]
+    flavor_module = getattr(mlflow, flavor)
+    flavor_module.save_model(**kwargs)
diff --git a/tests/pyfunc/docker/test_docker.py b/tests/pyfunc/docker/test_docker.py
new file mode 100644
index 0000000000000..e076449b4f424
--- /dev/null
+++ b/tests/pyfunc/docker/test_docker.py
@@ -0,0 +1,129 @@
+import difflib
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+from unittest import mock
+
+import pytest
+import sklearn
+import sklearn.neighbors
+
+import mlflow
+from mlflow.models import Model
+from mlflow.models.docker_utils import build_image_from_context
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.env_manager import VIRTUALENV
+from mlflow.version import VERSION
+
+from tests.pyfunc.docker.conftest import RESOURCE_DIR, get_released_mlflow_version
+
+
+def assert_dockerfiles_equal(actual_dockerfile_path: Path, expected_dockerfile_path: Path):
+    actual_dockerfile = actual_dockerfile_path.read_text().replace(
+        VERSION, get_released_mlflow_version()
+    )
+    expected_dockerfile = (
+        expected_dockerfile_path.read_text()
+        .replace("$MLFLOW_VERSION", get_released_mlflow_version())
+        .replace("$PYTHON_VERSION", PYTHON_VERSION)
+    )
+    assert (
+        actual_dockerfile == expected_dockerfile
+    ), "Generated Dockerfile does not match expected one. Diff:\n" + "\n".join(
+        difflib.unified_diff(expected_dockerfile.splitlines(), actual_dockerfile.splitlines())
+    )
+
+
+def save_model(tmp_path):
+    knn_model = sklearn.neighbors.KNeighborsClassifier()
+    model_path = os.path.join(tmp_path, "model")
+    mlflow.sklearn.save_model(
+        knn_model,
+        path=model_path,
+        pip_requirements=[
+            f"mlflow=={get_released_mlflow_version()}",
+            f"scikit-learn=={sklearn.__version__}",
+        ],  # Skip requirements inference for speed up
+    )
+    return model_path
+
+
+def add_spark_flavor_to_model(model_path):
+    model_config_path = os.path.join(model_path, "MLmodel")
+    model = Model.load(model_config_path)
+    model.add_flavor("spark", spark_version="3.5.0")
+    model.save(model_config_path)
+
+
+@dataclass
+class Param:
+    expected_dockerfile: str
+    env_manager: str = VIRTUALENV
+    mlflow_home: Optional[str] = None
+    install_mlflow: bool = False
+    enable_mlserver: bool = False
+    # If True, image is built with --model-uri param
+    specify_model_uri: bool = True
+
+
+_TEST_PARAMS = [
+    Param(expected_dockerfile="Dockerfile_default"),
+    Param(install_mlflow=True, expected_dockerfile="Dockerfile_install_mlflow"),
+    Param(enable_mlserver=True, expected_dockerfile="Dockerfile_enable_mlserver"),
+    Param(mlflow_home=".", expected_dockerfile="Dockerfile_with_mlflow_home"),
+    Param(specify_model_uri=False, expected_dockerfile="Dockerfile_no_model_uri"),
+]
+
+
+@pytest.mark.parametrize("params", _TEST_PARAMS)
+def test_build_image(tmp_path, params):
+    model_uri = save_model(tmp_path) if params.specify_model_uri else None
+
+    backend = get_flavor_backend(model_uri, docker_build=True, env_manager=params.env_manager)
+
+    # Copy the context dir to a temp dir so we can verify the generated Dockerfile
+    def _build_image_with_copy(context_dir, image_name):
+        # Replace mlflow dev version in Dockerfile with the latest released one
+        dockerfile = Path(context_dir) / "Dockerfile"
+        content = dockerfile.read_text()
+        content = content.replace(VERSION, get_released_mlflow_version())
+        dockerfile.write_text(content)
+
+        shutil.copytree(context_dir, dst_dir)
+        build_image_from_context(context_dir, image_name)
+
+    dst_dir = tmp_path / "context"
+    with mock.patch(
+        "mlflow.models.docker_utils.build_image_from_context",
+        side_effect=_build_image_with_copy,
+    ):
+        backend.build_image(
+            model_uri=model_uri,
+            image_name="test_image",
+            mlflow_home=params.mlflow_home,
+            install_mlflow=params.install_mlflow,
+            enable_mlserver=params.enable_mlserver,
+        )
+
+    actual = dst_dir / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / params.expected_dockerfile
+    assert_dockerfiles_equal(actual, expected)
+
+
+def test_generate_dockerfile_for_java_flavor(tmp_path):
+    model_path = save_model(tmp_path)
+    add_spark_flavor_to_model(model_path)
+
+    backend = get_flavor_backend(model_path, docker_build=True)
+
+    backend.generate_dockerfile(
+        model_uri=model_path,
+        output_dir=tmp_path,
+    )
+
+    actual = tmp_path / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / "Dockerfile_java_flavor"
+    assert_dockerfiles_equal(actual, expected)
diff --git a/tests/pyfunc/docker/test_docker_flavors.py b/tests/pyfunc/docker/test_docker_flavors.py
new file mode 100644
index 0000000000000..caa444ad2c934
--- /dev/null
+++ b/tests/pyfunc/docker/test_docker_flavors.py
@@ -0,0 +1,407 @@
+"""
+This test class is used for comprehensive testing of serving docker images for all MLflow flavors.
+As such, it is not intended to be run on a regular basis and is skipped by default. Rather, it
+should be run manually when making changes to the core docker logic.
+
+To run this test, run the following command manually
+
+    $ pytest tests/pyfunc/test_docker_flavors.py
+
+"""
+
+import json
+import os
+import time
+from operator import itemgetter
+
+import pandas as pd
+import pytest
+import requests
+
+import mlflow
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+
+from tests.catboost.test_catboost_model_export import reg_model
+from tests.diviner.test_diviner_model_export import (
+    diviner_data,
+    diviner_groups,
+    grouped_prophet,
+)
+from tests.fastai.test_fastai_model_export import fastai_model as fastai_model_raw
+from tests.h2o.test_h2o_model_export import h2o_iris_model
+from tests.helper_functions import get_safe_port
+from tests.langchain.test_langchain_model_export import fake_chat_model
+from tests.lightgbm.test_lightgbm_model_export import lgb_model
+from tests.models.test_model import iris_data, sklearn_knn_model
+from tests.paddle.test_paddle_model_export import pd_model
+from tests.pmdarima.test_pmdarima_model_export import (
+    auto_arima_object_model,
+    test_data,
+)
+from tests.prophet.test_prophet_model_export import prophet_model as prophet_raw_model
+from tests.pyfunc.docker.conftest import (
+    MLFLOW_ROOT,
+    TEST_IMAGE_NAME,
+    docker_client,
+    save_model_with_latest_mlflow_version,
+)
+from tests.spacy.test_spacy_model_export import spacy_model_with_data
+from tests.spark.test_spark_model_export import (
+    iris_df,
+    spark,
+    spark_model_iris,
+)
+from tests.statsmodels.model_fixtures import ols_model
+from tests.tensorflow.test_tensorflow2_core_model_export import tf2_toy_model
+from tests.transformers.helper import load_small_seq2seq_pipeline
+
+
+@pytest.mark.skipif(os.getenv("GITHUB_ACTIONS") == "true", reason="Time consuming tests")
+@pytest.mark.parametrize(
+    ("flavor"),
+    [
+        "catboost",
+        "diviner",
+        "fastai",
+        "h2o",
+        # "johnsnowlabs", # Couldn't test JohnSnowLab locally due to license issue
+        "keras",
+        "langchain",
+        "lightgbm",
+        # "mleap", # Mleap model logging is deprecated since 2.6.1
+        "onnx",
+        # "openai", # OPENAI API KEY is not necessarily available for everyone
+        "paddle",
+        "pmdarima",
+        "prophet",
+        "pyfunc",
+        "pytorch",
+        "sklearn",
+        "spacy",
+        "spark",
+        "statsmodels",
+        "tensorflow",
+        "transformers",
+    ],
+)
+def test_build_image_and_serve(flavor, request):
+    model_path = request.getfixturevalue(f"{flavor}_model")
+
+    # Build an image
+    backend = get_flavor_backend(model_uri=model_path, docker_build=True)
+    backend.build_image(
+        model_uri=model_path,
+        image_name=TEST_IMAGE_NAME,
+        mlflow_home=MLFLOW_ROOT,  # Required to prevent installing dev version of MLflow from PyPI
+    )
+
+    # Run a container
+    port = get_safe_port()
+    docker_client.containers.run(
+        image=TEST_IMAGE_NAME,
+        ports={8080: port},
+        detach=True,
+    )
+
+    # Wait until the container to start
+    timeout = 180
+    start_time = time.time()
+    success = False
+    while time.time() < start_time + timeout:
+        try:
+            # Send health check
+            response = requests.get(url=f"http://localhost:{port}/ping")
+            if response.status_code == 200:
+                success = True
+                break
+        except Exception:
+            time.sleep(5)
+
+    if not success:
+        raise TimeoutError(f"Container did not start in {timeout} seconds")
+
+    # Make a scoring request with a saved input example
+    with open(os.path.join(model_path, "input_example.json")) as f:
+        input_example = json.load(f)
+
+    # Wrap Pandas dataframe in a proper payload format
+    if "columns" in input_example or "data" in input_example:
+        input_example = {"dataframe_split": input_example}
+
+    response = requests.post(
+        url=f"http://localhost:{port}/invocations",
+        data=json.dumps(input_example),
+        headers={"Content-Type": "application/json"},
+    )
+
+    assert response.status_code == 200, f"Response: {response.text}"
+    assert "predictions" in response.json(), f"Response: {response.text}"
+
+
+@pytest.fixture
+def catboost_model(tmp_path, reg_model):
+    model_path = str(tmp_path / "catboost_model")
+    save_model_with_latest_mlflow_version(
+        flavor="catboost",
+        cb_model=reg_model.model,
+        path=model_path,
+        input_example=reg_model.inference_dataframe[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def diviner_model(tmp_path, grouped_prophet, diviner_groups):
+    model_path = str(tmp_path / "diviner_model")
+
+    save_model_with_latest_mlflow_version(
+        flavor="diviner",
+        diviner_model=grouped_prophet,
+        path=model_path,
+        input_example=pd.DataFrame(
+            {"groups": [diviner_groups], "horizon": 10, "frequency": "D"}, index=[0]
+        ),
+    )
+    return model_path
+
+
+@pytest.fixture
+def fastai_model(tmp_path, fastai_model_raw):
+    model_path = str(tmp_path / "fastai_model")
+    save_model_with_latest_mlflow_version(
+        flavor="fastai",
+        fastai_learner=fastai_model_raw.model,
+        path=model_path,
+        input_example=fastai_model_raw.inference_dataframe[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def h2o_model(tmp_path, h2o_iris_model):
+    model_path = str(tmp_path / "h2o_model")
+    save_model_with_latest_mlflow_version(
+        flavor="h2o",
+        h2o_model=h2o_iris_model.model,
+        path=model_path,
+        input_example=h2o_iris_model.inference_data.as_data_frame()[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def keras_model(tmp_path, iris_data):
+    from sklearn import datasets
+    from tensorflow.keras.layers import Dense
+    from tensorflow.keras.models import Sequential
+
+    model_path = str(tmp_path / "keras_model")
+
+    model = Sequential()
+    model.add(Dense(3, input_dim=4))
+    model.add(Dense(1))
+
+    X, y = datasets.load_iris(return_X_y=True)
+    save_model_with_latest_mlflow_version(
+        flavor="tensorflow",
+        model=model,
+        path=model_path,
+        input_example=X[:3, :],
+    )
+    return model_path
+
+
+@pytest.fixture
+def langchain_model(tmp_path):
+    from langchain.schema.runnable import RunnablePassthrough
+
+    chain = RunnablePassthrough() | itemgetter("messages")
+    model_path = str(tmp_path / "langchain_model")
+    save_model_with_latest_mlflow_version(
+        flavor="langchain", lc_model=chain, path=model_path, input_example={"messages": "Hi"}
+    )
+    return model_path
+
+
+@pytest.fixture
+def lightgbm_model(tmp_path, lgb_model):
+    model_path = str(tmp_path / "lightgbm_model")
+    save_model_with_latest_mlflow_version(
+        flavor="lightgbm",
+        lgb_model=lgb_model.model,
+        path=model_path,
+        input_example=lgb_model.inference_dataframe.to_numpy()[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def onnx_model(tmp_path):
+    import numpy as np
+    import onnx
+    import torch
+    from torch import nn
+
+    model = torch.nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))
+    onnx_model_path = os.path.join(tmp_path, "torch_onnx")
+    torch.onnx.export(
+        model,
+        torch.randn(1, 4),
+        onnx_model_path,
+        dynamic_axes={"input": {0: "batch"}},
+        input_names=["input"],
+    )
+    onnx_model = onnx.load(onnx_model_path)
+
+    model_path = str(tmp_path / "onnx_model")
+    save_model_with_latest_mlflow_version(
+        flavor="onnx",
+        onnx_model=onnx_model,
+        path=model_path,
+        input_example=np.random.rand(1, 4).astype(np.float32),
+    )
+    return model_path
+
+
+@pytest.fixture
+def paddle_model(tmp_path, pd_model):
+    model_path = str(tmp_path / "paddle_model")
+    save_model_with_latest_mlflow_version(
+        flavor="paddle",
+        pd_model=pd_model.model,
+        path=model_path,
+        input_example=pd_model.inference_dataframe[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def pmdarima_model(tmp_path, auto_arima_object_model):
+    model_path = str(tmp_path / "pmdarima_model")
+    save_model_with_latest_mlflow_version(
+        flavor="pmdarima",
+        pmdarima_model=auto_arima_object_model,
+        path=model_path,
+        input_example=pd.DataFrame({"n_periods": [30]}),
+    )
+    return model_path
+
+
+@pytest.fixture
+def prophet_model(tmp_path, prophet_raw_model):
+    model_path = str(tmp_path / "prophet_model")
+    save_model_with_latest_mlflow_version(
+        flavor="prophet",
+        pr_model=prophet_raw_model.model,
+        path=model_path,
+        input_example=prophet_raw_model.data[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def pyfunc_model(tmp_path):
+    class CustomModel(mlflow.pyfunc.PythonModel):
+        def __init__(self):
+            pass
+
+        def predict(self, context, model_input):
+            return model_input
+
+    model_path = str(tmp_path / "pyfunc_model")
+    save_model_with_latest_mlflow_version(
+        flavor="pyfunc",
+        python_model=CustomModel(),
+        path=model_path,
+        input_example=[1, 2, 3],
+    )
+    return model_path
+
+
+@pytest.fixture
+def pytorch_model(tmp_path):
+    from torch import nn, randn
+
+    model = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))
+    model_path = str(tmp_path / "pytorch_model")
+    save_model_with_latest_mlflow_version(
+        flavor="pytorch",
+        pytorch_model=model,
+        path=model_path,
+        input_example=randn(1, 4).numpy(),
+    )
+    return model_path
+
+
+@pytest.fixture
+def sklearn_model(tmp_path, sklearn_knn_model, iris_data):
+    model_path = str(tmp_path / "sklearn_model")
+    save_model_with_latest_mlflow_version(
+        flavor="sklearn",
+        sk_model=sklearn_knn_model,
+        path=model_path,
+        input_example=iris_data[0][:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def spacy_model(tmp_path, spacy_model_with_data):
+    model_path = str(tmp_path / "spacy_model")
+    save_model_with_latest_mlflow_version(
+        flavor="spacy",
+        spacy_model=spacy_model_with_data.model,
+        path=model_path,
+        input_example=spacy_model_with_data.inference_data[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def spark_model(tmp_path, spark_model_iris):
+    model_path = str(tmp_path / "spark_model")
+    save_model_with_latest_mlflow_version(
+        flavor="spark",
+        spark_model=spark_model_iris.model,
+        path=model_path,
+        input_example=spark_model_iris.spark_df.toPandas()[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def statsmodels_model(tmp_path):
+    model = ols_model()
+    model_path = str(tmp_path / "statsmodels_model")
+    save_model_with_latest_mlflow_version(
+        flavor="statsmodels",
+        statsmodels_model=model.model,
+        path=model_path,
+        input_example=model.inference_dataframe[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def tensorflow_model(tmp_path, tf2_toy_model):
+    model_path = str(tmp_path / "tensorflow_model")
+    save_model_with_latest_mlflow_version(
+        flavor="tensorflow",
+        model=tf2_toy_model.model,
+        path=model_path,
+        input_example=tf2_toy_model.inference_data[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def transformers_model(tmp_path):
+    pipeline = load_small_seq2seq_pipeline()
+    model_path = str(tmp_path / "transformers_model")
+    save_model_with_latest_mlflow_version(
+        flavor="transformers",
+        transformers_model=pipeline,
+        path=model_path,
+        input_example="hi",
+    )
+    return model_path
diff --git a/tests/pyfunc/test_docker.py b/tests/pyfunc/test_docker.py
deleted file mode 100644
index cdab820ff844a..0000000000000
--- a/tests/pyfunc/test_docker.py
+++ /dev/null
@@ -1,128 +0,0 @@
-import difflib
-import os
-import shutil
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Optional
-from unittest import mock
-
-import pytest
-import sklearn
-import sklearn.neighbors
-
-import mlflow
-from mlflow.models.flavor_backend_registry import get_flavor_backend
-from mlflow.utils.env_manager import CONDA, VIRTUALENV
-from mlflow.version import VERSION
-
-_MLFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
-_RESOURCE_DIR = os.path.join(_MLFLOW_ROOT, "tests", "resources", "dockerfile")
-
-
-def assert_dockerfiles_equal(actual_dockerfile_path: Path, expected_dockerfile_path: Path):
-    actual_dockerfile = actual_dockerfile_path.read_text()
-    expected_dockerfile = expected_dockerfile_path.read_text().replace("$VERSION", VERSION)
-
-    assert (
-        actual_dockerfile == expected_dockerfile
-    ), "Generated Dockerfile does not match expected one. Diff:\n" + "\n".join(
-        difflib.unified_diff(expected_dockerfile.splitlines(), actual_dockerfile.splitlines())
-    )
-
-
-def log_model():
-    with mlflow.start_run():
-        knn_model = sklearn.neighbors.KNeighborsClassifier()
-        model_info = mlflow.sklearn.log_model(
-            knn_model,
-            "model",
-            pip_requirements=[
-                f"mlflow=={mlflow.__version__}",
-                f"scikit-learn=={sklearn.__version__}",
-            ],  # Skip requirements inference for speed up
-        )
-        return model_info.model_uri
-
-
-@dataclass
-class Param:
-    expected_dockerfile: str
-    env_manager: str = VIRTUALENV
-    mlflow_home: Optional[str] = None
-    install_mlflow: bool = False
-    enable_mlserver: bool = False
-    # If True, image is built with --model-uri param
-    specify_model_uri: bool = True
-
-
-_TEST_PARAMS = [
-    Param(
-        expected_dockerfile="Dockerfile_default",
-    ),
-    Param(
-        env_manager=CONDA,
-        expected_dockerfile="Dockerfile_conda",
-    ),
-    Param(
-        mlflow_home=".",
-        expected_dockerfile="Dockerfile_with_mlflow_home",
-    ),
-    Param(
-        install_mlflow=True,
-        expected_dockerfile="Dockerfile_install_mlflow",
-    ),
-    Param(
-        mlflow_home=_MLFLOW_ROOT,
-        install_mlflow=True,
-        expected_dockerfile="Dockerfile_install_mlflow_from_mlflow_home",
-    ),
-    Param(
-        enable_mlserver=True,
-        expected_dockerfile="Dockerfile_enable_mlserver",
-    ),
-    Param(specify_model_uri=False, expected_dockerfile="Dockerfile_no_model_uri"),
-]
-
-
-@pytest.mark.parametrize("params", _TEST_PARAMS)
-def test_generate_dockerfile(tmp_path, params):
-    model_uri = log_model() if params.specify_model_uri else None
-
-    backend = get_flavor_backend(model_uri, docker_build=True, env_manager=params.env_manager)
-
-    backend.generate_dockerfile(
-        model_uri=model_uri,
-        output_dir=tmp_path,
-        mlflow_home=params.mlflow_home,
-        install_mlflow=params.install_mlflow,
-        enable_mlserver=params.enable_mlserver,
-    )
-
-    actual = tmp_path / "Dockerfile"
-    expected = Path(_RESOURCE_DIR) / params.expected_dockerfile
-    assert_dockerfiles_equal(actual, expected)
-
-
-@pytest.mark.parametrize("params", _TEST_PARAMS)
-def test_build_image(tmp_path, params):
-    model_uri = log_model() if params.specify_model_uri else None
-
-    backend = get_flavor_backend(model_uri, docker_build=True, env_manager=params.env_manager)
-
-    # Copy the context dir to a temp dir so we can verify the generated Dockerfile
-    dst_dir = tmp_path / "context"
-    with mock.patch(
-        "mlflow.models.docker_utils.build_image_from_context",
-        side_effect=lambda context_dir, image_name: shutil.copytree(context_dir, dst_dir),
-    ):
-        backend.build_image(
-            model_uri=model_uri,
-            image_name="test_image",
-            mlflow_home=params.mlflow_home,
-            install_mlflow=params.install_mlflow,
-            enable_mlserver=params.enable_mlserver,
-        )
-
-    actual = dst_dir / "Dockerfile"
-    expected = Path(_RESOURCE_DIR) / params.expected_dockerfile
-    assert_dockerfiles_equal(actual, expected)
diff --git a/tests/resources/dockerfile/Dockerfile_conda b/tests/resources/dockerfile/Dockerfile_conda
deleted file mode 100644
index a2fe3e4649f59..0000000000000
--- a/tests/resources/dockerfile/Dockerfile_conda
+++ /dev/null
@@ -1,39 +0,0 @@
-
-# Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
-
-# Setup miniconda
-RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh >> miniconda.sh
-RUN bash ./miniconda.sh -b -p /miniconda && rm ./miniconda.sh
-ENV PATH="/miniconda/bin:$PATH"
-
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
-WORKDIR /opt/mlflow
-
-RUN pip install mlflow==$VERSION
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:pom -DoutputDirectory=/opt/java 
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:jar -DoutputDirectory=/opt/java/jars 
-RUN cp /opt/java/mlflow-scoring-$VERSION.pom /opt/java/pom.xml
-RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirectory=/opt/java/jars 
-
-
-COPY model_dir/model /opt/ml/model
-RUN python -c "from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='conda');"
-
-ENV MLFLOW_DISABLE_ENV_CREATION=True
-ENV ENABLE_MLSERVER=False
-
-
-# granting read/write access and conditional execution authority to all child directories
-# and files to allow for deployment to AWS Sagemaker Serverless Endpoints
-# (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
-RUN chmod o+rwX /opt/mlflow/
-
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('conda')"]
diff --git a/tests/resources/dockerfile/Dockerfile_default b/tests/resources/dockerfile/Dockerfile_default
index 4d2244bc3602a..d31c6e2748592 100644
--- a/tests/resources/dockerfile/Dockerfile_default
+++ b/tests/resources/dockerfile/Dockerfile_default
@@ -1,52 +1,29 @@
-
 # Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
-
-# Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
-    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
-    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
-RUN git clone \
-    --depth 1 \
-    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
-    https://github.com/pyenv/pyenv.git /root/.pyenv
-ENV PYENV_ROOT="/root/.pyenv"
-ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
-RUN pip install virtualenv
-
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
-WORKDIR /opt/mlflow
+FROM python:$PYTHON_VERSION-slim
+
+RUN apt-get -y update && apt-get install -y --no-install-recommends nginx
+
 
-RUN pip install mlflow==$VERSION
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:pom -DoutputDirectory=/opt/java 
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:jar -DoutputDirectory=/opt/java/jars 
-RUN cp /opt/java/mlflow-scoring-$VERSION.pom /opt/java/pom.xml
-RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirectory=/opt/java/jars 
 
+WORKDIR /opt/mlflow
+
+# Install MLflow
+RUN pip install mlflow==$MLFLOW_VERSION
 
+# Copy model to image and install dependencies
 COPY model_dir/model /opt/ml/model
-RUN python -c "from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='virtualenv');"
+RUN python -c "from mlflow.models import container as C; C._install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='local');"
 
 ENV MLFLOW_DISABLE_ENV_CREATION=True
 ENV ENABLE_MLSERVER=False
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('virtualenv')"]
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('local')"]
diff --git a/tests/resources/dockerfile/Dockerfile_enable_mlserver b/tests/resources/dockerfile/Dockerfile_enable_mlserver
index 66f3c977a2b19..d4a39a69a164c 100644
--- a/tests/resources/dockerfile/Dockerfile_enable_mlserver
+++ b/tests/resources/dockerfile/Dockerfile_enable_mlserver
@@ -1,52 +1,29 @@
-
 # Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
-
-# Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
-    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
-    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
-RUN git clone \
-    --depth 1 \
-    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
-    https://github.com/pyenv/pyenv.git /root/.pyenv
-ENV PYENV_ROOT="/root/.pyenv"
-ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
-RUN pip install virtualenv
-
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
-WORKDIR /opt/mlflow
+FROM python:$PYTHON_VERSION-slim
+
+RUN apt-get -y update && apt-get install -y --no-install-recommends nginx
+
 
-RUN pip install mlflow==$VERSION
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:pom -DoutputDirectory=/opt/java 
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:jar -DoutputDirectory=/opt/java/jars 
-RUN cp /opt/java/mlflow-scoring-$VERSION.pom /opt/java/pom.xml
-RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirectory=/opt/java/jars 
 
+WORKDIR /opt/mlflow
+
+# Install MLflow
+RUN pip install mlflow==$MLFLOW_VERSION
 
+# Copy model to image and install dependencies
 COPY model_dir/model /opt/ml/model
-RUN python -c "from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=True, env_manager='virtualenv');"
+RUN python -c "from mlflow.models import container as C; C._install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=True, env_manager='local');"
 
 ENV MLFLOW_DISABLE_ENV_CREATION=True
 ENV ENABLE_MLSERVER=True
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('virtualenv')"]
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('local')"]
diff --git a/tests/resources/dockerfile/Dockerfile_install_mlflow b/tests/resources/dockerfile/Dockerfile_install_mlflow
index f252b7468cbfa..f3787629ebaaa 100644
--- a/tests/resources/dockerfile/Dockerfile_install_mlflow
+++ b/tests/resources/dockerfile/Dockerfile_install_mlflow
@@ -1,52 +1,29 @@
-
 # Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
-
-# Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
-    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
-    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
-RUN git clone \
-    --depth 1 \
-    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
-    https://github.com/pyenv/pyenv.git /root/.pyenv
-ENV PYENV_ROOT="/root/.pyenv"
-ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
-RUN pip install virtualenv
-
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
-WORKDIR /opt/mlflow
+FROM python:$PYTHON_VERSION-slim
+
+RUN apt-get -y update && apt-get install -y --no-install-recommends nginx
+
 
-RUN pip install mlflow==$VERSION
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:pom -DoutputDirectory=/opt/java 
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:jar -DoutputDirectory=/opt/java/jars 
-RUN cp /opt/java/mlflow-scoring-$VERSION.pom /opt/java/pom.xml
-RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirectory=/opt/java/jars 
 
+WORKDIR /opt/mlflow
+
+# Install MLflow
+RUN pip install mlflow==$MLFLOW_VERSION
 
+# Copy model to image and install dependencies
 COPY model_dir/model /opt/ml/model
-RUN python -c "from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=True, enable_mlserver=False, env_manager='virtualenv');"
+RUN python -c "from mlflow.models import container as C; C._install_pyfunc_deps('/opt/ml/model', install_mlflow=True, enable_mlserver=False, env_manager='local');"
 
 ENV MLFLOW_DISABLE_ENV_CREATION=True
 ENV ENABLE_MLSERVER=False
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('virtualenv')"]
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('local')"]
diff --git a/tests/resources/dockerfile/Dockerfile_install_mlflow_from_mlflow_home b/tests/resources/dockerfile/Dockerfile_install_mlflow_from_mlflow_home
deleted file mode 100644
index 17e15a9557f16..0000000000000
--- a/tests/resources/dockerfile/Dockerfile_install_mlflow_from_mlflow_home
+++ /dev/null
@@ -1,50 +0,0 @@
-
-# Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
-
-# Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
-    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
-    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
-RUN git clone \
-    --depth 1 \
-    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
-    https://github.com/pyenv/pyenv.git /root/.pyenv
-ENV PYENV_ROOT="/root/.pyenv"
-ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
-RUN pip install virtualenv
-
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
-WORKDIR /opt/mlflow
-
-COPY mlflow-project /opt/mlflow
-RUN pip install /opt/mlflow
-RUN cd /opt/mlflow/mlflow/java/scoring && mvn --batch-mode package -DskipTests  && mkdir -p /opt/java/jars && mv /opt/mlflow/mlflow/java/scoring/target/mlflow-scoring-*-with-dependencies.jar /opt/java/jars
-
-
-COPY model_dir/model /opt/ml/model
-RUN python -c "from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=True, enable_mlserver=False, env_manager='virtualenv');"
-
-ENV MLFLOW_DISABLE_ENV_CREATION=True
-ENV ENABLE_MLSERVER=False
-
-
-# granting read/write access and conditional execution authority to all child directories
-# and files to allow for deployment to AWS Sagemaker Serverless Endpoints
-# (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
-RUN chmod o+rwX /opt/mlflow/
-
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('virtualenv')"]
diff --git a/tests/resources/dockerfile/Dockerfile_java_flavor b/tests/resources/dockerfile/Dockerfile_java_flavor
new file mode 100644
index 0000000000000..a3883fb76b190
--- /dev/null
+++ b/tests/resources/dockerfile/Dockerfile_java_flavor
@@ -0,0 +1,55 @@
+# Build an image that can serve mlflow models.
+FROM ubuntu:20.04
+
+RUN apt-get -y update && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends wget curl nginx ca-certificates bzip2 build-essential cmake git-core
+
+# Setup pyenv
+RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata \
+    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
+    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
+RUN git clone \
+    --depth 1 \
+    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
+    https://github.com/pyenv/pyenv.git /root/.pyenv
+ENV PYENV_ROOT="/root/.pyenv"
+ENV PATH="$PYENV_ROOT/bin:$PATH"
+RUN apt install -y python3.8 python3.8-distutils \
+    && ln -s -f $(which python3.8) /usr/bin/python \
+    && wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py \
+    && python /tmp/get-pip.py
+RUN pip install virtualenv
+
+
+# Setup Java
+RUN apt-get install -y --no-install-recommends openjdk-8-jdk maven
+ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
+
+WORKDIR /opt/mlflow
+
+# Install MLflow
+RUN pip install mlflow==$MLFLOW_VERSION
+
+# Install Java mlflow-scoring from Maven Central
+RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$MLFLOW_VERSION:pom -DoutputDirectory=/opt/java 
+RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$MLFLOW_VERSION:jar -DoutputDirectory=/opt/java/jars 
+RUN cp /opt/java/mlflow-scoring-$MLFLOW_VERSION.pom /opt/java/pom.xml
+RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirectory=/opt/java/jars 
+
+
+# Copy model to image and install dependencies
+COPY model_dir/model /opt/ml/model
+RUN python -c "from mlflow.models import container as C; C._install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='virtualenv');"
+
+ENV MLFLOW_DISABLE_ENV_CREATION=True
+ENV ENABLE_MLSERVER=False
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
+
+# granting read/write access and conditional execution authority to all child directories
+# and files to allow for deployment to AWS Sagemaker Serverless Endpoints
+# (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
+RUN chmod o+rwX /opt/mlflow/
+
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('virtualenv')"]
diff --git a/tests/resources/dockerfile/Dockerfile_no_model_uri b/tests/resources/dockerfile/Dockerfile_no_model_uri
index 8573c1dcaa9e4..0488784b916a2 100644
--- a/tests/resources/dockerfile/Dockerfile_no_model_uri
+++ b/tests/resources/dockerfile/Dockerfile_no_model_uri
@@ -1,15 +1,10 @@
-
 # Build an image that can serve mlflow models.
 FROM ubuntu:20.04
 
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
+RUN apt-get -y update && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends wget curl nginx ca-certificates bzip2 build-essential cmake git-core
 
 # Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
+RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata \
     libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
     libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
 RUN git clone \
@@ -18,22 +13,26 @@ RUN git clone \
     https://github.com/pyenv/pyenv.git /root/.pyenv
 ENV PYENV_ROOT="/root/.pyenv"
 ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
+RUN apt install -y python3.8 python3.8-distutils \
+    && ln -s -f $(which python3.8) /usr/bin/python \
+    && wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py \
+    && python /tmp/get-pip.py
 RUN pip install virtualenv
 
 
+# Setup Java
+RUN apt-get install -y --no-install-recommends openjdk-8-jdk maven
 ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
+
 WORKDIR /opt/mlflow
 
-RUN pip install mlflow==$VERSION
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:pom -DoutputDirectory=/opt/java 
-RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$VERSION:jar -DoutputDirectory=/opt/java/jars 
-RUN cp /opt/java/mlflow-scoring-$VERSION.pom /opt/java/pom.xml
+# Install MLflow
+RUN pip install mlflow==$MLFLOW_VERSION
+
+# Install Java mlflow-scoring from Maven Central
+RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$MLFLOW_VERSION:pom -DoutputDirectory=/opt/java 
+RUN mvn --batch-mode dependency:copy -Dartifact=org.mlflow:mlflow-scoring:$MLFLOW_VERSION:jar -DoutputDirectory=/opt/java/jars 
+RUN cp /opt/java/mlflow-scoring-$MLFLOW_VERSION.pom /opt/java/pom.xml
 RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirectory=/opt/java/jars 
 
 
@@ -41,11 +40,14 @@ RUN cd /opt/java && mvn --batch-mode dependency:copy-dependencies -DoutputDirect
 
 ENV MLFLOW_DISABLE_ENV_CREATION=True
 ENV ENABLE_MLSERVER=False
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C;from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='virtualenv'); C._serve('virtualenv')"]
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='virtualenv'); C._serve('virtualenv')"]
diff --git a/tests/resources/dockerfile/Dockerfile_sagemaker_conda b/tests/resources/dockerfile/Dockerfile_sagemaker_conda
index 6f291ae878152..08e0af71d290d 100644
--- a/tests/resources/dockerfile/Dockerfile_sagemaker_conda
+++ b/tests/resources/dockerfile/Dockerfile_sagemaker_conda
@@ -1,10 +1,7 @@
-
 # Build an image that can serve mlflow models.
 FROM ubuntu:20.04
 
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
+RUN apt-get -y update && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends wget curl nginx ca-certificates bzip2 build-essential cmake git-core
 
 # Setup miniconda
 RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh >> miniconda.sh
@@ -12,25 +9,33 @@ RUN bash ./miniconda.sh -b -p /miniconda && rm ./miniconda.sh
 ENV PATH="/miniconda/bin:$PATH"
 
 
+# Setup Java
+RUN apt-get install -y --no-install-recommends openjdk-8-jdk maven
 ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
+
 WORKDIR /opt/mlflow
 
+# Install MLflow from local source
 COPY mlflow-project /opt/mlflow
 RUN pip install /opt/mlflow
+
+# Install Java mlflow-scoring from local source
 RUN cd /opt/mlflow/mlflow/java/scoring && mvn --batch-mode package -DskipTests  && mkdir -p /opt/java/jars && mv /opt/mlflow/mlflow/java/scoring/target/mlflow-scoring-*-with-dependencies.jar /opt/java/jars
 
 
+# Install minimal serving dependencies
 RUN python -c "from mlflow.models.container import _install_pyfunc_deps;_install_pyfunc_deps(None, False)"
 
 ENV MLFLOW_DISABLE_ENV_CREATION=False
 ENV ENABLE_MLSERVER=False
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
 ENTRYPOINT ["python", "-c", "import sys; from mlflow.models import container as C; C._init(sys.argv[1], 'conda')"]
diff --git a/tests/resources/dockerfile/Dockerfile_sagemaker_virtualenv b/tests/resources/dockerfile/Dockerfile_sagemaker_virtualenv
index 66ef8be19bf39..06b0c290a0765 100644
--- a/tests/resources/dockerfile/Dockerfile_sagemaker_virtualenv
+++ b/tests/resources/dockerfile/Dockerfile_sagemaker_virtualenv
@@ -1,15 +1,10 @@
-
 # Build an image that can serve mlflow models.
 FROM ubuntu:20.04
 
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
+RUN apt-get -y update && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends wget curl nginx ca-certificates bzip2 build-essential cmake git-core
 
 # Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
+RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata \
     libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
     libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
 RUN git clone \
@@ -18,32 +13,40 @@ RUN git clone \
     https://github.com/pyenv/pyenv.git /root/.pyenv
 ENV PYENV_ROOT="/root/.pyenv"
 ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
+RUN apt install -y python3.8 python3.8-distutils \
+    && ln -s -f $(which python3.8) /usr/bin/python \
+    && wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py \
+    && python /tmp/get-pip.py
 RUN pip install virtualenv
 
 
+# Setup Java
+RUN apt-get install -y --no-install-recommends openjdk-8-jdk maven
 ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
+
 WORKDIR /opt/mlflow
 
+# Install MLflow from local source
 COPY mlflow-project /opt/mlflow
 RUN pip install /opt/mlflow
+
+# Install Java mlflow-scoring from local source
 RUN cd /opt/mlflow/mlflow/java/scoring && mvn --batch-mode package -DskipTests  && mkdir -p /opt/java/jars && mv /opt/mlflow/mlflow/java/scoring/target/mlflow-scoring-*-with-dependencies.jar /opt/java/jars
 
 
+# Install minimal serving dependencies
 RUN python -c "from mlflow.models.container import _install_pyfunc_deps;_install_pyfunc_deps(None, False)"
 
 ENV MLFLOW_DISABLE_ENV_CREATION=False
 ENV ENABLE_MLSERVER=False
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
 ENTRYPOINT ["python", "-c", "import sys; from mlflow.models import container as C; C._init(sys.argv[1], 'virtualenv')"]
diff --git a/tests/resources/dockerfile/Dockerfile_with_mlflow_home b/tests/resources/dockerfile/Dockerfile_with_mlflow_home
index 2090528525552..7d995dc38e537 100644
--- a/tests/resources/dockerfile/Dockerfile_with_mlflow_home
+++ b/tests/resources/dockerfile/Dockerfile_with_mlflow_home
@@ -1,50 +1,30 @@
-
 # Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends          wget          curl          nginx          ca-certificates          bzip2          build-essential          cmake          openjdk-8-jdk          git-core          maven     && rm -rf /var/lib/apt/lists/*
-
-
-# Setup pyenv
-RUN apt -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata
-RUN apt-get install -y \
-    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
-    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
-RUN git clone \
-    --depth 1 \
-    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
-    https://github.com/pyenv/pyenv.git /root/.pyenv
-ENV PYENV_ROOT="/root/.pyenv"
-ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
-RUN pip install virtualenv
-
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
+FROM python:$PYTHON_VERSION-slim
+
+RUN apt-get -y update && apt-get install -y --no-install-recommends nginx
+
+
+
 WORKDIR /opt/mlflow
 
+# Install MLflow from local source
 COPY mlflow-project /opt/mlflow
 RUN pip install /opt/mlflow
-RUN cd /opt/mlflow/mlflow/java/scoring && mvn --batch-mode package -DskipTests  && mkdir -p /opt/java/jars && mv /opt/mlflow/mlflow/java/scoring/target/mlflow-scoring-*-with-dependencies.jar /opt/java/jars
-
 
+# Copy model to image and install dependencies
 COPY model_dir/model /opt/ml/model
-RUN python -c "from mlflow.models.container import _install_pyfunc_deps; _install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='virtualenv');"
+RUN python -c "from mlflow.models import container as C; C._install_pyfunc_deps('/opt/ml/model', install_mlflow=False, enable_mlserver=False, env_manager='local');"
 
 ENV MLFLOW_DISABLE_ENV_CREATION=True
 ENV ENABLE_MLSERVER=False
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('virtualenv')"]
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "from mlflow.models import container as C; C._serve('local')"]
diff --git a/tests/sagemaker/test_cli.py b/tests/sagemaker/test_cli.py
index 95e1404d6233a..134bd7191bd70 100644
--- a/tests/sagemaker/test_cli.py
+++ b/tests/sagemaker/test_cli.py
@@ -6,14 +6,18 @@
 import pytest
 from click.testing import CliRunner
 
+import docker
 import mlflow
 from mlflow.models.docker_utils import build_image_from_context
 from mlflow.sagemaker.cli import build_and_push_container
 
-from tests.pyfunc.test_docker import assert_dockerfiles_equal
+from tests.pyfunc.docker.test_docker import assert_dockerfiles_equal
 
 _MLFLOW_ROOT = Path(mlflow.__file__).parent.parent
 _RESOURCE_DIR = os.path.join(_MLFLOW_ROOT, "tests", "resources", "dockerfile")
+_TEST_IMAGE_NAME = "test-sagemaker-image"
+
+_docker_client = docker.from_env()
 
 
 @pytest.mark.parametrize("env_manager", ["conda", "virtualenv"])
@@ -36,6 +40,8 @@ def _build_image_with_copy(context_dir, image_name):
                 ".",
                 "--env-manager",
                 env_manager,
+                "--container",
+                _TEST_IMAGE_NAME,
             ],
             catch_exceptions=False,
         )
@@ -44,3 +50,6 @@ def _build_image_with_copy(context_dir, image_name):
     actual = dst_dir / "Dockerfile"
     expected = Path(_RESOURCE_DIR) / f"Dockerfile_sagemaker_{env_manager}"
     assert_dockerfiles_equal(actual, expected)
+
+    # Clean up generated image
+    _docker_client.images.remove(_TEST_IMAGE_NAME, force=True)


Code review comments:
path: tests/pyfunc/docker/__init__.py
text: @user1:
do we need this file?

@author:
nope, will remove!

diff_hunk: @@ -30,106 +22,112 @@
     https://github.com/pyenv/pyenv.git /root/.pyenv
 ENV PYENV_ROOT="/root/.pyenv"
 ENV PATH="$PYENV_ROOT/bin:$PATH"
-RUN apt install -y python3.8 python3.8-distutils
-RUN ln -s -f $(which python3.8) /usr/bin/python
-RUN wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py
-RUN python /tmp/get-pip.py
+RUN apt install -y python3.8 python3.8-distutils \
+    && ln -s -f $(which python3.8) /usr/bin/python \
+    && wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py \
+    && python /tmp/get-pip.py
 RUN pip install virtualenv
 """
 
+_DOCKERFILE_TEMPLATE = """# Build an image that can serve mlflow models.
+FROM {base_image}
 
-def _get_maven_proxy():
-    http_proxy = os.getenv("http_proxy")
-    https_proxy = os.getenv("https_proxy")
-    if not http_proxy or not https_proxy:
-        return ""
-
-    # Expects proxies as either PROTOCOL://{USER}:{PASSWORD}@HOSTNAME:PORT
-    # or PROTOCOL://HOSTNAME:PORT
-    parsed_http_proxy = urlparse(http_proxy)
-    assert parsed_http_proxy.hostname is not None, "Invalid `http_proxy` hostname."
-    assert parsed_http_proxy.port is not None, f"Invalid proxy port: {parsed_http_proxy.port}"
-
-    parsed_https_proxy = urlparse(https_proxy)
-    assert parsed_https_proxy.hostname is not None, "Invalid `https_proxy` hostname."
-    assert parsed_https_proxy.port is not None, f"Invalid proxy port: {parsed_https_proxy.port}"
-
-    maven_proxy_options = (
-        "-DproxySet=true",
-        f"-Dhttp.proxyHost={parsed_http_proxy.hostname}",
-        f"-Dhttp.proxyPort={parsed_http_proxy.port}",
-        f"-Dhttps.proxyHost={parsed_https_proxy.hostname}",
-        f"-Dhttps.proxyPort={parsed_https_proxy.port}",
-        "-Dhttps.nonProxyHosts=repo.maven.apache.org",
-    )
+{setup_python_venv}
 
-    if parsed_http_proxy.username is None or parsed_http_proxy.password is None:
-        return " ".join(maven_proxy_options)
+{setup_java}
 
-    return " ".join(
-        (
-            *maven_proxy_options,
-            f"-Dhttp.proxyUser={parsed_http_proxy.username}",
-            f"-Dhttp.proxyPassword={parsed_http_proxy.password}",
-        )
-    )
-
-
-_DOCKERFILE_TEMPLATE = """
-# Build an image that can serve mlflow models.
-FROM ubuntu:20.04
-
-RUN apt-get -y update
-RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y --no-install-recommends \
-         wget \
-         curl \
-         nginx \
-         ca-certificates \
-         bzip2 \
-         build-essential \
-         cmake \
-         openjdk-8-jdk \
-         git-core \
-         maven \
-    && rm -rf /var/lib/apt/lists/*
-
-{setup_python_env}
-
-ENV JAVA_HOME=/usr/lib/jvm/java-8-openjdk-amd64
-ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
-# Set up the program in the image
 WORKDIR /opt/mlflow
 
 {install_mlflow}
 
-{custom_setup_steps}
+{install_model_and_deps}
 
 ENV MLFLOW_DISABLE_ENV_CREATION={disable_env_creation}
 ENV ENABLE_MLSERVER={enable_mlserver}
-
+ENV GUNICORN_CMD_ARGS="--timeout 60 -k gevent"
 
 # granting read/write access and conditional execution authority to all child directories
 # and files to allow for deployment to AWS Sagemaker Serverless Endpoints
 # (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
 RUN chmod o+rwX /opt/mlflow/
 
-{entrypoint}
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "{entrypoint}"]
 """
 
 
-def _get_mlflow_install_step(dockerfile_context_dir, mlflow_home):
+SETUP_MINICONDA = """# Setup miniconda
+RUN curl -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh >> miniconda.sh
+RUN bash ./miniconda.sh -b -p /miniconda && rm ./miniconda.sh
+ENV PATH="/miniconda/bin:$PATH"
+"""
+
+
+def generate_dockerfile(
+    output_dir: str,
+    base_image: str,
+    model_install_steps: Optional[str],
+    entrypoint: str,
+    env_manager: Union[em.CONDA, em.LOCAL, em.VIRTUALENV],
+    mlflow_home: Optional[str] = None,
+    enable_mlserver: bool = False,
+    disable_env_creation_at_runtime: bool = True,
+):
     """
-    Get docker build commands for installing MLflow given a Docker context dir and optional source
-    directory
+    Generates a Dockerfile that can be used to build a docker image, that serves ML model
+    stored and tracked in MLflow.
     """
-    mlflow_home = os.path.abspath(mlflow_home) if mlflow_home else None
+    if base_image.startswith("python:"):
+        setup_python_venv_steps = (
+            "RUN apt-get -y update && apt-get install -y --no-install-recommends nginx"
+        )
+        setup_java_steps = ""
+        install_mlflow_steps = _pip_mlflow_install_step(output_dir, mlflow_home)
+
+    elif base_image == "ubuntu:20.04":
path: mlflow/models/docker_utils.py
text: @user1:
Do we need this to be hard-coded? Can this just check against the constant `UBUNTU_BASE_IMAGE` ?

diff_hunk: @@ -212,3 +212,9 @@ def _create_env_manager_option(help_string, default=None):
     is_flag=True,
     help="If specified, use local environment.",
 )
+
+INSTALL_JAVA = click.option(
+    "--install-java",
+    is_flag=True,
+    help="Install Java in the image. Default is False.",
path: mlflow/utils/cli_args.py
text: @user1:
```suggestion
    help="Install Java in the image. Default is False in order to reduce both the image size and the build time. Model flavors requiring Java will enable this setting automatically.",
```

diff_hunk: @@ -0,0 +1,407 @@
+"""
+This test class is used for comprehensive testing of serving docker images for all MLflow flavors.
+As such, it is not intended to be run on a regular basis and is skipped by default. Rather, it
+should be run manually when making changes to the core docker logic.
+
+To run this test, run the following command manually
+
+    $ pytest tests/pyfunc/test_docker_flavors.py
+
+"""
+
+import json
+import os
+import time
+from operator import itemgetter
+
+import pandas as pd
+import pytest
+import requests
+
+import mlflow
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+
+from tests.catboost.test_catboost_model_export import reg_model
+from tests.diviner.test_diviner_model_export import (
+    diviner_data,
+    diviner_groups,
+    grouped_prophet,
+)
+from tests.fastai.test_fastai_model_export import fastai_model as fastai_model_raw
+from tests.h2o.test_h2o_model_export import h2o_iris_model
+from tests.helper_functions import get_safe_port
+from tests.langchain.test_langchain_model_export import fake_chat_model
+from tests.lightgbm.test_lightgbm_model_export import lgb_model
+from tests.models.test_model import iris_data, sklearn_knn_model
+from tests.paddle.test_paddle_model_export import pd_model
+from tests.pmdarima.test_pmdarima_model_export import (
+    auto_arima_object_model,
+    test_data,
+)
+from tests.prophet.test_prophet_model_export import prophet_model as prophet_raw_model
+from tests.pyfunc.docker.conftest import (
+    MLFLOW_ROOT,
+    TEST_IMAGE_NAME,
+    docker_client,
+    save_model_with_latest_mlflow_version,
+)
+from tests.spacy.test_spacy_model_export import spacy_model_with_data
+from tests.spark.test_spark_model_export import (
+    iris_df,
+    spark,
+    spark_model_iris,
+)
+from tests.statsmodels.model_fixtures import ols_model
+from tests.tensorflow.test_tensorflow2_core_model_export import tf2_toy_model
+from tests.transformers.helper import load_small_seq2seq_pipeline
+
+
+@pytest.mark.skipif(os.getenv("GITHUB_ACTIONS") == "true", reason="Time consuming tests")
+@pytest.mark.parametrize(
+    ("flavor"),
+    [
+        "catboost",
+        "diviner",
+        "fastai",
+        "h2o",
+        # "johnsnowlabs", # Couldn't test JohnSnowLab locally due to license issue
+        "keras",
+        "langchain",
+        "lightgbm",
+        # "mleap", # Mleap model logging is deprecated since 2.6.1
+        "onnx",
+        # "openai", # OPENAI API KEY is not necessarily available for everyone
+        "paddle",
+        "pmdarima",
+        "prophet",
+        "pyfunc",
+        "pytorch",
+        "sklearn",
+        "spacy",
+        "spark",
+        "statsmodels",
+        "tensorflow",
+        "transformers",
+    ],
+)
+def test_build_image_and_serve(flavor, request):
+    model_path = request.getfixturevalue(f"{flavor}_model")
+
+    # Build an image
+    backend = get_flavor_backend(model_uri=model_path, docker_build=True)
+    backend.build_image(
+        model_uri=model_path,
+        image_name=TEST_IMAGE_NAME,
+        mlflow_home=MLFLOW_ROOT,  # Required to prevent installing dev version of MLflow from PyPI
+    )
+
+    # Run a container
+    port = get_safe_port()
+    docker_client.containers.run(
+        image=TEST_IMAGE_NAME,
+        ports={8080: port},
+        detach=True,
+    )
+
+    # Wait until the container to start
path: tests/pyfunc/docker/test_docker_flavors.py
text: @user1:
```suggestion
    # Wait for the container to start
```

diff_hunk: @@ -0,0 +1,129 @@
+import difflib
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+from unittest import mock
+
+import pytest
+import sklearn
+import sklearn.neighbors
+
+import mlflow
+from mlflow.models import Model
+from mlflow.models.docker_utils import build_image_from_context
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.env_manager import VIRTUALENV
+from mlflow.version import VERSION
+
+from tests.pyfunc.docker.conftest import RESOURCE_DIR, get_released_mlflow_version
+
+
+def assert_dockerfiles_equal(actual_dockerfile_path: Path, expected_dockerfile_path: Path):
+    actual_dockerfile = actual_dockerfile_path.read_text().replace(
+        VERSION, get_released_mlflow_version()
+    )
+    expected_dockerfile = (
+        expected_dockerfile_path.read_text()
+        .replace("$MLFLOW_VERSION", get_released_mlflow_version())
path: tests/pyfunc/docker/test_docker.py
text: @user1:
```suggestion
        .replace("${{ MLFLOW_VERSION }}", get_released_mlflow_version())
```

Can we use `${{ ... }}` as a placeholder? `$...` looks like it's an environment variable.

@author:
Makes sense! Will replace them.

diff_hunk: @@ -89,6 +89,7 @@ extend-exclude = [
 "examples/*" = ["T20"]
 "docs/broken_links.py" = ["T20"]
 "mlflow/*" = ["PT018"]
+"tests/pyfunc/docker/test_docker_flavors.py" = ["F401"] # For pytest fixture imports
path: pyproject.toml
text: @user1:
```suggestion
```

We can add `noqa: F401` on lines violating this rule. 

@author:
I did that at first but was a bit tedious to repeat it for many lines....

(on second thoughts, adding this line-by-line as suggested is better as it doesn't disable checks for other lines, will apply the change)

diff_hunk: @@ -0,0 +1,129 @@
+import difflib
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional
+from unittest import mock
+
+import pytest
+import sklearn
+import sklearn.neighbors
+
+import mlflow
+from mlflow.models import Model
+from mlflow.models.docker_utils import build_image_from_context
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.env_manager import VIRTUALENV
+from mlflow.version import VERSION
+
+from tests.pyfunc.docker.conftest import RESOURCE_DIR, get_released_mlflow_version
+
+
+def assert_dockerfiles_equal(actual_dockerfile_path: Path, expected_dockerfile_path: Path):
+    actual_dockerfile = actual_dockerfile_path.read_text().replace(
+        VERSION, get_released_mlflow_version()
+    )
+    expected_dockerfile = (
+        expected_dockerfile_path.read_text()
+        .replace("$MLFLOW_VERSION", get_released_mlflow_version())
+        .replace("$PYTHON_VERSION", PYTHON_VERSION)
+    )
+    assert (
+        actual_dockerfile == expected_dockerfile
+    ), "Generated Dockerfile does not match expected one. Diff:\n" + "\n".join(
+        difflib.unified_diff(expected_dockerfile.splitlines(), actual_dockerfile.splitlines())
+    )
+
+
+def save_model(tmp_path):
+    knn_model = sklearn.neighbors.KNeighborsClassifier()
+    model_path = os.path.join(tmp_path, "model")
+    mlflow.sklearn.save_model(
+        knn_model,
+        path=model_path,
+        pip_requirements=[
+            f"mlflow=={get_released_mlflow_version()}",
+            f"scikit-learn=={sklearn.__version__}",
+        ],  # Skip requirements inference for speed up
+    )
+    return model_path
+
+
+def add_spark_flavor_to_model(model_path):
+    model_config_path = os.path.join(model_path, "MLmodel")
+    model = Model.load(model_config_path)
+    model.add_flavor("spark", spark_version="3.5.0")
+    model.save(model_config_path)
+
+
+@dataclass
+class Param:
+    expected_dockerfile: str
+    env_manager: str = VIRTUALENV
+    mlflow_home: Optional[str] = None
+    install_mlflow: bool = False
+    enable_mlserver: bool = False
+    # If True, image is built with --model-uri param
+    specify_model_uri: bool = True
+
+
+_TEST_PARAMS = [
+    Param(expected_dockerfile="Dockerfile_default"),
+    Param(install_mlflow=True, expected_dockerfile="Dockerfile_install_mlflow"),
+    Param(enable_mlserver=True, expected_dockerfile="Dockerfile_enable_mlserver"),
+    Param(mlflow_home=".", expected_dockerfile="Dockerfile_with_mlflow_home"),
+    Param(specify_model_uri=False, expected_dockerfile="Dockerfile_no_model_uri"),
+]
+
+
+@pytest.mark.parametrize("params", _TEST_PARAMS)
path: tests/pyfunc/docker/test_docker.py
text: @user2:
```suggestion
@user1.mark.parametrize("params", [
    Param(expected_dockerfile="Dockerfile_default"),
    Param(install_mlflow=True, expected_dockerfile="Dockerfile_install_mlflow"),
    Param(enable_mlserver=True, expected_dockerfile="Dockerfile_enable_mlserver"),
    Param(mlflow_home=".", expected_dockerfile="Dockerfile_with_mlflow_home"),
    Param(specify_model_uri=False, expected_dockerfile="Dockerfile_no_model_uri"),
])
```

diff_hunk: @@ -0,0 +1,56 @@
+import os
+from functools import lru_cache
+
+import pytest
+import requests
+from packaging.version import Version
+
+import docker
+import mlflow
+
+TEST_IMAGE_NAME = "test_image"
+MLFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+RESOURCE_DIR = os.path.join(MLFLOW_ROOT, "tests", "resources", "dockerfile")
+
+docker_client = docker.from_env()
+
+
+@pytest.fixture(autouse=True)
+def clean_up_docker_image():
+    try:
+        # Get all containers using the test image
+        containers = docker_client.containers.list(filters={"ancestor": TEST_IMAGE_NAME})
+        for container in containers:
+            container.remove(force=True)
+
+        # Clean up the image
+        docker_client.images.remove(TEST_IMAGE_NAME, force=True)
+    except Exception:
+        pass
+
+
+@lru_cache(maxsize=1)
+def get_released_mlflow_version():
+    url = "https://pypi.org/pypi/mlflow/json"
+    response = requests.get(url)
+    response.raise_for_status()
+    data = response.json()
+    versions = [Version(v) for v in data["releases"].keys() if not v.endswith(("a", "b", "rc"))]
path: tests/pyfunc/docker/conftest.py
text: @user1:
```suggestion
    versions = [v for v in map(Version, data["releases"]) if not (v.is_devreleae or v.is_prerelease)]
```