[AGENT-3621] Adds support for extending DRUM server mode (#698)

* Add new hook file for flask extension

If this file is present in the custom model dir then we will use it to
extend our Flask application

* Add new example for extending Flask server

I took the sklearn example and just added the flask customizer code to it.
I felt like sklearn is pretty simple so that should not confuse from
main part of the example which is the `custom_flask.py` code. I'm using
3rd party Flask extension to help with some of the authentication code to
show that this solution is highly customizable.

* Bump version 1.9.11

* Fix lint

* Add new line

* Change print to logging

* Improve README formatting

* Make it a dev release

* Small renaming tweaks to flask hook code

* Improve comments for sample custom flask hook code

* Add test case for custom flask hook

* Fix changelog

* Fix black

* Bump to 1.9.12

* Fix comment

* Improve readme for new sample

* Rename and update docs

I don't want people confusing the flask extension with a real model template.

* Don't use assert for user driven logic

Author: elatt

custom_model_runner/CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+#### [1.9.12] - 2022-10-25
+##### Added
+- Add support for a new hook (`custom_flask.py`) in the model-dir to allow extending the Flask
+  application when drum is running in server mode.
+- Add a new model template sample (`flask_extension_httpauth`) to illustrate a potential
+  authentication use-case using the new `custom_flask.py` hook.
+
 #### [1.9.11] - 2022-10-24
 ##### Changed
 - Add support to initialize DR Python client in order to allow DR API access.
@@ -18,7 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Bump com.datarobot.datarobot-prediction package to 2.2.1
 ##### Fixed
 - Pin `datarobot==2.27.0`
-- Handle missing values in image typeschema validator 
+- Handle missing values in image typeschema validator
 
 #### [1.9.8] - 2022-08-04
 ##### Changed
@@ -62,7 +69,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 #### [1.8.0] - 2022-02-18
 ##### Added
-- Built-in support for ONNX models 
+- Built-in support for ONNX models
 - Support for new custom (training) task templates
 
 #### [1.7.2dev1] - 2022-02-16
@@ -164,7 +171,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 #### [1.5.11] - 2021-08-19
 ##### Fixed
 - Apply default schema to transforms
-##### Added 
+##### Added
 - type schema to pipeline examples
 
 #### [1.5.10] - 2021-07-30
@@ -306,7 +313,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 #### [1.4.5] - 2020-12-02
 ##### Added
--  **/transform** endpoint added to prediction server 
+-  **/transform** endpoint added to prediction server
 ##### Changes
 - Allow multiclass to function with only 2 labels
 
@@ -375,12 +382,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 #### [1.1.4] - 2020-08-04
 ##### Added
 - the docker flag now takes a directory, and will build a docker image
-- the `push` verb lets you add your code into DataRobot. 
+- the `push` verb lets you add your code into DataRobot.
 - H2O models support
 - r_lang fit component, pipeline, and template
 ##### Changed
 - search custom.py recursively in the code dir
-- set rpy2 dependcy <= 3.2.7 to avoid pandas import error 
+- set rpy2 dependcy <= 3.2.7 to avoid pandas import error
 
 ## [1.1.3] - 2020-07-17
 ### Added
@@ -427,15 +434,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - change command `cmrun predict` to `cmrun score`
 
 ### Added
-- `new` subcommand for model templates creation 
+- `new` subcommand for model templates creation
 
 ## [1.0.16] - 2020-05-05
 ### Changed
 - unpin rpy2 dependency version
 
 ## [1.0.15] - 2020-05-04
 ### Changed
-- require to use sub-command, e.g. `cmrun predict`  
+- require to use sub-command, e.g. `cmrun predict`
 
 ## [1.0.14] - 2020-04-30
 ### Changed

custom_model_runner/datarobot_drum/drum/description.py

@@ -4,6 +4,6 @@
 This is proprietary source code of DataRobot, Inc. and its affiliates.
 Released under the terms of DataRobot Tool and Utility Agreement.
 """
-version = "1.9.11"
+version = "1.9.12"
 __version__ = version
 project_name = "datarobot-drum"

custom_model_runner/datarobot_drum/drum/enum.py

@@ -21,6 +21,9 @@
 CUSTOM_FILE_NAME = "custom"
 
 
+FLASK_EXT_FILE_NAME = "custom_flask"
+
+
 POSITIVE_CLASS_LABEL_ARG_KEYWORD = "positive_class_label"
 
 

custom_model_runner/datarobot_drum/drum/model_adapter.py

@@ -138,7 +138,8 @@ def _load_custom_hooks_for_legacy_drum(self, custom_module):
 
     def load_custom_hooks(self):
         custom_file_paths = list(Path(self._model_dir).rglob("{}.py".format(CUSTOM_FILE_NAME)))
-        assert len(custom_file_paths) <= 1
+        if len(custom_file_paths) > 1:
+            raise RuntimeError("Found too many custom hook files: {}".format(custom_file_paths))
 
         if len(custom_file_paths) == 0:
             print("No {}.py file detected in {}".format(CUSTOM_FILE_NAME, self._model_dir))

custom_model_runner/datarobot_drum/resource/components/Python/prediction_server/prediction_server.py

@@ -6,6 +6,7 @@
 """
 import logging
 import sys
+from pathlib import Path
 from mlpiper.components.connectable_component import ConnectableComponent
 
 from datarobot_drum.drum.common import (
@@ -18,6 +19,7 @@
     ModelInfoKeys,
     RunLanguage,
     TargetType,
+    FLASK_EXT_FILE_NAME,
 )
 from datarobot_drum.drum.description import version as drum_version
 from datarobot_drum.drum.exceptions import DrumCommonException
@@ -192,6 +194,7 @@ def handle_exception(e):
         cli.show_server_banner = lambda *x: None
 
         app = get_flask_app(model_api)
+        self.load_flask_extensions(app)
 
         host = self._params.get("host", None)
         port = self._params.get("port", None)
@@ -209,3 +212,25 @@ def terminate(self):
         terminate_op = getattr(self._predictor, "terminate", None)
         if callable(terminate_op):
             terminate_op()
+
+    def load_flask_extensions(self, app):
+        custom_file_paths = list(Path(self._code_dir).rglob("{}.py".format(FLASK_EXT_FILE_NAME)))
+        if len(custom_file_paths) > 1:
+            raise RuntimeError("Found too many custom hook files: {}".format(custom_file_paths))
+
+        if len(custom_file_paths) == 0:
+            logger.info("No %s.py file detected in %s", FLASK_EXT_FILE_NAME, self._code_dir)
+            return
+
+        custom_file_path = custom_file_paths[0]
+        logger.info("Detected %s .. trying to load Flask extensions", custom_file_path)
+        sys.path.insert(0, str(custom_file_path.parent))
+
+        try:
+            custom_module = __import__(FLASK_EXT_FILE_NAME)
+            custom_module.init_app(app)
+        except ImportError as e:
+            logger.error("Could not load hooks", exc_info=True)
+            raise DrumCommonException(
+                "Failed to extend Flask app from [{}] : {}".format(custom_file_path, e)
+            )

model_templates/flask_extension_httpauth/README.md

@@ -0,0 +1,20 @@
+## Extending Web Server Behavior
+
+This sample is meant to illustrate how you can add thirdparty or your own custom Flask extensions
+when drum is in server-mode. To demonstrate one potential usecase, the `custom_flask.py` file in
+this model directory will extend the HTTP server to require a specific [Bearer Token](https://swagger.io/docs/specification/authentication/bearer-authentication/)
+when making any requests to it.
+
+For completeness, we also include all the model related files from the [Python Sklearn Inference Model Template](../python3_sklearn/).
+
+Note: it is **not** necessary (nor recommended) to add authentication to custom models that are created in DataRobot MLOps.
+This example is simply to demonstration the flexibility of the `custom_flask.py` hook.
+
+## Instructions
+Create a new custom model with these files and use the Python Drop-In Environment with it
+
+### To run locally using 'drum'
+Paths are relative to `./datarobot-user-models`:
+```
+drum server --docker public_dropin_environments/python3_sklearn --code-dir model_templates/python3_sklearn_flask_ext/ --target-type regression --address localhost:8080
+```

model_templates/flask_extension_httpauth/custom.py

@@ -0,0 +1,33 @@
+"""
+Copyright 2021 DataRobot, Inc. and its affiliates.
+All rights reserved.
+This is proprietary source code of DataRobot, Inc. and its affiliates.
+Released under the terms of DataRobot Tool and Utility Agreement.
+"""
+import pandas as pd
+
+
+def transform(data, model):
+    """
+    Note: This hook may not have to be implemented for your model.
+    In this case implemented for the model used in the example.
+
+    Modify this method to add data transformation before scoring calls. For example, this can be
+    used to implement one-hot encoding for models that don't include it on their own.
+
+    Parameters
+    ----------
+    data: pd.DataFrame
+    model: object, the deserialized model
+
+    Returns
+    -------
+    pd.DataFrame
+    """
+    # Execute any steps you need to do before scoring
+    # Remove target columns if they're in the dataset
+    for target_col in ["Grade 2014", "Species"]:
+        if target_col in data:
+            data.pop(target_col)
+    data = data.fillna(0)
+    return data

model_templates/flask_extension_httpauth/custom_flask.py

@@ -0,0 +1,57 @@
+"""
+Copyright 2022 DataRobot, Inc. and its affiliates.
+All rights reserved.
+This is proprietary source code of DataRobot, Inc. and its affiliates.
+Released under the terms of DataRobot Tool and Utility Agreement.
+"""
+import logging
+
+from flask import request
+from flask_httpauth import HTTPTokenAuth
+
+logger = logging.getLogger(__name__)
+
+token_auth = HTTPTokenAuth()
+
+AUTHENTICATION_TOKEN = "DUMMY_TOKEN_123"
+
+
+@token_auth.verify_token
+def verify(token):
+    # Hard-code users for demo purposes but in a real setup this data would be
+    # fetched from a database or secure key vault, for example.
+    if token == AUTHENTICATION_TOKEN:
+        # flask_httpauth requires this function to return a username on successful authentication
+        # so it can be used for authorization (but we aren't implementing that for this sample).
+        return "dummy_user"
+
+
+def init_app(app):
+    """
+    Below is a sample hook that illustrates how to add simple token based
+    authentication to most of the routes served by the custom model runner.
+
+    Parameters
+    ----------
+    app: Flask
+    """
+    # Health check endpoints shouldn't require auth
+    no_auth_endpoints = {"model_api.ping", "model_api.health"}
+    logger.info("Setting up authentication on all routes except ping routes")
+
+    @app.before_request
+    def check_auth():
+        auth = token_auth.get_auth()
+
+        # Flask normally handles OPTIONS requests on its own, but in the case it is configured to
+        # forward those to the application, we need to ignore authentication headers and let the
+        # request through to avoid unwanted interactions with CORS.
+        if request.method != "OPTIONS" and request.endpoint not in no_auth_endpoints:
+            user = token_auth.authenticate(auth, None)
+            if user in (False, None):
+                return token_auth.auth_error_callback(401)
+
+    logger.info(
+        'Please authenticate to the server:\n\t`curl -H "Authorization: Bearer %s" ...`',
+        AUTHENTICATION_TOKEN,
+    )

model_templates/flask_extension_httpauth/model-metadata.yaml

@@ -0,0 +1,10 @@
+name: drumpush-regression
+type: inference
+targetType: regression
+modelID: 5f1f15a4d6111f01cb7f91fd
+environmentID: 5e8c889607389fe0f466c72d
+inferenceModel:
+  targetName: Grade 2014
+validation:
+  # Path is relative to this file
+  input: ../../../tests/testdata/juniors_3_year_stats_regression.csv

model_templates/flask_extension_httpauth/requirements.txt

@@ -0,0 +1 @@
+Flask-HTTPAuth==4.7.0

model_templates/flask_extension_httpauth/sklearn_reg.pkl

@@ -2,4 +2,4 @@
 -r requirements_lint.txt
 -r requirements_test.txt
 
-datarobot-oss-java-jdk11==1.11.0.13.post1+dr
+datarobot-oss-java-jdk11==1.11.0.13.post1+dr

requirements_dev.txt

@@ -0,0 +1,69 @@
+"""
+Copyright 2022 DataRobot, Inc. and its affiliates.
+All rights reserved.
+This is proprietary source code of DataRobot, Inc. and its affiliates.
+Released under the terms of DataRobot Tool and Utility Agreement.
+"""
+from pathlib import Path
+import shutil
+
+import pytest
+import requests
+
+from .constants import (
+    PYTHON_UNSTRUCTURED,
+    UNSTRUCTURED,
+    TESTS_FIXTURES_PATH,
+)
+from datarobot_drum.resource.drum_server_utils import DrumServerRun
+from datarobot_drum.resource.utils import _create_custom_model_dir
+from datarobot_drum.drum.utils.drum_utils import unset_drum_supported_env_vars
+
+
+class TestDrumServerCustomAuth:
+    @pytest.fixture(scope="class")
+    def custom_flask_script(self):
+        return (Path(TESTS_FIXTURES_PATH) / "custom_flask_demo_auth.py", "custom_flask.py")
+
+    @pytest.fixture(scope="class")
+    def custom_model_dir(self, custom_flask_script, resources, tmp_path_factory):
+        tmp_dir = tmp_path_factory.mktemp("model_dir")
+        custom_model_dir = _create_custom_model_dir(
+            resources, tmp_dir, None, UNSTRUCTURED, PYTHON_UNSTRUCTURED,
+        )
+        fixture_filename, target_name = custom_flask_script
+        shutil.copy2(fixture_filename, custom_model_dir / target_name)
+        return custom_model_dir
+
+    @pytest.fixture(scope="class")
+    def drum_server(self, resources, custom_model_dir):
+        unset_drum_supported_env_vars()
+        with DrumServerRun(
+            resources.target_types(UNSTRUCTURED),
+            resources.class_labels(None, UNSTRUCTURED),
+            custom_model_dir,
+        ) as run:
+            yield run
+
+    def test_auth_passthrough(self, drum_server):
+        response = requests.get(drum_server.url_server_address + "/ping/")
+        assert response.ok
+
+    def test_missing_auth_header(self, drum_server):
+        response = requests.get(drum_server.url_server_address + "/info/")
+        assert response.status_code == 401
+        assert response.json()["message"] == "Missing X-Auth header"
+
+    def test_bad_auth_token(self, drum_server):
+        response = requests.get(
+            drum_server.url_server_address + "/info/", headers={"X-Auth": "token"}
+        )
+        assert response.status_code == 401
+        assert response.json()["message"] == "Auth token is invalid"
+
+    def test_successful_auth(self, drum_server):
+        response = requests.get(
+            drum_server.url_server_address + "/info/", headers={"X-Auth": "t0k3n"}
+        )
+        assert response.ok
+        assert response.json()["drumServer"] == "flask"

tests/drum/test_drum_server_custom_flask.py

@@ -0,0 +1,21 @@
+"""
+Copyright 2022 DataRobot, Inc. and its affiliates.
+All rights reserved.
+This is proprietary source code of DataRobot, Inc. and its affiliates.
+Released under the terms of DataRobot Tool and Utility Agreement.
+"""
+from flask import request, jsonify
+
+
+def init_app(app):
+    @app.before_request
+    def check_header():
+        # Allow ping route with no Auth otherwise test setup would fail
+        if request.endpoint != "model_api.ping":
+            try:
+                token = request.headers["X-Auth"]
+            except KeyError:
+                return jsonify({"message": "Missing X-Auth header"}), 401
+            else:
+                if token != "t0k3n":
+                    return jsonify({"message": "Auth token is invalid"}), 401