config.py 45.6 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

16
import json
17
import logging
18
import math
19
import re
20
import shlex
21
from typing import Literal, Optional, Protocol
22

23
import yaml
24
from pydantic import BaseModel
25

26
27
28
29
from benchmarks.profiler.utils.defaults import (
    DEFAULT_MODEL_NAME,
    DYNAMO_RUN_DEFAULT_PORT,
)
30
from benchmarks.profiler.utils.planner_utils import build_planner_args_from_namespace
31
from dynamo.planner.defaults import WORKER_COMPONENT_NAMES, SubComponentType
32
33
34
35
36
37
38
39
40
41
42
43

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
console_handler = logging.StreamHandler()
console_handler.setLevel(logging.INFO)
formatter = logging.Formatter(
    "%(asctime)s - %(name)s - %(levelname)s - %(message)s", "%Y-%m-%d %H:%M:%S"
)
console_handler.setFormatter(formatter)
logger.addHandler(console_handler)


44
45
46
47
48
class VolumeMount(BaseModel):
    name: str = "dynamo-pvc"
    mountPoint: str = "/data"


49
class Container(BaseModel):
50
51
52
    image: Optional[str] = None
    workingDir: Optional[str] = None
    command: Optional[list[str]] = None
53
54
    args: Optional[list[str]] = None
    model_config = {"extra": "allow"}
55
56
57


class PodSpec(BaseModel):
58
59
    mainContainer: Optional[Container] = None
    model_config = {"extra": "allow"}
60
61
62


class ServiceResources(BaseModel):
63
    requests: Optional[dict[str, str]] = None
64
65
66
67
    limits: Optional[dict[str, str]] = None


class Service(BaseModel):
68
69
70
    replicas: Optional[int] = None
    resources: Optional[ServiceResources] = None
    extraPodSpec: Optional[PodSpec] = None
71
    subComponentType: Optional[str] = None
72
    model_config = {"extra": "allow"}
73
74
75
76


class Services(BaseModel):
    Frontend: Service
77
    model_config = {"extra": "allow"}
78
79


80
81
82
83
84
85
class PVCConfig(BaseModel):
    name: str = "dynamo-pvc"
    create: Optional[bool] = False
    model_config = {"extra": "allow"}


86
87
class Spec(BaseModel):
    services: dict[str, Service]
88
89
    pvcs: Optional[list[PVCConfig]] = None
    model_config = {"extra": "allow"}
90
91
92
93


class Metadata(BaseModel):
    name: str
94
    model_config = {"extra": "allow"}
95
96
97
98
99


class Config(BaseModel):
    metadata: Metadata
    spec: Spec
100
    model_config = {"extra": "allow"}
101
102


103
104
105
106
class MultinodeConfig(BaseModel):
    nodeCount: int


107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
class DgdPlannerServiceConfig(BaseModel):
    dynamoNamespace: str = "dynamo"  # placeholder
    componentType: str = "planner"
    replicas: int = 1
    volumeMounts: list[VolumeMount] = [VolumeMount()]
    extraPodSpec: PodSpec = PodSpec(
        mainContainer=Container(
            image="my-registry/dynamo-runtime:my-tag",  # placeholder
            workingDir="/workspace/components/src/dynamo/planner",
            command=["python3", "-m", "planner_sla"],
            args=[],
        )
    )
    model_config = {"extra": "allow"}


123
124
125
126
def break_arguments(args: list[str] | None) -> list[str]:
    ans: list[str] = []
    if args is None:
        return ans
127
    if isinstance(args, str):
128
129
        # Use shlex.split to properly handle quoted arguments and JSON values
        ans = shlex.split(args)
130
131
    else:
        for arg in args:
132
            if arg is not None:
133
134
                # Use shlex.split to properly handle quoted arguments
                ans.extend(shlex.split(arg))
135
    return ans
136
137


138
139
140
141
142
143
144
145
146
147
def remove_valued_arguments(args: list[str], key: str) -> list[str]:
    """Remove a valued argument (e.g., --key value) from the arguments list if exists."""
    if key in args:
        idx = args.index(key)
        if idx + 1 < len(args):
            del args[idx : idx + 2]

    return args


148
def join_arguments(args: list[str]) -> list[str]:
149
150
    # Use shlex.join to properly quote arguments that contain spaces or special characters
    return [shlex.join(args)]
151
152


153
154
155
156
157
158
159
def append_argument(args: list[str], to_append) -> list[str]:
    idx = find_arg_index(args)
    if isinstance(to_append, list):
        args[idx:idx] = to_append
    else:
        args.insert(idx, to_append)
    return args
160
161


162
163
164
def find_arg_index(args: list[str]) -> int:
    # find the correct index to insert an argument
    idx = len(args)
165

166
167
168
169
170
    try:
        new_idx = args.index("|")
        idx = min(idx, new_idx)
    except ValueError:
        pass
171

172
173
174
175
176
    try:
        new_idx = args.index("2>&1")
        idx = min(idx, new_idx)
    except ValueError:
        pass
177

178
    return idx
179
180


181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def parse_override_engine_args(args: list[str]) -> tuple[dict, list[str]]:
    """
    Parse and extract --override-engine-args from argument list.

    Returns:
        tuple: (override_dict, modified_args) where override_dict is the parsed JSON
               and modified_args is the args list with --override-engine-args removed
    """
    override_dict = {}
    try:
        idx = args.index("--override-engine-args")
        if idx + 1 < len(args):
            # Parse existing override
            override_dict = json.loads(args[idx + 1])
            # Remove the old override args
            del args[idx : idx + 2]
    except (ValueError, json.JSONDecodeError):
        pass  # No existing override or invalid JSON

    return override_dict, args


203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
def set_multinode_config(worker_service, gpu_count: int, num_gpus_per_node: int):
    """Helper function to set multinode configuration based on GPU count and GPUs per node."""
    if gpu_count <= num_gpus_per_node:
        # Single node: remove multinode configuration if present
        if (
            hasattr(worker_service, "multinode")
            and worker_service.multinode is not None
        ):
            worker_service.multinode = None
    else:
        # Multi-node: set nodeCount = math.ceil(gpu_count / num_gpus_per_node)
        node_count = math.ceil(gpu_count / num_gpus_per_node)
        if not hasattr(worker_service, "multinode") or worker_service.multinode is None:
            # Create multinode configuration if it doesn't exist
            worker_service.multinode = MultinodeConfig(nodeCount=node_count)
        else:
            # Handle both dict (from YAML) and MultinodeConfig object cases
            if isinstance(worker_service.multinode, dict):
                worker_service.multinode["nodeCount"] = node_count
            else:
                worker_service.multinode.nodeCount = node_count


226
def get_service_name_by_type(
227
    config: Config, backend: str, sub_component_type: SubComponentType
228
229
230
231
232
233
) -> str:
    """Helper function to get service name by subComponentType.

    First tries to find service by subComponentType, then falls back to component name.

    Args:
234
        config: Configuration object
235
236
237
238
239
240
241
        backend: Backend name (e.g., "sglang", "vllm", "trtllm")
        sub_component_type: The type of sub-component to look for (PREFILL or DECODE)

    Returns:
        The service name
    """
    # Check if config has the expected structure
242
    if not config.spec or not config.spec.services:
243
244
245
246
247
248
249
        # Fall back to default name if structure is unexpected
        if sub_component_type == SubComponentType.DECODE:
            return WORKER_COMPONENT_NAMES[backend].decode_worker_k8s_name
        else:
            return WORKER_COMPONENT_NAMES[backend].prefill_worker_k8s_name

    # Look through services to find one with matching subComponentType
250
    services = config.spec.services
251
    for service_name, service_config in services.items():
252
        if service_config.subComponentType == sub_component_type.value:
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
            return service_name

    # Fall back to default component names
    if sub_component_type == SubComponentType.DECODE:
        default_name = WORKER_COMPONENT_NAMES[backend].decode_worker_k8s_name
    else:
        default_name = WORKER_COMPONENT_NAMES[backend].prefill_worker_k8s_name

    # Check if the default name exists in services
    if default_name in services:
        return default_name

    # Last resort: return the default name anyway
    return default_name


def get_worker_service_from_config(
270
    config: Config,
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
    backend: str = "sglang",
    sub_component_type: SubComponentType = SubComponentType.DECODE,
):
    """Helper function to get a worker service from config.

    First tries to find service by subComponentType, then falls back to component name.

    Args:
        config: Configuration dictionary
        backend: Backend name (e.g., "sglang", "vllm", "trtllm"). Defaults to "sglang".
        sub_component_type: The type of sub-component to look for (PREFILL or DECODE). Defaults to DECODE.

    Returns:
        The worker service from the configuration
    """
    if backend not in WORKER_COMPONENT_NAMES:
        raise ValueError(
            f"Unsupported backend: {backend}. Supported backends: {list(WORKER_COMPONENT_NAMES.keys())}"
        )

    # Get the service name using the type-aware logic
    service_name = get_service_name_by_type(config, backend, sub_component_type)

    # Get the actual service from the config
295
    return config.spec.services[service_name]
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326


def setup_worker_service_resources(
    worker_service, gpu_count: int, num_gpus_per_node: Optional[int] = None
):
    """Helper function to set up worker service resources (requests and limits)."""
    # Handle multinode configuration if num_gpus_per_node is provided
    if num_gpus_per_node is not None:
        set_multinode_config(worker_service, gpu_count, num_gpus_per_node)

    # Ensure resources exists
    if worker_service.resources is None:
        worker_service.resources = ServiceResources()

    # Ensure requests exists
    if worker_service.resources.requests is None:
        worker_service.resources.requests = {}

    # Set GPU requests
    gpu_value = (
        min(gpu_count, num_gpus_per_node)
        if num_gpus_per_node is not None
        else gpu_count
    )
    worker_service.resources.requests["gpu"] = str(gpu_value)

    # Update limits if they exist
    if worker_service.resources.limits is not None:
        worker_service.resources.limits["gpu"] = str(gpu_value)


327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
def validate_and_get_worker_args(worker_service, backend):
    """Helper function to validate worker service and get its arguments.

    Args:
        worker_service: Worker service object to validate
        backend: Backend name (e.g., "sglang", "vllm", "trtllm"). Defaults to "sglang".

    Returns:
        List of arguments from the worker service
    """
    if backend not in WORKER_COMPONENT_NAMES:
        raise ValueError(
            f"Unsupported backend: {backend}. Supported backends: {list(WORKER_COMPONENT_NAMES.keys())}"
        )

342
343
    if not worker_service.extraPodSpec or not worker_service.extraPodSpec.mainContainer:
        raise ValueError(
344
            f"Missing extraPodSpec or mainContainer in {backend} decode worker service '{WORKER_COMPONENT_NAMES[backend].decode_worker_k8s_name}'"
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
        )

    args = worker_service.extraPodSpec.mainContainer.args
    return break_arguments(args)


def set_argument_value(args: list, arg_name: str, value: str):
    """Helper function to set an argument value, adding it if not present."""
    try:
        idx = args.index(arg_name)
        args[idx + 1] = value
    except ValueError:
        args = append_argument(args, [arg_name, value])
    return args


361
362
class ConfigModifierProtocol(Protocol):
    @classmethod
363
364
365
366
367
368
    def convert_config(
        cls,
        config: dict,
        target: Literal["prefill", "decode"],
        is_moe_model: bool = False,
    ) -> dict:
369
370
371
        ...

    @classmethod
372
373
374
375
376
377
    def set_config_tp_size(
        cls,
        config: dict,
        tp_size: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ) -> dict:
378
379
        ...

380
381
    @classmethod
    def set_config_tep_size(
382
383
384
385
386
        cls,
        config: dict,
        tep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
387
388
389
390
391
    ) -> dict:
        ...

    @classmethod
    def set_config_dep_size(
392
393
394
395
396
        cls,
        config: dict,
        dep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
397
398
399
    ) -> dict:
        ...

400
401
402
403
404
405
406
407
408
    @classmethod
    def get_model_name(cls, config: dict) -> str:
        ...

    @classmethod
    def get_port(cls, config: dict) -> int:
        ...

    @classmethod
409
410
411
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
412
413
414
        ...


415
416
class VllmV1ConfigModifier:
    @classmethod
417
418
419
420
421
422
423
424
425
426
427
    def convert_config(
        cls,
        config: dict,
        target: Literal["prefill", "decode"],
        is_moe_model: bool = False,
    ) -> dict:
        if is_moe_model:
            raise NotImplementedError(
                "MoE model support is not implemented for VLLM backend"
            )

428
        cfg = Config.model_validate(config)
429

430
        # set metadata name
431
        cfg.metadata.name = "vllm-agg"
432

433
        # disable planner
434
435
        if "Planner" in cfg.spec.services:
            del cfg.spec.services["Planner"]
436
437

        if target == "prefill":
438
439
            # Get service names by inferring from subComponentType first
            prefill_service_name = get_service_name_by_type(
440
                cfg, "vllm", SubComponentType.PREFILL
441
442
            )
            decode_service_name = get_service_name_by_type(
443
                cfg, "vllm", SubComponentType.DECODE
444
445
            )

446
            # convert prefill worker into decode worker
447
448
            cfg.spec.services[decode_service_name] = cfg.spec.services[
                prefill_service_name
449
            ]
450
            del cfg.spec.services[prefill_service_name]
451

452
453
            # Set subComponentType for aggregated mode (using decode worker for prefill-only)
            cfg.spec.services[decode_service_name].subComponentType = "decode"
454

455
            worker_service = get_worker_service_from_config(
456
                cfg,
457
458
459
460
                backend="vllm",
                sub_component_type=SubComponentType.DECODE,
            )
            args = validate_and_get_worker_args(worker_service, backend="vllm")
461
462
463
464
465
466
467
468
469
470
471
            args = break_arguments(args)

            # remove --is-prefill-worker flag
            args.remove("--is-prefill-worker")

            # disable prefix caching
            if "--enable-prefix-caching" in args:
                args.remove("--enable-prefix-caching")
            if "--no-enable-prefix-caching" not in args:
                args = append_argument(args, "--no-enable-prefix-caching")

472
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
473

474
        elif target == "decode":
475
476
            # Get service names by inferring from subComponentType first
            prefill_service_name = get_service_name_by_type(
477
                cfg, "vllm", SubComponentType.PREFILL
478
479
            )
            decode_service_name = get_service_name_by_type(
480
                cfg, "vllm", SubComponentType.DECODE
481
482
            )

483
            # delete prefill worker
484
            del cfg.spec.services[prefill_service_name]
485

486
487
            # Set subComponentType for aggregated decode-only mode
            cfg.spec.services[decode_service_name].subComponentType = "decode"
488

489
            worker_service = get_worker_service_from_config(
490
                cfg,
491
492
493
494
                backend="vllm",
                sub_component_type=SubComponentType.DECODE,
            )
            args = validate_and_get_worker_args(worker_service, backend="vllm")
495
            args = break_arguments(args)
496

497
498
499
500
501
502
            # enable prefix caching
            if "--enable-prefix-caching" not in args:
                args = append_argument(args, "--enable-prefix-caching")
            if "--no-enable-prefix-caching" in args:
                args.remove("--no-enable-prefix-caching")

503
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
504
505

        # set num workers to 1
506
507
        # Use the inferred decode service name
        final_decode_service_name = get_service_name_by_type(
508
            cfg, "vllm", SubComponentType.DECODE
509
510
        )
        decode_worker_config = cfg.spec.services[final_decode_service_name]
511
        decode_worker_config.replicas = 1
512

513
        return cfg.model_dump()
514
515

    @classmethod
516
517
518
519
520
521
    def set_config_tp_size(
        cls,
        config: dict,
        tp_size: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
522
        cfg = Config.model_validate(config)
523
        worker_service = get_worker_service_from_config(
524
            cfg, backend="vllm", sub_component_type=component_type
525
        )
526

527
528
        # Set up resources
        setup_worker_service_resources(worker_service, tp_size)
529

530
531
        # Get and validate args
        args = validate_and_get_worker_args(worker_service, backend="vllm")
532
533
534
535
536
537
538
539
        args = break_arguments(args)

        try:
            idx = args.index("--tensor-parallel-size")
            args[idx + 1] = str(tp_size)
        except ValueError:
            args = append_argument(args, ["--tensor-parallel-size", str(tp_size)])

540
        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
541

542
        return cfg.model_dump()
543

544
    @classmethod
545
546
547
548
549
550
551
    def set_config_tep_size(
        cls,
        config: dict,
        tep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
552
553
554
555
556
        raise NotImplementedError(
            "TEP (Tensor Expert Parallelism) is not implemented for VLLM backend"
        )

    @classmethod
557
558
559
560
561
562
563
    def set_config_dep_size(
        cls,
        config: dict,
        dep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
564
565
566
567
        raise NotImplementedError(
            "DEP (Data Expert Parallelism) is not implemented for VLLM backend"
        )

568
569
    @classmethod
    def get_model_name(cls, config: dict) -> str:
570
        cfg = Config.model_validate(config)
571
        try:
572
            worker_service = get_worker_service_from_config(cfg, backend="vllm")
573
574
            args = validate_and_get_worker_args(worker_service, backend="vllm")
        except (ValueError, KeyError):
575
            logger.warning(
576
                f"Worker service missing or invalid, using default model name: {DEFAULT_MODEL_NAME}"
577
578
            )
            return DEFAULT_MODEL_NAME
579
580
581
582
583
584
585
586
587
588

        args = break_arguments(args)
        for i, arg in enumerate(args):
            if arg == "--model" and i + 1 < len(args):
                return args[i + 1]

        logger.warning(
            f"Model name not found in configuration args, using default model name: {DEFAULT_MODEL_NAME}"
        )
        return DEFAULT_MODEL_NAME
589
590
591

    @classmethod
    def get_port(cls, config: dict) -> int:
592
        cfg = Config.model_validate(config)
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
        frontend_service = cfg.spec.services.get("Frontend")
        if (
            not frontend_service
            or not frontend_service.extraPodSpec
            or not frontend_service.extraPodSpec.mainContainer
        ):
            logger.warning(
                f"Frontend service or container not found, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

        args = frontend_service.extraPodSpec.mainContainer.args
        if not args:
            logger.warning(
                f"No args found in Frontend configuration, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

611
612
613
614
        args = break_arguments(args)
        try:
            idx = args.index("--http-port")
            return int(args[idx + 1])
615
        except (ValueError, IndexError):
616
617
618
619
            logger.warning(
                f"Port not found in configuration args, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT
620
621

    @classmethod
622
623
624
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
        try:
            with open(dynamo_log_fn, "r") as f:
                for line in f:
                    if "Maximum concurrency for" in line:
                        line = line.strip().split("Maximum concurrency for ")[1]
                        token_count = int(
                            line.split(" tokens per request: ")[0].replace(",", "")
                        )
                        concurrency = float(line.split(" tokens per request: ")[1][:-1])

                        logger.info(
                            f"Found KV cache info: {token_count} x {concurrency} = {int(token_count * concurrency)}"
                        )
                        return int(token_count * concurrency)
        except Exception as e:
            logger.warning(
                f"Failed to parse KV cache size from line: {line}. Error: {e}"
            )
        return 0


646
647
class SGLangConfigModifier:
    @classmethod
648
649
650
651
652
653
    def convert_config(
        cls,
        config: dict,
        target: Literal["prefill", "decode"],
        is_moe_model: bool = False,
    ) -> dict:
654
        cfg = Config.model_validate(config)
655
656

        # set metadata name
657
        cfg.metadata.name = "sglang-agg"
658
659

        # disable planner
660
661
        if "Planner" in cfg.spec.services:
            del cfg.spec.services["Planner"]
662
663

        if target == "prefill":
664
665
            # Get service names by inferring from subComponentType first
            prefill_service_name = get_service_name_by_type(
666
                cfg, "sglang", SubComponentType.PREFILL
667
668
            )
            decode_service_name = get_service_name_by_type(
669
                cfg, "sglang", SubComponentType.DECODE
670
671
            )

672
            # convert prefill worker into decode worker
673
674
            cfg.spec.services[decode_service_name] = cfg.spec.services[
                prefill_service_name
675
            ]
676
            del cfg.spec.services[prefill_service_name]
677

678
679
            # Set subComponentType for aggregated mode (using decode worker for prefill-only)
            cfg.spec.services[decode_service_name].subComponentType = "decode"
680

681
            worker_service = get_worker_service_from_config(
682
                cfg,
683
684
685
686
                backend="sglang",
                sub_component_type=SubComponentType.DECODE,
            )
            args = validate_and_get_worker_args(worker_service, backend="sglang")
687
688
            args = break_arguments(args)

689
            # remove disagg flags
690
691
            args = remove_valued_arguments(args, "--disaggregation-mode")
            args = remove_valued_arguments(args, "--disaggregation-transfer-backend")
692
            args = remove_valued_arguments(args, "--disaggregation-bootstrap-port")
693
694
695
696
697

            # disable prefix caching
            if "--disable-radix-cache" not in args:
                args = append_argument(args, "--disable-radix-cache")

698
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
699
700

        elif target == "decode":
701
702
            # Get service names by inferring from subComponentType first
            prefill_service_name = get_service_name_by_type(
703
                cfg, "sglang", SubComponentType.PREFILL
704
705
            )
            decode_service_name = get_service_name_by_type(
706
                cfg, "sglang", SubComponentType.DECODE
707
708
            )

709
            # delete prefill worker
710
            del cfg.spec.services[prefill_service_name]
711

712
713
            # Set subComponentType for aggregated decode-only mode
            cfg.spec.services[decode_service_name].subComponentType = "decode"
714

715
            worker_service = get_worker_service_from_config(
716
                cfg,
717
718
719
720
                backend="sglang",
                sub_component_type=SubComponentType.DECODE,
            )
            args = validate_and_get_worker_args(worker_service, backend="sglang")
721
722
            args = break_arguments(args)

723
            # remove disagg flags
724
725
            args = remove_valued_arguments(args, "--disaggregation-mode")
            args = remove_valued_arguments(args, "--disaggregation-transfer-backend")
726
            args = remove_valued_arguments(args, "--disaggregation-bootstrap-port")
727
728
729
730
731

            # enable prefix caching
            if "--disable-radix-cache" in args:
                args.remove("--disable-radix-cache")

732
733
734
735
736
737
738
739
740
741
            if is_moe_model:
                # need to use round_robin dp attention routing for MoE models to ensure kv reuse can skip prefill
                if "--load-balance-method" in args:
                    idx = args.index("--load-balance-method")
                    args[idx + 1] = "round_robin"
                else:
                    args = append_argument(
                        args, ["--load-balance-method", "round_robin"]
                    )

742
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
743
744

        # set num workers to 1
745
746
        # Use the inferred decode service name
        final_decode_service_name = get_service_name_by_type(
747
            cfg, "sglang", SubComponentType.DECODE
748
749
750
        )
        decode_worker_config = cfg.spec.services[final_decode_service_name]
        decode_worker_config.replicas = 1
751

752
        return cfg.model_dump()
753
754

    @classmethod
755
756
757
758
759
760
    def set_config_tp_size(
        cls,
        config: dict,
        tp_size: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
761
        cfg = Config.model_validate(config)
762
        worker_service = get_worker_service_from_config(
763
            cfg, backend="sglang", sub_component_type=component_type
764
        )
765

766
767
        # Set up resources
        setup_worker_service_resources(worker_service, tp_size)
768

769
        # Get and validate args
770
        args = validate_and_get_worker_args(worker_service, backend="sglang")
771

772
773
        # Set --tp argument
        args = set_argument_value(args, "--tp", str(tp_size))
774

775
776
        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
        return cfg.model_dump()
777

778
    @classmethod
779
780
781
782
783
784
785
    def set_config_tep_size(
        cls,
        config: dict,
        tep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
786
        cfg = Config.model_validate(config)
787
        worker_service = get_worker_service_from_config(
788
            cfg, backend="sglang", sub_component_type=component_type
789
        )
790

791
792
        # Set up resources with multinode configuration
        setup_worker_service_resources(worker_service, tep_size, num_gpus_per_node)
793

794
        # Get and validate args
795
        args = validate_and_get_worker_args(worker_service, backend="sglang")
796

797
798
799
800
801
802
803
804
805
806
807
808
        # 1. Set --tp=tep_size, if not present add it
        args = set_argument_value(args, "--tp", str(tep_size))

        # 2. Set --ep-size=tep_size, if not present add it
        args = set_argument_value(args, "--ep-size", str(tep_size))

        # 3. Remove --dp if present
        args = remove_valued_arguments(args, "--dp")

        # 4. Remove --enable-dp-attention if present
        if "--enable-dp-attention" in args:
            args.remove("--enable-dp-attention")
809

810
        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
811
812
813
        return cfg.model_dump()

    @classmethod
814
815
816
817
818
819
820
    def set_config_dep_size(
        cls,
        config: dict,
        dep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
821
        cfg = Config.model_validate(config)
822
        worker_service = get_worker_service_from_config(
823
            cfg, backend="sglang", sub_component_type=component_type
824
        )
825
826
827
828
829

        # Set up resources with multinode configuration
        setup_worker_service_resources(worker_service, dep_size, num_gpus_per_node)

        # Get and validate args
830
        args = validate_and_get_worker_args(worker_service, backend="sglang")
831
832
833
834
835
836
837
838
839
840

        # 1. Set --tp=dep_size
        args = set_argument_value(args, "--tp", str(dep_size))

        # 2. Set --dp=dep_size (data parallelism across experts)
        args = set_argument_value(args, "--dp", str(dep_size))

        # 3. Enable --enable-dp-attention
        if "--enable-dp-attention" not in args:
            args = append_argument(args, "--enable-dp-attention")
841

842
843
844
845
        # 4. Set --ep-size=dep_size (expert parallelism size)
        args = set_argument_value(args, "--ep-size", str(dep_size))

        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
846
        return cfg.model_dump()
847
848
849

    @classmethod
    def get_model_name(cls, config: dict) -> str:
850
        cfg = Config.model_validate(config)
851
        try:
852
            worker_service = get_worker_service_from_config(cfg, backend="sglang")
853
854
            args = validate_and_get_worker_args(worker_service, backend="sglang")
        except (ValueError, KeyError):
855
            logger.warning(
856
                f"Worker service missing or invalid, using default model name: {DEFAULT_MODEL_NAME}"
857
858
            )
            return DEFAULT_MODEL_NAME
859
860
861
862
863
864
865
866
867
868
869
870
871

        args = break_arguments(args)
        for i, arg in enumerate(args):
            if arg == "--served-model-name" and i + 1 < len(args):
                return args[i + 1]

        logger.warning(
            f"Model name not found in configuration args, using default model name: {DEFAULT_MODEL_NAME}"
        )
        return DEFAULT_MODEL_NAME

    @classmethod
    def get_port(cls, config: dict) -> int:
872
        cfg = Config.model_validate(config)
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
        frontend_service = cfg.spec.services.get("Frontend")
        if (
            not frontend_service
            or not frontend_service.extraPodSpec
            or not frontend_service.extraPodSpec.mainContainer
        ):
            logger.warning(
                f"Frontend service or container not found, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

        args = frontend_service.extraPodSpec.mainContainer.args
        if not args:
            logger.warning(
                f"No args found in Frontend configuration, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

891
892
893
894
        args = break_arguments(args)
        try:
            idx = args.index("--http-port")
            return int(args[idx + 1])
895
        except (ValueError, IndexError):
896
897
898
899
900
901
            logger.warning(
                f"Port not found in configuration args, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

    @classmethod
902
903
904
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
905
906
907
908
909
910
911
        try:
            with open(dynamo_log_fn, "r") as f:
                for line in f:
                    if "KV Cache is allocated" in line and "#tokens:" in line:
                        # Extract the number after "#tokens:"
                        match = re.search(r"#tokens:\s*(\d+)", line)
                        if match:
912
                            return int(match.group(1)) * attention_dp_size
913
914
915
916
917
        except Exception as e:
            logger.warning(f"Failed to parse KV cache size from log file. Error: {e}")
        return 0


918
919
class TrtllmConfigModifier:
    @classmethod
920
921
922
923
924
925
926
927
928
929
930
    def convert_config(
        cls,
        config: dict,
        target: Literal["prefill", "decode"],
        is_moe_model: bool = False,
    ) -> dict:
        if is_moe_model:
            raise NotImplementedError(
                "MoE model support is not implemented for TrtLLM backend"
            )

931
932
933
934
935
936
937
938
939
940
        cfg = Config.model_validate(config)

        # set metadata name
        cfg.metadata.name = "trtllm-agg"

        # disable planner
        if "Planner" in cfg.spec.services:
            del cfg.spec.services["Planner"]

        if target == "prefill":
941
942
            # Get service names by inferring from subComponentType first
            prefill_service_name = get_service_name_by_type(
943
                cfg, "trtllm", SubComponentType.PREFILL
944
945
            )
            decode_service_name = get_service_name_by_type(
946
                cfg, "trtllm", SubComponentType.DECODE
947
948
            )

949
            # Convert to prefill-only aggregated setup
950
951
952
953
954
955
956
957
            # Rename prefill worker to decode worker name
            cfg.spec.services[decode_service_name] = cfg.spec.services[
                prefill_service_name
            ]
            del cfg.spec.services[prefill_service_name]

            # Set subComponentType for aggregated mode (using decode worker for prefill-only)
            cfg.spec.services[decode_service_name].subComponentType = "decode"
958

959
            worker_service = get_worker_service_from_config(
960
                cfg,
961
962
963
964
                backend="trtllm",
                sub_component_type=SubComponentType.DECODE,
            )
            args = validate_and_get_worker_args(worker_service, backend="trtllm")
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
            args = break_arguments(args)

            # Remove disaggregation args
            args = remove_valued_arguments(args, "--disaggregation-mode")
            args = remove_valued_arguments(args, "--disaggregation-strategy")

            # Keep the original extra-engine-args (prefill.yaml) which may contain user settings
            # Check if user already has override-engine-args and merge with our changes
            override_dict, args = parse_override_engine_args(args)

            # Merge our overrides for converting prefill-only disagg to aggregated:
            # - Disable enable_block_reuse (no KV reuse for prefill-only)
            # - Enable overlap scheduler (disabled in prefill.yaml but needed for agg)
            # - Remove cache_transceiver_config (not needed in agg mode)
            if "kv_cache_config" not in override_dict:
                override_dict["kv_cache_config"] = {}
            override_dict["kv_cache_config"]["enable_block_reuse"] = False
            override_dict[
                "disable_overlap_scheduler"
            ] = False  # Enable overlap scheduler for agg
            override_dict[
                "cache_transceiver_config"
            ] = None  # Remove cache transceiver for agg

            override_str = json.dumps(override_dict)
            args = append_argument(args, ["--override-engine-args", override_str])

            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)

        elif target == "decode":
995
996
            # Get service names by inferring from subComponentType first
            prefill_service_name = get_service_name_by_type(
997
                cfg, "trtllm", SubComponentType.PREFILL
998
999
            )
            decode_service_name = get_service_name_by_type(
1000
                cfg, "trtllm", SubComponentType.DECODE
1001
            )
1002

1003
            # Convert to decode-only aggregated setup
1004
            # Remove prefill worker if exists
1005
            del cfg.spec.services[prefill_service_name]
1006

1007
1008
1009
1010
1011
            # Set subComponentType for aggregated decode-only mode
            cfg.spec.services[decode_service_name].subComponentType = "decode"

            # Decode worker already has the correct name
            worker_service = get_worker_service_from_config(
1012
                cfg,
1013
1014
1015
1016
                backend="trtllm",
                sub_component_type=SubComponentType.DECODE,
            )
            args = validate_and_get_worker_args(worker_service, backend="trtllm")
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
            args = break_arguments(args)

            # Remove disaggregation args
            args = remove_valued_arguments(args, "--disaggregation-mode")
            args = remove_valued_arguments(args, "--disaggregation-strategy")

            # Keep the original extra-engine-args (decode.yaml) which may contain user settings
            # Check if user already has override-engine-args and merge with our changes
            override_dict, args = parse_override_engine_args(args)

            # Merge our overrides for converting decode-only disagg to aggregated:
            # - Enable enable_block_reuse (to skip prefill in decode-only)
            # - Remove cache_transceiver_config (not needed in agg mode)
            if "kv_cache_config" not in override_dict:
                override_dict["kv_cache_config"] = {}
            override_dict["kv_cache_config"]["enable_block_reuse"] = True
            override_dict[
                "cache_transceiver_config"
            ] = None  # Remove cache transceiver for agg

            override_str = json.dumps(override_dict)
            args = append_argument(args, ["--override-engine-args", override_str])

            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)

        # Set num workers to 1
1043
1044
        # Use the inferred decode service name
        final_decode_service_name = get_service_name_by_type(
1045
            cfg, "trtllm", SubComponentType.DECODE
1046
1047
        )
        worker_config = cfg.spec.services[final_decode_service_name]
1048
1049
1050
1051
1052
        worker_config.replicas = 1

        return cfg.model_dump()

    @classmethod
1053
1054
1055
1056
1057
1058
    def set_config_tp_size(
        cls,
        config: dict,
        tp_size: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
1059
1060
        cfg = Config.model_validate(config)

1061
1062
        # Get the worker service using helper function
        # This assumes convert_config has been called, so the service is named decode_worker_k8s_name
1063
        worker_service = get_worker_service_from_config(
1064
            cfg, backend="trtllm", sub_component_type=component_type
1065
        )
1066

1067
1068
        # Set up resources
        setup_worker_service_resources(worker_service, tp_size)
1069

1070
1071
        # Validate and get args
        args = validate_and_get_worker_args(worker_service, backend="trtllm")
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088

        # Break arguments to handle both joined strings and lists
        args = break_arguments(args)

        # For TRT-LLM, we need to update the override-engine-args
        # to set the tensor_parallel_size
        override_dict, args = parse_override_engine_args(args)

        # Add/update tensor_parallel_size in the override
        override_dict["tensor_parallel_size"] = tp_size
        override_str = json.dumps(override_dict)
        args = append_argument(args, ["--override-engine-args", override_str])

        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)

        return cfg.model_dump()

1089
    @classmethod
1090
1091
1092
1093
1094
1095
1096
    def set_config_tep_size(
        cls,
        config: dict,
        tep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
1097
1098
1099
1100
1101
        raise NotImplementedError(
            "TEP (Tensor Expert Parallelism) is not implemented for TrtLLM backend"
        )

    @classmethod
1102
1103
1104
1105
1106
1107
1108
    def set_config_dep_size(
        cls,
        config: dict,
        dep_size: int,
        num_gpus_per_node: int,
        component_type: SubComponentType = SubComponentType.DECODE,
    ):
1109
1110
1111
1112
        raise NotImplementedError(
            "DEP (Data Expert Parallelism) is not implemented for TrtLLM backend"
        )

1113
1114
    @classmethod
    def get_model_name(cls, config: dict) -> str:
1115
        cfg = Config.model_validate(config)
1116
        try:
1117
            worker_service = get_worker_service_from_config(cfg, backend="trtllm")
1118
1119
            args = validate_and_get_worker_args(worker_service, backend="trtllm")
        except (ValueError, KeyError):
1120
            logger.warning(
1121
                f"Worker service missing or invalid, using default model name: {DEFAULT_MODEL_NAME}"
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
            )
            return DEFAULT_MODEL_NAME

        args = break_arguments(args)
        for i, arg in enumerate(args):
            if arg == "--served-model-name" and i + 1 < len(args):
                return args[i + 1]

        logger.warning(
            f"Model name not found in configuration args, using default model name: {DEFAULT_MODEL_NAME}"
        )
        return DEFAULT_MODEL_NAME

    @classmethod
    def get_port(cls, config: dict) -> int:
        cfg = Config.model_validate(config)
        frontend_service = cfg.spec.services.get("Frontend")
        if (
            not frontend_service
            or not frontend_service.extraPodSpec
            or not frontend_service.extraPodSpec.mainContainer
        ):
            logger.warning(
                f"Frontend service or container not found, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

        # TRT-LLM frontend doesn't have args, it uses the default port
        return DYNAMO_RUN_DEFAULT_PORT

    @classmethod
1153
1154
1155
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
        # TRT-LLM log parsing for KV cache size
        # Format: [TensorRT-LLM][INFO] [MemUsageChange] Allocated XX GiB for max tokens in paged KV cache (XXXXXX).
        try:
            with open(dynamo_log_fn, "r") as f:
                for line in f:
                    # Look for the specific TRT-LLM KV cache allocation log
                    if (
                        "Allocated" in line
                        and "for max tokens in paged KV cache" in line
                    ):
                        # Extract the number in parentheses at the end
                        match = re.search(r"paged KV cache \((\d+)\)", line)
                        if match:
                            max_tokens = int(match.group(1))
                            logger.info(
                                f"Found TRT-LLM KV cache max tokens: {max_tokens}"
                            )
                            return max_tokens
        except Exception as e:
            logger.warning(f"Failed to parse KV cache size from log file. Error: {e}")

        # Return a reasonable default if we couldn't find the KV cache size in logs
        logger.warning(
            "Could not find KV cache size in TRT-LLM logs, using default value of 100000"
        )
        return 100000  # Default fallback value for TRT-LLM


1184
CONFIG_MODIFIERS: dict[str, type[ConfigModifierProtocol]] = {
1185
    "vllm": VllmV1ConfigModifier,
1186
    "sglang": SGLangConfigModifier,
1187
    "trtllm": TrtllmConfigModifier,
1188
}
1189

1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302

def generate_dgd_config_with_planner(
    config_path: str,
    config_modifier,
    best_prefill_gpus: int,
    best_decode_gpus: int,
    output_dir: str,
    args,
    is_moe_model: bool = False,
    num_gpus_per_node: int = 8,
):
    """Generate DGD config with planner based on profiling results.

    Args:
        config_path: Path to the YAML config file
        config_modifier: Config modifier instance (e.g., SGLangConfigModifier)
        best_prefill_gpus: Number of GPUs for prefill engine
        best_decode_gpus: Number of GPUs for decode engine
        output_dir: Output directory for profile results
        args: Parsed arguments namespace from profile_sla
        is_moe_model: Whether this is an MoE model
        num_gpus_per_node: Number of GPUs per node (for MoE models)

    Returns:
        dict: Final DGD config with planner service configured
    """

    # Load config from file
    with open(config_path, "r") as f:
        config = yaml.safe_load(f)

    if not is_moe_model:
        # dense model, use TP for both prefill and decode
        config = config_modifier.set_config_tp_size(
            config, best_prefill_gpus, SubComponentType.PREFILL
        )
        config = config_modifier.set_config_tp_size(
            config, best_decode_gpus, SubComponentType.DECODE
        )
    else:
        # MoE model, use TEP for prefill and DEP for decode
        config = config_modifier.set_config_tep_size(
            config,
            best_prefill_gpus,
            num_gpus_per_node,
            SubComponentType.PREFILL,
        )
        config = config_modifier.set_config_dep_size(
            config,
            best_decode_gpus,
            num_gpus_per_node,
            SubComponentType.DECODE,
        )
    config = Config.model_validate(config)

    # add PVC config if not present
    if not config.spec.pvcs:
        config.spec.pvcs = [PVCConfig()]

    # add the planner service
    planner_config = DgdPlannerServiceConfig()
    frontend_service = config.spec.services["Frontend"]
    planner_config.dynamoNamespace = getattr(frontend_service, "dynamoNamespace", "dynamo")  # type: ignore[attr-defined]
    if frontend_service.extraPodSpec and frontend_service.extraPodSpec.mainContainer:
        frontend_image = frontend_service.extraPodSpec.mainContainer.image
        if frontend_image and planner_config.extraPodSpec.mainContainer:
            planner_config.extraPodSpec.mainContainer.image = frontend_image

    # Build planner args dynamically from parsed arguments
    # This includes shared args (ttft, itl, backend, namespace) from profile_sla
    # and planner-specific args (with planner_ prefix)
    planner_args = build_planner_args_from_namespace(args, prefix="planner_")

    # Override profiling-specific arguments with results from profiling
    # Remove and re-add to ensure correct values from profiling context
    planner_args = [
        arg
        for arg in planner_args
        if not any(
            arg.startswith(f"--{key}=")
            for key in [
                "namespace",
                "prefill-engine-num-gpu",
                "decode-engine-num-gpu",
                "profile-results-dir",
            ]
        )
    ]

    # Add arguments determined by profiling results
    frontend_namespace = getattr(config.spec.services["Frontend"], "dynamoNamespace", "dynamo")  # type: ignore[attr-defined]
    planner_args.extend(
        [
            f"--namespace={frontend_namespace}",
            f"--prefill-engine-num-gpu={best_prefill_gpus}",
            f"--decode-engine-num-gpu={best_decode_gpus}",
            f"--profile-results-dir={output_dir}",
        ]
    )

    if (
        planner_config.extraPodSpec.mainContainer
        and planner_config.extraPodSpec.mainContainer.args is not None
    ):
        planner_config.extraPodSpec.mainContainer.args.extend(planner_args)
    # Convert planner config to dict first, then the entire config to dict
    planner_dict = planner_config.model_dump(exclude_unset=False)
    config_dict = config.model_dump(exclude_unset=False)
    config_dict["spec"]["services"]["Planner"] = planner_dict

    return config_dict


1303
1304
# Re-export WORKER_COMPONENT_NAMES for profile_sla.py
__all__ = ["CONFIG_MODIFIERS", "WORKER_COMPONENT_NAMES"]