Repository: mlflow/mlflow

Issue:
[FR] Allow override of gunicorn app entrypoint for MLflow server
### Willingness to contribute

Yes. I can contribute this feature independently.

### Proposal Summary

The MLflow server executes gunicorn with a hard-coded value `mlflow.server:app` ([reference](https://github.com/mlflow/mlflow/blob/master/mlflow/server/__init__.py#L106)) gunicorn application cli parameter. This prevents easy customization of the MLflow server Flask app i.e. adding `bofore_request` hooks.

### Motivation

> #### What is the use case for this feature?
The goal is adding a `before_request` hook to perform authorization by checking JWT authenication header provided by an AWS application load balancer 

> #### Why is this use case valuable to support for MLflow users in general?
In order to allow extra flexibility for self-hosted MLflow users to customize the mlflow Flask server with pre/post request hooks 

> #### Why is this use case valuable to support for your project(s) or organization?
This would simplify the logic for authorization of users of the MLflow server

> #### Why is it currently difficult to achieve this use case?
Because MLflow server doesn't natively allow custom pre/post process  request hooks for the tracking  webserver component


### Details

It would be desirable that the aforementioned gunicorn app entrypoint can be overridden by a custom value (either via --gunicorn-opts or an additional `mlflow server` cli option)

Ideally one would be able to write a custom script as this:

```python
# customapp.py

from flask import request
from mlflow.server import app

@app.before_request
def hook():
    print(request.path)
```
and run:
`mlflow server ... --gunicorn-app customapp:app`


### What component(s) does this bug affect?

- [ ] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [ ] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [ ] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [X] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
Hi @lucascott 
1. What's your scenario for adding bofore_request hooks on the server request handler ? 

2. For the demand of adding hook, I think we should define a hook plugin interface instead of allow custom gunicorn entry app. e.g. we can define the plugin interface like:

```
class MLflowServingHookBase {
    def before_request(...)
    def post_request(...)
}
```

And we can allow MLFlow support an HOOK environment variable that specify the plugin class path, if the HOOK env variable defined, load the plugin.
Hi @WeichenXu123 

(to answer 1.) As mentioned above,  The goal is adding a `before_request` hook to perform authorization.
In details, the mlflow server deployment consists of a AWS application load balancer performing authentication and forwarding the authenticated user to the mlflow server running as a Fargate service. To keep the architecture simple, it would be nice not having to add middleware (between the ALB and the MLflow server container) to perform user authorization to the server. Therefore it would be convenient to add a custom hook checking the JWT authenication header provided by the ALB.

(to answer 2.) Your proposal seems sound however I wouldn't like to be blocked by/coupled to a plugin interface MLflow defines on top of Flask's one. One reason would be to take advantage of shared Flask middleware components (for auth, logging, etc.) without having to instantiate new MLflow server plugin classes. 

My proposal was somewhat more generic as it would have allowed one to have full access and customization capabilities to the `mlflow.server.app` instance before spinning up the gunicorn WSGI. This approach of "customizing" the Flask app before spinning up the server is already taken into use in the MLflow's Prometheus exporter ([ref](https://github.com/mlflow/mlflow/blob/f6c4eb4736542447ca2e0aa201f6f85e7e61bd60/mlflow/server/prometheus_exporter.py)).
We have the same use case. It would be nice to configure the flask app to eliminate middleware.

For the specific use case of implementing auth, I wonder if it's possible to implement in a [gunicorn pre_request](https://docs.gunicorn.org/en/latest/settings.html#pre-request) hook? And then `--gunicorn-opts='--config myconfig.py` or `GUNICORN_CMD_ARGS='--config myconfig.py'.

I have not used gunicorn's pre_request hooks, so might not be what I am thinking it is. I'll give this a spin.
> For the specific use case of implementing auth, I wonder if it's possible to implement in a [gunicorn pre_request](https://docs.gunicorn.org/en/latest/settings.html#pre-request) hook? And then `--gunicorn-opts='--config myconfig.py` or `GUNICORN_CMD_ARGS='--config myconfig.py'.

Hooking at the gunicorn level does not work. You can't return a response, redirect etc. Even raising doesn't really do anything. Guess we can rule that one out.

You can have quite some configuration of the Flask app via a hacky workaround. 
Import the `mlflow.server.app` instance in your gunicorn `myconfig.py` configuration script and customize it by adding your desired middleware:

```python
import logging

from flask import request
from mlflow.server import app


@app.before_request
def before_req_hook():
    # your custom logic here
    logging.info(request)
```
This worked for me but it's clearly not an elegant solution driven by the lack of support of such feature. I'd really appreciate to move away from this approach.


<!-- assign-maintainer -->
@BenWilson2 @dbczumar @harupy @WeichenXu123 Please assign a maintainer and start triaging this issue.
I'd be happy to work on this if there's interest. Currently manually running gunicorn to get the desired behavior. But that could break down if something were to change with mlflow.
@jmahlik Please feel free to submit a draft PR where we can discuss your plans for implementation (And see a mockup of your proposed design). Looking forward to it!
Thoughts on implementing via a plugin entrypoint? Thought I'd throw together an example before getting too far down the road.

* Seems simple
* Allows sharing of plugins between users
* Falls in line with the other plugin styles
* Keeps it clean by requiring a plugin interface with an entrypoint.
* Allows complete control over the config like @lucascott was mentioning a need for. There could be multiple things a user needs to do. Adding them via another interface seems overly complex.

Some some pseudo code. (possibly almost all of the code)

```python
# Default app name is mlflow.server.app
app_name = "{}:app".format(__name__)

# A user can have one app plugin installed
apps = [entrypoint.name for entrypoint in entrypoints.get_group_all("mlflow.app")]

if len(apps) > 1:
    raise MlflowException("There can only be one app, multiple detected. {}".format(apps))

if len(apps) == 1:
    app_name = apps[0]

# Change the commands to reference app_name

def _build_waitress_command(waitress_opts, host, port):
    opts = shlex.split(waitress_opts) if waitress_opts else []
    return (
        ["waitress-serve"]
        + opts
        + ["--host=%s" % host, "--port=%s" % port, "--ident=mlflow", app_name]
    )


def _build_gunicorn_command(gunicorn_opts, host, port, workers):
    bind_address = f"{host}:{port}"
    opts = shlex.split(gunicorn_opts) if gunicorn_opts else []
    return ["gunicorn"] + opts + ["-b", bind_address, "-w", "%s" % workers, app_name]
```

And an example plugin would look something like this

entrypoint: `mlflow.app = my_plugin:app`

```python
import logging

# This would be all that plugin author is required to import
from mlflow.server import app

# Could do custom logging on either the app or logging itself
# but you'll possibly have to clear the handlers or there will be duplicate output
logging.basicConfig(level=logging.INFO)
root_logger = logging.getLogger("root")
mlflow_logger = logging.getLogger("mlflow")

for handler in root_logger.handlers:
    root_logger.removeHandler(handler)

for handler in mlflow_logger.handlers:
    root_logger.removeHandler(handler)

# Change the format or implement json logging etc
f = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
st = logging.StreamHandler()
st.setFormatter(f)
root_logger.addHandler(st)


# Now configure the app to meet your needs
@app.before_request
def before_req_hook():
    # Implement jwt auth here
    return "This works!"
```
The plugin interface approach looks great! I think it will keep the mllfow implementation simple and clean as well as simplifying the UX for the submitted app configuration as well. Looking forward to your PR! :D 

Resolving patch:
diff --git a/mlflow/server/__init__.py b/mlflow/server/__init__.py
--- a/mlflow/server/__init__.py
+++ b/mlflow/server/__init__.py
@@ -1,3 +1,4 @@
+import entrypoints
 import os
 import shlex
 import sys
@@ -5,6 +6,7 @@
 
 from flask import Flask, send_from_directory, Response
 
+from mlflow.exceptions import MlflowException
 from mlflow.server import handlers
 from mlflow.server.handlers import (
     get_artifact_handler,
@@ -105,19 +107,37 @@ def serve():
     return Response(text, mimetype="text/plain")
 
 
-def _build_waitress_command(waitress_opts, host, port):
+def _get_app_name() -> str:
+    """Search for plugins for custom mlflow app, otherwise return default."""
+    apps = list(entrypoints.get_group_all("mlflow.app"))
+    # Default, nothing installed
+    if len(apps) == 0:
+        return f"{__name__}:app"
+    # Cannot install more than one
+    if len(apps) > 1:
+        raise MlflowException(
+            "Multiple server plugins detected. "
+            "Only one server plugin may be installed. "
+            f"Detected plugins: {', '.join([f'{a.module_name}.{a.object_name}' for a in apps])}"
+        )
+    # Has a plugin installed
+    plugin_app = apps[0]
+    return f"{plugin_app.module_name}:{plugin_app.object_name}"
+
+
+def _build_waitress_command(waitress_opts, host, port, app_name):
     opts = shlex.split(waitress_opts) if waitress_opts else []
     return (
         ["waitress-serve"]
         + opts
-        + ["--host=%s" % host, "--port=%s" % port, "--ident=mlflow", "mlflow.server:app"]
+        + ["--host=%s" % host, "--port=%s" % port, "--ident=mlflow", app_name]
     )
 
 
-def _build_gunicorn_command(gunicorn_opts, host, port, workers):
+def _build_gunicorn_command(gunicorn_opts, host, port, workers, app_name):
     bind_address = f"{host}:{port}"
     opts = shlex.split(gunicorn_opts) if gunicorn_opts else []
-    return ["gunicorn"] + opts + ["-b", bind_address, "-w", "%s" % workers, "mlflow.server:app"]
+    return ["gunicorn"] + opts + ["-b", bind_address, "-w", "%s" % workers, app_name]
 
 
 def _run_server(
@@ -160,9 +180,10 @@ def _run_server(
     if expose_prometheus:
         env_map[PROMETHEUS_EXPORTER_ENV_VAR] = expose_prometheus
 
+    app_name = _get_app_name()
     # TODO: eventually may want waitress on non-win32
     if sys.platform == "win32":
-        full_command = _build_waitress_command(waitress_opts, host, port)
+        full_command = _build_waitress_command(waitress_opts, host, port, app_name)
     else:
-        full_command = _build_gunicorn_command(gunicorn_opts, host, port, workers or 4)
+        full_command = _build_gunicorn_command(gunicorn_opts, host, port, workers or 4, app_name)
     _exec_cmd(full_command, extra_env=env_map, capture_output=False)


Test patch:
diff --git a/tests/resources/mlflow-test-plugin/mlflow_test_plugin/app.py b/tests/resources/mlflow-test-plugin/mlflow_test_plugin/app.py
new file mode 100644
--- /dev/null
+++ b/tests/resources/mlflow-test-plugin/mlflow_test_plugin/app.py
@@ -0,0 +1,28 @@
+import logging
+
+# This would be all that plugin author is required to import
+from mlflow.server import app as custom_app
+
+# Can do custom logging on either the app or logging itself
+# but you'll possibly have to clear the existing handlers or there will be duplicate output
+# See https://docs.python.org/3/howto/logging-cookbook.html
+
+app_logger = logging.getLogger(__name__)
+
+# Configure the app
+custom_app.config["MY_VAR"] = "config-var"
+
+
+def is_logged_in():
+    return True
+
+
+@custom_app.before_request
+def before_req_hook():
+    """A custom before request handler.
+
+    Can implement things such as authentication, special handling, etc.
+    """
+    app_logger.warning("Hello from before request!")
+    if not is_logged_in():
+        return "Unauthorized", 403
diff --git a/tests/resources/mlflow-test-plugin/setup.py b/tests/resources/mlflow-test-plugin/setup.py
--- a/tests/resources/mlflow-test-plugin/setup.py
+++ b/tests/resources/mlflow-test-plugin/setup.py
@@ -31,5 +31,7 @@
         "mlflow.deployments": "faketarget=mlflow_test_plugin.fake_deployment_plugin",
         # Define a Mlflow model evaluator with name "dummy_evaluator"
         "mlflow.model_evaluator": "dummy_evaluator=mlflow_test_plugin.dummy_evaluator:DummyEvaluator",  # pylint: disable=line-too-long
+        # Define a custom Mlflow application with name custom_app
+        "mlflow.app": "app=mlflow_test_plugin.app:custom_app",
     },
 )
diff --git a/tests/server/test_init.py b/tests/server/test_init.py
--- a/tests/server/test_init.py
+++ b/tests/server/test_init.py
@@ -0,0 +1,65 @@
+from unittest import mock
+
+import entrypoints
+import pytest
+
+from mlflow.exceptions import MlflowException
+from mlflow import server
+
+
+@pytest.fixture()
+def mock_exec_cmd():
+    with mock.patch("mlflow.server._exec_cmd") as m:
+        yield m
+
+
+def test_get_app_name():
+    # Hard to test if the test plugin is installed.
+    with mock.patch("mlflow.server.entrypoints.get_group_all", return_value=[]):
+        assert server._get_app_name() == f"{server.__name__}:app"
+
+
+def test_get_app_name_two_plugins():
+    """Two server apps cannot exist, which one would be served?"""
+    plugins = [
+        entrypoints.EntryPoint("one", "one.app", "app"),
+        entrypoints.EntryPoint("two", "two.app", "two"),
+    ]
+    with mock.patch("mlflow.server.entrypoints.get_group_all", return_value=plugins):
+        with pytest.raises(MlflowException, match=r"one.*two"):
+            server._get_app_name()
+
+
+def test_get_app_name_custom_app_plugin():
+    """This test requires the package in tests/resources/mlflow-test-plugin to be installed"""
+    assert server._get_app_name() == "mlflow_test_plugin.app:custom_app"
+
+
+def test_build_waitress_command():
+    assert server._build_waitress_command("", "localhost", "5000", f"{server.__name__}:app") == [
+        "waitress-serve",
+        "--host=localhost",
+        "--port=5000",
+        "--ident=mlflow",
+        "mlflow.server:app",
+    ]
+
+
+def test_build_gunicorn_command():
+    assert server._build_gunicorn_command(
+        "", "localhost", "5000", "4", f"{server.__name__}:app"
+    ) == ["gunicorn", "-b", "localhost:5000", "-w", "4", "mlflow.server:app"]
+
+
+def test_run_server(mock_exec_cmd):
+    """Make sure this runs."""
+    with mock.patch("sys.platform", return_value="linux"):
+        server._run_server("", "", "", "", "", "", "", "")
+    mock_exec_cmd.assert_called_once()
+
+
+def test_run_server_win32(mock_exec_cmd):
+    """Make sure this runs."""
+    with mock.patch("sys.platform", return_value="win32"):
+        server._run_server("", "", "", "", "", "", "", "")
+    mock_exec_cmd.assert_called_once()