"docs/pages/vscode:/vscode.git/clone" did not exist on "673822ea931352cebe5d6e49a9b66f66a8f1d272"
Unverified Commit 92064030 authored by hhzhang16's avatar hhzhang16 Committed by GitHub
Browse files

fix: json strings should remain intact through profiler arg processing (#3680)


Signed-off-by: default avatarHannah Zhang <hannahz@nvidia.com>
parent b01541c4
...@@ -125,8 +125,17 @@ def break_arguments(args: list[str] | None) -> list[str]: ...@@ -125,8 +125,17 @@ def break_arguments(args: list[str] | None) -> list[str]:
else: else:
for arg in args: for arg in args:
if arg is not None: if arg is not None:
# If the arg looks like it might be JSON (starts with { or [) or is already a single token,
# don't split it further. Only split if it contains spaces AND doesn't look like JSON.
if (
isinstance(arg, str)
and (" " in arg or "\t" in arg)
and not (arg.strip().startswith(("{", "[")))
):
# Use shlex.split to properly handle quoted arguments # Use shlex.split to properly handle quoted arguments
ans.extend(shlex.split(arg)) ans.extend(shlex.split(arg))
else:
ans.append(arg)
return ans return ans
......
...@@ -95,7 +95,9 @@ class TrtllmConfigModifier: ...@@ -95,7 +95,9 @@ class TrtllmConfigModifier:
# - Disable enable_block_reuse (no KV reuse for prefill-only) # - Disable enable_block_reuse (no KV reuse for prefill-only)
# - Enable overlap scheduler (disabled in prefill.yaml but needed for agg) # - Enable overlap scheduler (disabled in prefill.yaml but needed for agg)
# - Remove cache_transceiver_config (not needed in agg mode) # - Remove cache_transceiver_config (not needed in agg mode)
if "kv_cache_config" not in override_dict: if "kv_cache_config" not in override_dict or not isinstance(
override_dict["kv_cache_config"], dict
):
override_dict["kv_cache_config"] = {} override_dict["kv_cache_config"] = {}
override_dict["kv_cache_config"]["enable_block_reuse"] = False override_dict["kv_cache_config"]["enable_block_reuse"] = False
override_dict[ override_dict[
...@@ -146,7 +148,9 @@ class TrtllmConfigModifier: ...@@ -146,7 +148,9 @@ class TrtllmConfigModifier:
# Merge our overrides for converting decode-only disagg to aggregated: # Merge our overrides for converting decode-only disagg to aggregated:
# - Enable enable_block_reuse (to skip prefill in decode-only) # - Enable enable_block_reuse (to skip prefill in decode-only)
# - Remove cache_transceiver_config (not needed in agg mode) # - Remove cache_transceiver_config (not needed in agg mode)
if "kv_cache_config" not in override_dict: if "kv_cache_config" not in override_dict or not isinstance(
override_dict["kv_cache_config"], dict
):
override_dict["kv_cache_config"] = {} override_dict["kv_cache_config"] = {}
override_dict["kv_cache_config"]["enable_block_reuse"] = True override_dict["kv_cache_config"]["enable_block_reuse"] = True
override_dict[ override_dict[
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment