From 09190e7c968e6b6aac752a73002a374899083515 Mon Sep 17 00:00:00 2001 From: Yakov Goldberg Date: Fri, 12 Sep 2025 14:21:00 -0700 Subject: [PATCH 1/2] fix drum --- .../datarobot_drum/drum/entry_point.py | 14 ++++++-------- .../datarobot_drum/drum/gunicorn/gunicorn.conf.py | 2 +- .../datarobot_drum/drum/utils/setup.py | 6 ++++++ 3 files changed, 13 insertions(+), 9 deletions(-) diff --git a/custom_model_runner/datarobot_drum/drum/entry_point.py b/custom_model_runner/datarobot_drum/drum/entry_point.py index 02c26bd3a..874f6b059 100644 --- a/custom_model_runner/datarobot_drum/drum/entry_point.py +++ b/custom_model_runner/datarobot_drum/drum/entry_point.py @@ -1,19 +1,17 @@ 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 +from datarobot_drum.drum.enum import ArgumentsOptions -def run_drum_server(): - cmd_args = sys.argv - is_server = False - if len(cmd_args) > 1 and cmd_args[1] == "server": - is_server = True +def run_drum_server(): + options = setup_options() if ( - RuntimeParameters.has("DRUM_SERVER_TYPE") + options.subparser_name == ArgumentsOptions.SERVER + and RuntimeParameters.has("DRUM_SERVER_TYPE") and str(RuntimeParameters.get("DRUM_SERVER_TYPE")).lower() == "gunicorn" - and is_server ): main_gunicorn() else: diff --git a/custom_model_runner/datarobot_drum/drum/gunicorn/gunicorn.conf.py b/custom_model_runner/datarobot_drum/drum/gunicorn/gunicorn.conf.py index e4d186f8e..ba3076989 100644 --- a/custom_model_runner/datarobot_drum/drum/gunicorn/gunicorn.conf.py +++ b/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"' diff --git a/custom_model_runner/datarobot_drum/drum/utils/setup.py b/custom_model_runner/datarobot_drum/drum/utils/setup.py index 6bd80ba88..97c2c4fa8 100644 --- a/custom_model_runner/datarobot_drum/drum/utils/setup.py +++ b/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") From b95ce2206677139454e53088a1ac496716305941 Mon Sep 17 00:00:00 2001 From: Yakov Goldberg Date: Wed, 17 Dec 2025 11:58:32 -0800 Subject: [PATCH 2/2] add tests --- .../datarobot_drum/drum/utils/setup.py | 1 - tests/functional/test_runtime_parameters.py | 20 +- .../datarobot_drum/drum/utils/test_setup.py | 268 ++++++++++-------- 3 files changed, 158 insertions(+), 131 deletions(-) diff --git a/custom_model_runner/datarobot_drum/drum/utils/setup.py b/custom_model_runner/datarobot_drum/drum/utils/setup.py index 97c2c4fa8..012d187b1 100644 --- a/custom_model_runner/datarobot_drum/drum/utils/setup.py +++ b/custom_model_runner/datarobot_drum/drum/utils/setup.py @@ -32,7 +32,6 @@ def setup_options(args=None): Parsed command line options as an argparse.Namespace object. """ arg_parser = CMRunnerArgsRegistry.get_arg_parser() - try: import argcomplete except ImportError: diff --git a/tests/functional/test_runtime_parameters.py b/tests/functional/test_runtime_parameters.py index 8d0bcdbbc..c4d95814c 100644 --- a/tests/functional/test_runtime_parameters.py +++ b/tests/functional/test_runtime_parameters.py @@ -222,8 +222,8 @@ def _test_custom_model_with_runtime_params( if docker: cmd += " --docker {} --verbose ".format(docker) - _, stdout, _ = _exec_shell_cmd(cmd, err_msg=None, assert_if_fail=False, env=env) - return stdout + _, stdout, stderr = _exec_shell_cmd(cmd, err_msg=None, assert_if_fail=False, env=env) + return stdout, stderr @classmethod def _setup_runtime_parameters( @@ -251,20 +251,20 @@ def _setup_runtime_parameters( def test_runtime_parameters_invalid_yaml( self, resources, tmp_path, runtime_param_values_stream, use_runtime_params_env_var ): - stdout = self._test_custom_model_with_runtime_params( + stdout, stderr = self._test_custom_model_with_runtime_params( resources, tmp_path, runtime_param_values_stream, is_invalid_yaml=True, use_runtime_params_env_var=use_runtime_params_env_var, ) - assert "Invalid runtime parameter values YAML content!" in stdout + assert "Invalid runtime parameter values YAML content!" in stdout + stderr @pytest.mark.parametrize("use_runtime_params_env_var", [True, False]) def test_runtime_parameters_missing_attr( self, resources, tmp_path, runtime_param_values_stream, use_runtime_params_env_var ): - stdout = self._test_custom_model_with_runtime_params( + stdout, stderr = self._test_custom_model_with_runtime_params( resources, tmp_path, runtime_param_values_stream, @@ -273,27 +273,27 @@ def test_runtime_parameters_missing_attr( ) assert re.search( r".*Failed to load runtime parameter.*{\\'credentialType\\': DataError\(\\'is required\\'\)}.*", - stdout, + stdout + stderr, ) def test_runtime_parameters_boolean_invalid( self, resources, tmp_path, runtime_param_values_stream ): - stderr = self._test_custom_model_with_runtime_params( + stdout, stderr = self._test_custom_model_with_runtime_params( resources, tmp_path, runtime_param_values_stream, bool_var_value="text" ) assert re.search( r".*Failed to load runtime parameter.*value should be True or False.*", - stderr, + stdout + stderr, ) def test_runtime_parameters_numeric_invalid( self, resources, tmp_path, runtime_param_values_stream ): - stderr = self._test_custom_model_with_runtime_params( + stdout, stderr = self._test_custom_model_with_runtime_params( resources, tmp_path, runtime_param_values_stream, numeric_var_value="text" ) assert re.search( r".*Failed to load runtime parameter.*value can.*t be converted to float.*", - stderr, + stdout + stderr, ) diff --git a/tests/unit/datarobot_drum/drum/utils/test_setup.py b/tests/unit/datarobot_drum/drum/utils/test_setup.py index cb37cb9de..4e67e40ec 100644 --- a/tests/unit/datarobot_drum/drum/utils/test_setup.py +++ b/tests/unit/datarobot_drum/drum/utils/test_setup.py @@ -11,123 +11,151 @@ # Import the function under test from datarobot_drum.drum.utils.setup import setup_options - -def test_setup_options_default(monkeypatch): - # Mock CMRunnerArgsRegistry and dependencies - mock_parser = mock.Mock() - mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None) - mock_parser.parse_args.return_value = mock_options - - with mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", - return_value=mock_parser, - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" - ): - opts = setup_options([]) - assert hasattr(opts, "max_workers") - assert opts.max_workers == 1 # Default to 1 - - -def test_setup_options_with_max_workers(monkeypatch): - mock_parser = mock.Mock() - mock_options = Namespace(max_workers="3", runtime_params_file=None, lazy_loading_file=None) - mock_parser.parse_args.return_value = mock_options - - with mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", - return_value=mock_parser, - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" - ): - opts = setup_options([]) - assert opts.max_workers == 3 or opts.max_workers == "3" or int(opts.max_workers) == 3 - - -def test_setup_options_runtime_parameters(monkeypatch): - mock_parser = mock.Mock() - mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None) - mock_parser.parse_args.return_value = mock_options - - with mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", - return_value=mock_parser, - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" - ), mock.patch( - "datarobot_drum.RuntimeParameters.has", return_value=True - ), mock.patch( - "datarobot_drum.RuntimeParameters.get", return_value=5 - ): - opts = setup_options([]) - assert opts.max_workers == 5 - - -def test_setup_options_runtime_params_file(monkeypatch): - mock_parser = mock.Mock() - mock_options = Namespace( - max_workers=None, runtime_params_file="params.yaml", code_dir=".", lazy_loading_file=None - ) - mock_parser.parse_args.return_value = mock_options - - with mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", - return_value=mock_parser, - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" - ), mock.patch( - "datarobot_drum.drum.utils.setup.RuntimeParametersLoader" - ) as mock_loader: - instance = mock_loader.return_value - instance.setup_environment_variables = mock.Mock() - opts = setup_options([]) - instance.setup_environment_variables.assert_called_once() - - -def test_setup_options_lazy_loading_file(monkeypatch): - mock_parser = mock.Mock() - mock_options = Namespace( - max_workers=None, runtime_params_file=None, lazy_loading_file="lazy.yaml" - ) - mock_parser.parse_args.return_value = mock_options - - with mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", - return_value=mock_parser, - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" - ), mock.patch( - "datarobot_drum.drum.lazy_loading.lazy_loading_handler.LazyLoadingHandler.setup_environment_variables_from_values_file" - ) as mock_lazy: - opts = setup_options([]) - mock_lazy.assert_called_once_with("lazy.yaml") - - -def test_setup_options_argcomplete_missing(monkeypatch, capsys): - mock_parser = mock.Mock() - mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None) - mock_parser.parse_args.return_value = mock_options - - with mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", - return_value=mock_parser, - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" - ), mock.patch( - "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" - ), mock.patch.dict( - "sys.modules", {"argcomplete": None} - ): - opts = setup_options([]) - captured = capsys.readouterr() - assert "autocompletion of arguments is not supported" in captured.err +import os + + +class TestSetupOptions: + def test_setup_options_default(self, monkeypatch): + # Mock CMRunnerArgsRegistry and dependencies + mock_parser = mock.Mock() + mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None) + mock_parser.parse_args.return_value = mock_options + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ): + opts = setup_options([]) + assert hasattr(opts, "max_workers") + assert opts.max_workers == 1 # Default to 1 + + def test_setup_options_with_max_workers(self, monkeypatch): + mock_parser = mock.Mock() + mock_options = Namespace(max_workers="3", runtime_params_file=None, lazy_loading_file=None) + mock_parser.parse_args.return_value = mock_options + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ): + opts = setup_options([]) + assert opts.max_workers == 3 or opts.max_workers == "3" or int(opts.max_workers) == 3 + + def test_setup_options_runtime_parameters(self, monkeypatch): + mock_parser = mock.Mock() + mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None) + mock_parser.parse_args.return_value = mock_options + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ), mock.patch( + "datarobot_drum.RuntimeParameters.has", return_value=True + ), mock.patch( + "datarobot_drum.RuntimeParameters.get", return_value=5 + ): + opts = setup_options([]) + assert opts.max_workers == 5 + + def test_setup_options_runtime_params_file(self, monkeypatch): + mock_parser = mock.Mock() + mock_options = Namespace( + max_workers=None, + runtime_params_file="params.yaml", + code_dir=".", + lazy_loading_file=None, + ) + mock_parser.parse_args.return_value = mock_options + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ), mock.patch( + "datarobot_drum.drum.utils.setup.RuntimeParametersLoader" + ) as mock_loader: + instance = mock_loader.return_value + instance.setup_environment_variables = mock.Mock() + opts = setup_options([]) + instance.setup_environment_variables.assert_called_once() + + def test_setup_options_lazy_loading_file(self, monkeypatch): + mock_parser = mock.Mock() + mock_options = Namespace( + max_workers=None, runtime_params_file=None, lazy_loading_file="lazy.yaml" + ) + mock_parser.parse_args.return_value = mock_options + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ), mock.patch( + "datarobot_drum.drum.lazy_loading.lazy_loading_handler.LazyLoadingHandler.setup_environment_variables_from_values_file" + ) as mock_lazy: + opts = setup_options([]) + mock_lazy.assert_called_once_with("lazy.yaml") + + def test_setup_options_argcomplete_missing(self, monkeypatch, capsys): + mock_parser = mock.Mock() + mock_options = Namespace(max_workers=None, runtime_params_file=None, lazy_loading_file=None) + mock_parser.parse_args.return_value = mock_options + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ), mock.patch.dict( + "sys.modules", {"argcomplete": None} + ): + opts = setup_options([]) + captured = capsys.readouterr() + assert "autocompletion of arguments is not supported" in captured.err + + def test_setup_options_address_env_var(self, monkeypatch): + mock_parser = mock.Mock() + test_address = "127.0.0.1:9999" + mock_options = Namespace( + max_workers=None, runtime_params_file=None, lazy_loading_file=None, address=test_address + ) + mock_parser.parse_args.return_value = mock_options + + # Remove ADDRESS if it exists to avoid side effects + if "ADDRESS" in os.environ: + del os.environ["ADDRESS"] + + with mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.get_arg_parser", + return_value=mock_parser, + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.extend_sys_argv_with_env_vars" + ), mock.patch( + "datarobot_drum.drum.args_parser.CMRunnerArgsRegistry.verify_options" + ): + opts = setup_options([]) + assert os.environ["ADDRESS"] == test_address + + # Clean up + if "ADDRESS" in os.environ: + del os.environ["ADDRESS"]