diff --git a/morpheus/cli/commands.py b/morpheus/cli/commands.py index 99df8e1ff5..21c134657c 100644 --- a/morpheus/cli/commands.py +++ b/morpheus/cli/commands.py @@ -288,13 +288,22 @@ def install(**kwargs): type=bool, help=("Whether or not to use C++ node and message types or to prefer python. " "Only use as a last resort if bugs are encountered")) +@click.option('--manual_seed', + default=None, + type=click.IntRange(min=1), + envvar="MORPHEUS_MANUAL_SEED", + help=("Manually seed the random number generators used by Morpheus, useful for testing.")) @prepare_command(parse_config=True) def run(ctx: click.Context, **kwargs): """Run subcommand, used for running a pipeline""" # Since the option isnt the same name as `should_use_cpp` anymore, manually set the value here. CppConfig.set_should_use_cpp(kwargs.pop("use_cpp", CppConfig.get_should_use_cpp())) - pass + manual_seed_val = kwargs.pop("manual_seed", None) + if manual_seed_val is not None: + from morpheus.utils.seed import manual_seed + logger.debug("Manually seeding random number generators to %d", manual_seed_val) + manual_seed(manual_seed_val) @click.group(chain=True, diff --git a/tests/test_cli.py b/tests/test_cli.py index 3cb0efea43..aef467bf84 100755 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -16,6 +16,7 @@ import os import shutil +from unittest import mock import click import mlflow @@ -157,6 +158,24 @@ def test_autocomplete(self, tmp_path): env={'HOME': str(tmp_path)}) assert result.exit_code == 0, result.output + @pytest.mark.usefixtures("restore_environ") + @pytest.mark.parametrize('use_environ', [True, False]) + @pytest.mark.parametrize('value', [1, 13, 33]) + @mock.patch('morpheus.utils.seed.manual_seed') + def test_manual_seed(self, mock_manual_seed: mock.MagicMock, value: int, use_environ: bool): + flags = ['run'] + if use_environ: + os.environ['MORPHEUS_MANUAL_SEED'] = str(value) + else: + flags.append(f'--manual_seed={value}') + + flags.append('pipeline-other') + + runner = CliRunner() + result = runner.invoke(commands.cli, flags) + assert result.exit_code == 0, result.output + mock_manual_seed.assert_called_once_with(value) + @pytest.mark.replace_callback('pipeline_ae') def test_pipeline_ae(self, config, callback_values): """