fix drum

Author: yakov-g

custom_model_runner/datarobot_drum/drum/entry_point.py

@@ -1,19 +1,15 @@
 from datarobot_drum.drum.gunicorn.run_gunicorn import main_gunicorn
 from datarobot_drum.drum.main import main
-import sys
 from datarobot_drum import RuntimeParameters
+from datarobot_drum.drum.utils.setup import setup_options
 
 
 def run_drum_server():
-    cmd_args = sys.argv
-    is_server = False
-    if len(cmd_args) > 1 and cmd_args[1] == "server":
-        is_server = True
-
+    options = setup_options()
     if (
-        RuntimeParameters.has("DRUM_SERVER_TYPE")
+        options.subparser_name == "server"
+        and RuntimeParameters.has("DRUM_SERVER_TYPE")
         and str(RuntimeParameters.get("DRUM_SERVER_TYPE")).lower() == "gunicorn"
-        and is_server
     ):
         main_gunicorn()
     else:

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

@@ -66,7 +66,7 @@
     if temp_loglevel in {"debug", "info", "warning", "error", "critical"}:
         loglevel = temp_loglevel
 
-bind = os.environ.get("ADDRESS", "0.0.0.0:8080")
+bind = os.environ["ADDRESS"]
 # loglevel = "info"
 accesslog = "-"
 access_log_format = '%(h)s %(l)s %(u)s %(t)s "%(r)s" %(s)s %(b)s "%(f)s" "%(a)s"'

custom_model_runner/datarobot_drum/drum/utils/setup.py

@@ -5,6 +5,7 @@
 Released under the terms of DataRobot Tool and Utility Agreement.
 """
 
+import os
 import sys
 
 from datarobot_drum import RuntimeParameters
@@ -48,6 +49,11 @@ def setup_options(args=None):
 
     options = arg_parser.parse_args(args)
 
+    # Override env var from CLI parameter.
+    # Now only ADDRESS is supported.
+    if getattr(options, "address", None):
+        os.environ["ADDRESS"] = options.address
+
     """Set max workers from runtime parameters if available."""
     if RuntimeParameters.has("CUSTOM_MODEL_WORKERS"):
         options.max_workers = RuntimeParameters.get("CUSTOM_MODEL_WORKERS")