[Frontend] [Neuron] Parse literals out of override-neuron-config (#8959)

Co-authored-by: Jerzy Zagorski <jzagorsk@amazon.com>
This commit is contained in:
xendo 2024-10-03 20:02:07 +02:00 committed by GitHub
parent f5d72b2fc6
commit 63e39937f9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 37 additions and 20 deletions

View File

@ -42,22 +42,42 @@ def test_bad_nullable_kvs(arg):
nullable_kvs(arg) nullable_kvs(arg)
@pytest.mark.parametrize(("arg", "expected"), [ # yapf: disable
(None, None), @pytest.mark.parametrize(("arg", "expected", "option"), [
("{}", {}), (None, None, "mm-processor-kwargs"),
('{"num_crops": 4}', { ("{}", {}, "mm-processor-kwargs"),
(
'{"num_crops": 4}',
{
"num_crops": 4 "num_crops": 4
}), },
('{"foo": {"bar": "baz"}}', { "mm-processor-kwargs"
"foo": { ),
(
'{"foo": {"bar": "baz"}}',
{
"foo":
{
"bar": "baz" "bar": "baz"
} }
}), },
"mm-processor-kwargs"
),
(
'{"cast_logits_dtype":"bfloat16","sequence_parallel_norm":true,"sequence_parallel_norm_threshold":2048}',
{
"cast_logits_dtype": "bfloat16",
"sequence_parallel_norm": True,
"sequence_parallel_norm_threshold": 2048,
},
"override-neuron-config"
),
]) ])
def test_mm_processor_kwargs_prompt_parser(arg, expected): # yapf: enable
def test_composite_arg_parser(arg, expected, option):
parser = EngineArgs.add_cli_args(FlexibleArgumentParser()) parser = EngineArgs.add_cli_args(FlexibleArgumentParser())
if arg is None: if arg is None:
args = parser.parse_args([]) args = parser.parse_args([])
else: else:
args = parser.parse_args(["--mm-processor-kwargs", arg]) args = parser.parse_args([f"--{option}", arg])
assert args.mm_processor_kwargs == expected assert getattr(args, option.replace("-", "_")) == expected

View File

@ -800,13 +800,10 @@ class EngineArgs:
"lower performance.") "lower performance.")
parser.add_argument( parser.add_argument(
'--override-neuron-config', '--override-neuron-config',
type=lambda configs: { type=json.loads,
str(key): value
for key, value in
(config.split(':') for config in configs.split(','))
},
default=None, default=None,
help="override or set neuron device configuration.") help="Override or set neuron device configuration. "
"e.g. {\"cast_logits_dtype\": \"bloat16\"}.'")
parser.add_argument( parser.add_argument(
'--scheduling-policy', '--scheduling-policy',