minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] feat: add composable component weight path cli (#36078)

Browse Source
This commit is contained in:
Mick
2026-08-24 09:31:08 +08:00
committed by GitHub
parent 1c1c9d9b4e
commit fee00a41db
3 changed files with 105 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+58 -24
View File
@@ -2760,38 +2760,38 @@ class ServerArgs(DisaggServerArgsMixin):
)
@staticmethod
def _extract_component_paths(
def _extract_dynamic_component_map(
unknown_args: list[str],
*,
option_prefixes: tuple[str, ...],
alias_suffix: str,
) -> tuple[dict[str, str], list[str]]:
"""
Extract dynamic component path args from unrecognised CLI args.
Supported forms:
- ``--<component>-path /path/to/component``
- ``--component-paths.<component> /path/to/component`` (expanded from config)
"""
component_paths: dict[str, str] = {}
component_values: dict[str, str] = {}
remaining: list[str] = []
i = 0
while i < len(unknown_args):
arg = unknown_args[i]
key_part = arg.split("=", 1)[0] if "=" in arg else arg
component = None
if key_part.startswith("--component-paths."):
component = key_part[len("--component-paths.") :].replace("-", "_")
elif key_part.startswith("--component_paths."):
component = key_part[len("--component_paths.") :].replace("-", "_")
elif key_part.startswith("--") and key_part.endswith("-path"):
component = key_part[2:-5].replace("-", "_")
for option_prefix in option_prefixes:
if key_part.startswith(option_prefix):
component = key_part[len(option_prefix) :].replace("-", "_")
break
if (
component is None
and key_part.startswith("--")
and key_part.endswith(alias_suffix)
):
component = key_part[2 : -len(alias_suffix)].replace("-", "_")
if component is not None:
if "=" in arg:
component_paths[component] = arg.split("=", 1)[1]
component_values[component] = arg.split("=", 1)[1]
elif i + 1 < len(unknown_args) and not unknown_args[i + 1].startswith(
"-"
):
i += 1
component_paths[component] = unknown_args[i]
component_values[component] = unknown_args[i]
else:
remaining.append(arg)
i += 1
@@ -2800,11 +2800,37 @@ class ServerArgs(DisaggServerArgsMixin):
remaining.append(arg)
i += 1
# canonicalize and validate
for component, path in component_paths.items():
path = os.path.expanduser(path)
component_paths[component] = path
return component_paths, remaining
return {
component: os.path.expanduser(value)
for component, value in component_values.items()
}, remaining
@classmethod
def _extract_component_paths(
cls,
unknown_args: list[str],
) -> tuple[dict[str, str], list[str]]:
"""Extract dynamic component configuration paths and aliases."""
return cls._extract_dynamic_component_map(
unknown_args,
option_prefixes=("--component-paths.", "--component_paths."),
alias_suffix="-path",
)
@classmethod
def _extract_component_weights_paths(
cls,
unknown_args: list[str],
) -> tuple[dict[str, str], list[str]]:
"""Extract dynamic component weight-file paths and aliases."""
return cls._extract_dynamic_component_map(
unknown_args,
option_prefixes=(
"--component-weights-paths.",
"--component_weights_paths.",
),
alias_suffix="-weights-path",
)
@staticmethod
def _extract_component_attention_backends(
@@ -2853,8 +2879,11 @@ class ServerArgs(DisaggServerArgsMixin):
if unknown_args is None:
unknown_args = []
# extract dynamic --<component>-path from unknown args
dynamic_paths, remaining = cls._extract_component_paths(unknown_args)
# Extract the more specific weights suffix before the generic path alias.
dynamic_weights_paths, remaining = cls._extract_component_weights_paths(
unknown_args
)
dynamic_paths, remaining = cls._extract_component_paths(remaining)
dynamic_attention_backends, remaining = (
cls._extract_component_attention_backends(remaining)
)
@@ -2880,6 +2909,11 @@ class ServerArgs(DisaggServerArgsMixin):
existing.update(dynamic_paths)
provided_args["component_paths"] = existing
explicit_arg_names.add("component_paths")
if dynamic_weights_paths:
existing = dict(provided_args.get("component_weights_paths") or {})
existing.update(dynamic_weights_paths)
provided_args["component_weights_paths"] = existing
explicit_arg_names.add("component_weights_paths")
if dynamic_attention_backends:
existing = cls._parse_component_attention_backend_map(
provided_args.get("component_attention_backends")
python/sglang/multimodal_gen/test/unit/test_server_args.py
+21 -1
View File
@@ -464,7 +464,16 @@ class TestServerArgsPathExpansion(unittest.TestCase):
)
with tempfile.NamedTemporaryFile("w", suffix=".json") as config_file:
json.dump({"model_path": "/from/config", "num_gpus": 2}, config_file)
json.dump(
{
"model_path": "/from/config",
"num_gpus": 2,
"component_weights_paths": {
"transformer": "owner/repo/transformer.safetensors"
},
},
config_file,
)
config_file.flush()
parser = FlexibleArgumentParser()
add_multimodal_gen_serve_args(parser)
@@ -475,6 +484,9 @@ class TestServerArgsPathExpansion(unittest.TestCase):
"/from/cli",
"--vae-path",
"/custom/vae",
"--component-weights-paths.text_encoder",
"owner/repo/text_encoder.safetensors",
"--image-encoder-weights-path=/custom/image_encoder.safetensors",
"--component-attention-backends.transformer",
"fa3",
]
@@ -505,6 +517,14 @@ class TestServerArgsPathExpansion(unittest.TestCase):
self.assertEqual("/from/cli", server_args.model_path)
self.assertEqual(2, server_args.num_gpus)
self.assertEqual("/custom/vae", server_args.component_paths["vae"])
self.assertEqual(
{
"transformer": "owner/repo/transformer.safetensors",
"text_encoder": "owner/repo/text_encoder.safetensors",
"image_encoder": "/custom/image_encoder.safetensors",
},
server_args.component_weights_paths,
)
self.assertEqual(
{"transformer": "fa"},
server_args.component_attention_backends,