[RAPTOR-14353] add gunicorn with gevent to DRUM (#1655)

* added

* working

* lint

* add server selector

* run drum server

* corrected

* import flask

* rollback

* == "gunicorn"

* reformat

* workers = 1 for gunicorn

* add logger

* fix error

* fix error

* fix sys.argv

Author: s-gavrenkov

custom_model_runner/bin/drum

@@ -1,8 +1,6 @@
 #!/usr/bin/env python3
 # PYTHON_ARGCOMPLETE_OK
-
-import os
-from datarobot_drum.drum.main import main
+from datarobot_drum.drum.entry_point import run_drum_server
 
 if __name__ == "__main__":
-    main()
+    run_drum_server()

custom_model_runner/datarobot_drum/drum/entry_point.py

@@ -0,0 +1,17 @@
+from datarobot_drum.drum.gunicorn.run_gunicorn import main_gunicorn
+from datarobot_drum.drum.main import main
+from datarobot_drum import RuntimeParameters
+
+
+def run_drum_server():
+    if (
+        RuntimeParameters.has("DRUM_SERVER_TYPE")
+        and str(RuntimeParameters.get("DRUM_SERVER_TYPE")).lower() == "gunicorn"
+    ):
+        main_gunicorn()
+    else:
+        main()
+
+
+if __name__ == "__main__":
+    run_drum_server()

custom_model_runner/datarobot_drum/drum/gunicorn/__init__.py

@@ -0,0 +1,14 @@
+from datarobot_drum.drum.server import create_flask_app
+
+app = create_flask_app()
+
+_worker_ctx = None
+
+
+def set_worker_ctx(ctx):
+    global _worker_ctx
+    _worker_ctx = ctx
+
+
+def get_worker_ctx():
+    return _worker_ctx

custom_model_runner/datarobot_drum/drum/gunicorn/app.py

@@ -0,0 +1,182 @@
+import logging
+import threading
+from typing import Callable, Any, List, Tuple, Optional
+
+from datarobot_drum.drum.enum import LOGGER_NAME_PREFIX
+
+logger = logging.getLogger(LOGGER_NAME_PREFIX + "." + __name__)
+
+
+class WorkerCtx:
+    """
+    Context of a single gunicorn worker:
+    - .start() — start background tasks
+    - .stop()  — stop background tasks (graceful)
+    - .cleanup() — close resources (DB/clients) in the correct order
+    - add_* methods for registering objects/callbacks for stopping/cleaning up
+    """
+
+    def __init__(self, app):
+        self.app = app
+        self._running = False
+        self._threads: List[threading.Thread] = []
+        self._greenlets: List[Any] = []
+        self._on_stop: List[Tuple[int, Callable[[], None], str]] = []
+        self._on_cleanup: List[Tuple[int, Callable[[], None], str]] = []
+
+    def add_thread(
+        self, t: threading.Thread, *, join_timeout: float = 2.0, name: Optional[str] = None
+    ):
+        """
+        Adds a thread to the worker context and registers it for graceful stopping.
+
+        Args:
+            t (threading.Thread): The thread to be added.
+            join_timeout (float, optional): The timeout for joining the thread during stop. Defaults to 2.0 seconds.
+            name (Optional[str], optional): A descriptive name for the thread. Defaults to None.
+        """
+        t.daemon = True
+        self._threads.append(t)
+
+        def _join():
+            if t.is_alive():
+                t.join(join_timeout)
+
+        self.defer_stop(_join, desc=name or f"thread:{id(t)}")
+
+    def add_greenlet(self, g, *, kill_timeout: float = 2.0, name: Optional[str] = None):
+        """
+        Adds a greenlet to the worker context and registers it for graceful stopping.
+
+        Args:
+            g: The greenlet to be added.
+            kill_timeout (float, optional): The timeout for killing the greenlet during stop. Defaults to 2.0 seconds.
+            name (Optional[str], optional): A descriptive name for the greenlet. Defaults to None.
+        """
+        self._greenlets.append(g)
+
+        def _kill():
+            try:
+                g.kill(block=True, timeout=kill_timeout)
+            except Exception:
+                pass
+
+        self.defer_stop(_kill, desc=name or f"greenlet:{id(g)}")
+
+    def add_closeable(
+        self, obj: Any, method: str = "close", *, order: int = 0, desc: Optional[str] = None
+    ):
+        """
+        Registers an object to be closed using a specified method (e.g., close/quit/disconnect/shutdown) during cleanup (after stop).
+
+        Args:
+            obj (Any): The object to be closed.
+            method (str, optional): The method to call on the object for closing. Defaults to "close".
+            order (int, optional): The priority order for cleanup. Lower values are executed later. Defaults to 0.
+            desc (Optional[str], optional): A description for the cleanup action. Defaults to None.
+        """
+        fn = getattr(obj, method, None)
+        if callable(fn):
+            self.defer_cleanup(fn, order=order, desc=desc or f"{obj}.{method}")
+
+    def defer_stop(self, fn: Callable[[], None], *, order: int = 0, desc: str = "on_stop"):
+        self._on_stop.append((order, fn, desc))
+
+    def defer_cleanup(self, fn: Callable[[], None], *, order: int = 0, desc: str = "on_cleanup"):
+        self._on_cleanup.append((order, fn, desc))
+
+    def start(self):
+        """
+        Starts background tasks for the worker context.
+
+        This method sets the running flag to True and initializes the main application logic.
+        It imports and calls the `main` function from `datarobot_drum.drum.main`, passing
+        the application instance and the worker context as arguments.
+        """
+        self._running = True
+
+        # fix RuntimeError: asyncio.run() cannot be called from a running event loop
+        import asyncio
+        import opentelemetry.instrumentation.openai.utils as otel_utils
+
+        def fixed_run_async(method):
+            try:
+                loop = asyncio.get_running_loop()
+            except RuntimeError:
+                loop = None
+
+            if loop and loop.is_running():
+
+                def handle_task_result(task):
+                    try:
+                        task.result()  # retrieve exception if any
+                    except Exception as e:
+                        # Log or handle exceptions here if needed
+                        logger.error("OpenTelemetry async task error")
+
+                # Schedule the coroutine safely and add done callback to handle errors
+                task = loop.create_task(method)
+                task.add_done_callback(handle_task_result)
+            else:
+                asyncio.run(method)
+
+        # Apply monkey patch
+        otel_utils.run_async = fixed_run_async
+
+        from datarobot_drum.drum.main import main
+
+        main(self.app, self)
+
+    def stop(self):
+        """
+        Gracefully stops the worker context.
+
+        This method sets the running flag to False and executes all registered
+        stop callbacks in the order of their priority. Threads and greenlets
+        are then joined or killed as part of the stopping process.
+        """
+        self._running = False
+        for _, fn, desc in sorted(self._on_stop, key=lambda x: x[0]):
+            try:
+                fn()
+            except Exception:
+                pass
+
+    def cleanup(self):
+        """
+        Cleans up resources in the worker context.
+
+        This method is called at the end of the worker's lifecycle to release resources.
+        It executes all registered cleanup callbacks in reverse priority order, ensuring
+        that resources are closed in the correct sequence.
+
+        Exceptions during cleanup are caught and ignored to prevent interruption.
+        """
+        for _, fn, desc in sorted(self._on_cleanup, key=lambda x: x[0], reverse=True):
+            try:
+                logger.info("WorkerCtx cleanup: %s", desc)
+                fn()
+            except Exception as e:
+                logger.error("Tracing shutdown failed: %s", e)
+
+    def running(self) -> bool:
+        return self._running
+
+
+def create_ctx(app):
+    """
+    Factory method to create a WorkerCtx instance.
+
+    This method initializes the worker context for the application and allows
+    registration of objects or callbacks for stopping and cleanup. It ensures
+    that no long-running processes are started here; the actual startup occurs
+    in `post_worker_init` via `ctx.start()`.
+
+    Args:
+        app: The application instance.
+
+    Returns:
+        WorkerCtx: The initialized worker context.
+    """
+    ctx = WorkerCtx(app)
+    return ctx

custom_model_runner/datarobot_drum/drum/gunicorn/context.py

@@ -0,0 +1,120 @@
+# Import DRUM's WSGI application
+import os
+from datarobot_drum import RuntimeParameters
+
+workers = 1
+if RuntimeParameters.has("CUSTOM_MODEL_WORKERS"):
+    temp_workers = int(RuntimeParameters.get("CUSTOM_MODEL_WORKERS"))
+    if 0 < temp_workers < 200:
+        workers = temp_workers
+elif os.environ.get("MAX_WORKERS"):
+    temp_workers = int(os.environ.get("MAX_WORKERS"))
+    if 0 < temp_workers < 200:
+        workers = temp_workers
+else:
+    pass
+
+backlog = 2048
+if RuntimeParameters.has("DRUM_WEBSERVER_BACKLOG"):
+    temp_backlog = int(RuntimeParameters.get("DRUM_WEBSERVER_BACKLOG"))
+    if 1 <= temp_backlog <= 10000:
+        backlog = temp_backlog
+
+timeout = 120
+if RuntimeParameters.has("DRUM_CLIENT_REQUEST_TIMEOUT"):
+    temp_timeout = int(RuntimeParameters.get("DRUM_CLIENT_REQUEST_TIMEOUT"))
+    if 0 <= temp_timeout <= 3600:
+        timeout = temp_timeout
+
+max_requests = 0
+if RuntimeParameters.has("DRUM_GUNICORN_MAX_REQUESTS"):
+    temp_max_requests = int(RuntimeParameters.get("DRUM_GUNICORN_MAX_REQUESTS"))
+    if 0 <= temp_max_requests <= 10000:
+        max_requests = temp_max_requests
+
+max_requests_jitter = 0
+if RuntimeParameters.has("DRUM_GUNICORN_MAX_REQUESTS_JITTER"):
+    temp_max_requests_jitter = int(RuntimeParameters.get("DRUM_GUNICORN_MAX_REQUESTS_JITTER"))
+    if 1 <= temp_max_requests_jitter <= 10000:
+        max_requests_jitter = temp_max_requests_jitter
+
+worker_connections = 5
+if RuntimeParameters.has("DRUM_WORKER_CONNECTIONS"):
+    temp_worker_connections = int(RuntimeParameters.get("DRUM_WORKER_CONNECTIONS"))
+    if 1 <= temp_worker_connections <= 10000:
+        worker_connections = temp_worker_connections
+
+worker_class = "gevent"
+if RuntimeParameters.has("DRUM_GUNICORN_WORKER_CLASS"):
+    temp_worker_class = str(RuntimeParameters.get("DRUM_GUNICORN_WORKER_CLASS")).lower()
+    if temp_worker_class in {"sync", "gevent"}:
+        worker_class = temp_worker_class
+
+if RuntimeParameters.has("DRUM_GUNICORN_GRACEFUL_TIMEOUT"):
+    temp_graceful_timeout = int(RuntimeParameters.get("DRUM_GUNICORN_GRACEFUL_TIMEOUT"))
+    if 1 <= temp_graceful_timeout <= 3600:
+        graceful_timeout = temp_graceful_timeout
+
+if RuntimeParameters.has("DRUM_GUNICORN_KEEP_ALIVE"):
+    temp_keepalive = int(RuntimeParameters.get("DRUM_GUNICORN_KEEP_ALIVE"))
+    if 1 <= temp_keepalive <= 3600:
+        keepalive = temp_keepalive
+
+loglevel = "info"
+if RuntimeParameters.has("DRUM_GUNICORN_LOG_LEVEL"):
+    temp_loglevel = str(RuntimeParameters.get("DRUM_GUNICORN_LOG_LEVEL")).lower()
+    if temp_loglevel in {"debug", "info", "warning", "error", "critical"}:
+        loglevel = temp_loglevel
+
+bind = os.environ.get("ADDRESS", "0.0.0.0:8080")
+# loglevel = "info"
+accesslog = "-"
+access_log_format = '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s "%(f)s" "%(a)s"'
+
+
+def post_worker_init(worker):
+    from app import app, set_worker_ctx
+    from datarobot_drum.drum.gunicorn.context import create_ctx
+    import sys, shlex
+
+    sys.argv = shlex.split(os.environ.get("DRUM_GUNICORN_DRUM_ARGS"))
+
+    os.environ["MAX_WORKERS"] = "1"
+    if RuntimeParameters.has("CUSTOM_MODEL_WORKERS"):
+        os.environ.pop("MLOPS_RUNTIME_PARAM_CUSTOM_MODEL_WORKERS", None)
+
+    ctx = create_ctx(app)
+    set_worker_ctx(ctx)
+    ctx.start()
+
+
+def worker_exit(worker, code):
+    """
+    Handles the shutdown and cleanup operations for a single Gunicorn worker.
+
+    When terminating or restarting a single Gunicorn worker, this method ensures
+    that the shutdown and cleanup operations are applied only to that specific worker.
+    The following steps should be executed per worker:
+      - predictor.terminate()
+      - stats_collector.mark("end")
+      - runtime.cm_runner.terminate()
+      - runtime.trace_provider.shutdown()
+      - runtime.metric_provider.shutdown()
+      - runtime.log_provider.shutdown()
+      - etc.
+
+    This is necessary because the Gunicorn server is started from the command line
+    (outside of DRUM), and worker-specific cleanup must be handled explicitly.
+
+    Args:
+        worker: The Gunicorn worker instance being terminated.
+        code: The exit code for the worker.
+    """
+    from app import get_worker_ctx
+
+    ctx = get_worker_ctx()
+    if ctx:
+        try:
+            ctx.stop()
+        finally:
+            ctx.cleanup()

custom_model_runner/datarobot_drum/drum/gunicorn/gunicorn.conf.py

@@ -0,0 +1,51 @@
+import logging
+import subprocess
+from pathlib import Path
+import sys
+import os
+import shlex
+
+from datarobot_drum.drum.enum import LOGGER_NAME_PREFIX
+
+logger = logging.getLogger(LOGGER_NAME_PREFIX + "." + __name__)
+
+
+def main_gunicorn():
+    # Resolve directory containing this script so we can always find config and app.py
+    base_dir = Path(__file__).resolve().parent
+    config_path = base_dir / "gunicorn.conf.py"
+
+    if not config_path.is_file():
+        raise FileNotFoundError(f"Gunicorn config not found: {config_path}")
+
+    # Export all provided CLI args (excluding script) into DRUM_GUNICORN_DRUM_ARGS
+    extra_args = sys.argv
+    if extra_args:
+        try:
+            os.environ["DRUM_GUNICORN_DRUM_ARGS"] = shlex.join(extra_args)
+        except AttributeError:
+            os.environ["DRUM_GUNICORN_DRUM_ARGS"] = " ".join(shlex.quote(a) for a in extra_args)
+
+    # Use the gunicorn module explicitly to avoid issues where a shadowed
+    # console script named "gunicorn" actually invokes the DRUM CLI.
+    gunicorn_command = [
+        sys.executable,
+        "-m",
+        "gunicorn",
+        "-c",
+        str(config_path),
+        "app:app",  # module:variable; app.py sits next to this script
+    ]
+
+    try:
+        subprocess.run(gunicorn_command, cwd=base_dir, check=True)
+    except FileNotFoundError:
+        logger.error("gunicorn module not found. Ensure it is installed.")
+        raise
+    except subprocess.CalledProcessError as e:
+        logger.error("Gunicorn exited with non-zero status %s", e.returncode)
+        raise
+
+
+if __name__ == "__main__":
+    main_gunicorn()