config.py 37.3 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
from pydantic import BaseModel
24

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

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)


42
class Container(BaseModel):
43
44
    args: Optional[list[str]] = None
    model_config = {"extra": "allow"}
45
46
47


class PodSpec(BaseModel):
48
49
    mainContainer: Optional[Container] = None
    model_config = {"extra": "allow"}
50
51
52


class ServiceResources(BaseModel):
53
    requests: Optional[dict[str, str]] = None
54
55
56
57
    limits: Optional[dict[str, str]] = None


class Service(BaseModel):
58
59
60
61
    replicas: Optional[int] = None
    resources: Optional[ServiceResources] = None
    extraPodSpec: Optional[PodSpec] = None
    model_config = {"extra": "allow"}
62
63
64
65


class Services(BaseModel):
    Frontend: Service
66
    model_config = {"extra": "allow"}
67
68
69
70
71
72
73
74
75
76
77
78
79


class Spec(BaseModel):
    services: dict[str, Service]


class Metadata(BaseModel):
    name: str


class Config(BaseModel):
    metadata: Metadata
    spec: Spec
80
    model_config = {"extra": "allow"}
81
82


83
84
85
86
class MultinodeConfig(BaseModel):
    nodeCount: int


87
88
89
90
def break_arguments(args: list[str] | None) -> list[str]:
    ans: list[str] = []
    if args is None:
        return ans
91
    if isinstance(args, str):
92
93
        # Use shlex.split to properly handle quoted arguments and JSON values
        ans = shlex.split(args)
94
95
    else:
        for arg in args:
96
            if arg is not None:
97
98
                # Use shlex.split to properly handle quoted arguments
                ans.extend(shlex.split(arg))
99
    return ans
100
101


102
103
104
105
106
107
108
109
110
111
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


112
def join_arguments(args: list[str]) -> list[str]:
113
114
    # Use shlex.join to properly quote arguments that contain spaces or special characters
    return [shlex.join(args)]
115
116


117
118
119
120
121
122
123
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
124
125


126
127
128
def find_arg_index(args: list[str]) -> int:
    # find the correct index to insert an argument
    idx = len(args)
129

130
131
132
133
134
    try:
        new_idx = args.index("|")
        idx = min(idx, new_idx)
    except ValueError:
        pass
135

136
137
138
139
140
    try:
        new_idx = args.index("2>&1")
        idx = min(idx, new_idx)
    except ValueError:
        pass
141

142
    return idx
143
144


145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
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


167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
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


# TODO: make is work for all frameworks
def get_worker_service_from_config(config: dict):
    """Helper function to get the SGLang decode worker service from config."""
    cfg = Config.model_validate(config)
    return cfg.spec.services[WORKER_COMPONENT_NAMES["sglang"].decode_worker_k8s_name]


# TODO: make is work for all frameworks
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)


# TODO: make is work for all frameworks
def validate_and_get_worker_args(worker_service):
    """Helper function to validate worker service and get its arguments."""
    if not worker_service.extraPodSpec or not worker_service.extraPodSpec.mainContainer:
        raise ValueError(
            f"Missing extraPodSpec or mainContainer in SGLang decode worker service '{WORKER_COMPONENT_NAMES['sglang'].decode_worker_k8s_name}'"
        )

    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


249
250
class ConfigModifierProtocol(Protocol):
    @classmethod
251
252
253
254
255
256
    def convert_config(
        cls,
        config: dict,
        target: Literal["prefill", "decode"],
        is_moe_model: bool = False,
    ) -> dict:
257
258
259
260
261
262
        ...

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

263
264
265
266
267
268
269
270
271
272
273
274
    @classmethod
    def set_config_tep_size(
        cls, config: dict, tep_size: int, num_gpus_per_node: int
    ) -> dict:
        ...

    @classmethod
    def set_config_dep_size(
        cls, config: dict, dep_size: int, num_gpus_per_node: int
    ) -> dict:
        ...

275
276
277
278
279
280
281
282
283
    @classmethod
    def get_model_name(cls, config: dict) -> str:
        ...

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

    @classmethod
284
285
286
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
287
288
289
        ...


290
291
class VllmV1ConfigModifier:
    @classmethod
292
293
294
295
296
297
298
299
300
301
302
    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"
            )

303
        cfg = Config.model_validate(config)
304

305
        # set metadata name
306
        cfg.metadata.name = "vllm-agg"
307

308
        # disable planner
309
310
        if "Planner" in cfg.spec.services:
            del cfg.spec.services["Planner"]
311
312

        if target == "prefill":
313
            # convert prefill worker into decode worker
314
            cfg.spec.services[
315
                WORKER_COMPONENT_NAMES["vllm"].decode_worker_k8s_name
316
            ] = cfg.spec.services[
317
                WORKER_COMPONENT_NAMES["vllm"].prefill_worker_k8s_name
318
            ]
319
            del cfg.spec.services[
320
                WORKER_COMPONENT_NAMES["vllm"].prefill_worker_k8s_name
321
322
            ]

323
            worker_service = cfg.spec.services[
324
                WORKER_COMPONENT_NAMES["vllm"].decode_worker_k8s_name
325
326
327
328
329
330
            ]
            if (
                not worker_service.extraPodSpec
                or not worker_service.extraPodSpec.mainContainer
            ):
                raise ValueError(
331
                    f"Missing extraPodSpec or mainContainer in VLLM decode worker service '{WORKER_COMPONENT_NAMES['vllm'].decode_worker_k8s_name}'"
332
333
                )
            args = worker_service.extraPodSpec.mainContainer.args
334
335
336
337
338
339
340
341
342
343
344
345

            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")

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

348
        elif target == "decode":
349
            # delete prefill worker
350
            del cfg.spec.services[
351
                WORKER_COMPONENT_NAMES["vllm"].prefill_worker_k8s_name
352
353
            ]

354
            worker_service = cfg.spec.services[
355
                WORKER_COMPONENT_NAMES["vllm"].decode_worker_k8s_name
356
357
358
359
360
361
            ]
            if (
                not worker_service.extraPodSpec
                or not worker_service.extraPodSpec.mainContainer
            ):
                raise ValueError(
362
                    f"Missing extraPodSpec or mainContainer in VLLM decode worker service '{WORKER_COMPONENT_NAMES['vllm'].decode_worker_k8s_name}'"
363
364
                )
            args = worker_service.extraPodSpec.mainContainer.args
365
366

            args = break_arguments(args)
367

368
369
370
371
372
373
            # 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")

374
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
375
376

        # set num workers to 1
377
        decode_worker_config = cfg.spec.services[
378
            WORKER_COMPONENT_NAMES["vllm"].decode_worker_k8s_name
379
        ]
380
        decode_worker_config.replicas = 1
381

382
        return cfg.model_dump()
383
384
385

    @classmethod
    def set_config_tp_size(cls, config: dict, tp_size: int):
386
        cfg = Config.model_validate(config)
387

388
        worker_service = cfg.spec.services[
389
            WORKER_COMPONENT_NAMES["vllm"].decode_worker_k8s_name
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
        ]

        # 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 = {}

        worker_service.resources.requests["gpu"] = str(tp_size)

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

406
        if (
407
408
            not worker_service.extraPodSpec
            or not worker_service.extraPodSpec.mainContainer
409
        ):
410
411
412
            raise ValueError(
                f"Missing extraPodSpec or mainContainer in VLLM decode worker service '{WORKER_COMPONENT_NAMES['vllm'].decode_worker_k8s_name}'"
            )
413
        args = worker_service.extraPodSpec.mainContainer.args
414
415
416
417
418
419
420
421
422

        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)])

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

425
        return cfg.model_dump()
426

427
428
429
430
431
432
433
434
435
436
437
438
    @classmethod
    def set_config_tep_size(cls, config: dict, tep_size: int, num_gpus_per_node: int):
        raise NotImplementedError(
            "TEP (Tensor Expert Parallelism) is not implemented for VLLM backend"
        )

    @classmethod
    def set_config_dep_size(cls, config: dict, dep_size: int, num_gpus_per_node: int):
        raise NotImplementedError(
            "DEP (Data Expert Parallelism) is not implemented for VLLM backend"
        )

439
440
    @classmethod
    def get_model_name(cls, config: dict) -> str:
441
        cfg = Config.model_validate(config)
442
        worker_name = WORKER_COMPONENT_NAMES["vllm"].decode_worker_k8s_name
443
444
445
446
447
448
449
450
451
452
        worker_service = cfg.spec.services[worker_name]
        if (
            not worker_service.extraPodSpec
            or not worker_service.extraPodSpec.mainContainer
        ):
            logger.warning(
                f"Worker service missing extraPodSpec or mainContainer, using default model name: {DEFAULT_MODEL_NAME}"
            )
            return DEFAULT_MODEL_NAME
        args = worker_service.extraPodSpec.mainContainer.args
453
454
455
456
457
458
459
460
461
462

        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
463
464
465

    @classmethod
    def get_port(cls, config: dict) -> int:
466
        cfg = Config.model_validate(config)
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
        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

485
486
487
488
        args = break_arguments(args)
        try:
            idx = args.index("--http-port")
            return int(args[idx + 1])
489
        except (ValueError, IndexError):
490
491
492
493
            logger.warning(
                f"Port not found in configuration args, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT
494
495

    @classmethod
496
497
498
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
        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


520
521
class SGLangConfigModifier:
    @classmethod
522
523
524
525
526
527
    def convert_config(
        cls,
        config: dict,
        target: Literal["prefill", "decode"],
        is_moe_model: bool = False,
    ) -> dict:
528
        cfg = Config.model_validate(config)
529
530

        # set metadata name
531
        cfg.metadata.name = "sglang-agg"
532
533

        # disable planner
534
535
        if "Planner" in cfg.spec.services:
            del cfg.spec.services["Planner"]
536
537
538

        if target == "prefill":
            # convert prefill worker into decode worker
539
            cfg.spec.services[
540
                WORKER_COMPONENT_NAMES["sglang"].decode_worker_k8s_name
541
            ] = cfg.spec.services[
542
543
                WORKER_COMPONENT_NAMES["sglang"].prefill_worker_k8s_name
            ]
544
            del cfg.spec.services[
545
546
547
                WORKER_COMPONENT_NAMES["sglang"].prefill_worker_k8s_name
            ]

548
            worker_service = cfg.spec.services[
549
                WORKER_COMPONENT_NAMES["sglang"].decode_worker_k8s_name
550
551
552
553
554
555
            ]
            if (
                not worker_service.extraPodSpec
                or not worker_service.extraPodSpec.mainContainer
            ):
                raise ValueError(
556
                    f"Missing extraPodSpec or mainContainer in SGLang decode worker service '{WORKER_COMPONENT_NAMES['sglang'].decode_worker_k8s_name}'"
557
558
                )
            args = worker_service.extraPodSpec.mainContainer.args
559
560
561

            args = break_arguments(args)

562
            # remove disagg flags
563
564
            args = remove_valued_arguments(args, "--disaggregation-mode")
            args = remove_valued_arguments(args, "--disaggregation-transfer-backend")
565
            args = remove_valued_arguments(args, "--disaggregation-bootstrap-port")
566
567
568
569
570

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

571
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
572
573
574

        elif target == "decode":
            # delete prefill worker
575
            del cfg.spec.services[
576
577
578
                WORKER_COMPONENT_NAMES["sglang"].prefill_worker_k8s_name
            ]

579
            worker_service = cfg.spec.services[
580
                WORKER_COMPONENT_NAMES["sglang"].decode_worker_k8s_name
581
582
583
584
585
586
            ]
            if (
                not worker_service.extraPodSpec
                or not worker_service.extraPodSpec.mainContainer
            ):
                raise ValueError(
587
                    f"Missing extraPodSpec or mainContainer in SGLang decode worker service '{WORKER_COMPONENT_NAMES['sglang'].decode_worker_k8s_name}'"
588
589
                )
            args = worker_service.extraPodSpec.mainContainer.args
590
591
592

            args = break_arguments(args)

593
            # remove disagg flags
594
595
            args = remove_valued_arguments(args, "--disaggregation-mode")
            args = remove_valued_arguments(args, "--disaggregation-transfer-backend")
596
            args = remove_valued_arguments(args, "--disaggregation-bootstrap-port")
597
598
599
600
601

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

602
603
604
605
606
607
608
609
610
611
            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"]
                    )

612
            worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
613
614
615
616
617
618
619

        # set num workers to 1
        decode_worker_config = config["spec"]["services"][
            WORKER_COMPONENT_NAMES["sglang"].decode_worker_k8s_name
        ]
        decode_worker_config["replicas"] = 1

620
        return cfg.model_dump()
621
622
623

    @classmethod
    def set_config_tp_size(cls, config: dict, tp_size: int):
624
        cfg = Config.model_validate(config)
625
        worker_service = get_worker_service_from_config(config)
626

627
628
        # Set up resources
        setup_worker_service_resources(worker_service, tp_size)
629

630
631
        # Get and validate args
        args = validate_and_get_worker_args(worker_service)
632

633
634
        # Set --tp argument
        args = set_argument_value(args, "--tp", str(tp_size))
635

636
637
        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
        return cfg.model_dump()
638

639
640
641
642
    @classmethod
    def set_config_tep_size(cls, config: dict, tep_size: int, num_gpus_per_node: int):
        cfg = Config.model_validate(config)
        worker_service = get_worker_service_from_config(config)
643

644
645
        # Set up resources with multinode configuration
        setup_worker_service_resources(worker_service, tep_size, num_gpus_per_node)
646

647
648
        # Get and validate args
        args = validate_and_get_worker_args(worker_service)
649

650
651
652
653
654
655
656
657
658
659
660
661
        # 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")
662

663
        worker_service.extraPodSpec.mainContainer.args = join_arguments(args)
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
        return cfg.model_dump()

    @classmethod
    def set_config_dep_size(cls, config: dict, dep_size: int, num_gpus_per_node: int):
        cfg = Config.model_validate(config)
        worker_service = get_worker_service_from_config(config)

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

        # Get and validate args
        args = validate_and_get_worker_args(worker_service)

        # 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")
686

687
688
689
690
        # 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)
691
        return cfg.model_dump()
692
693
694

    @classmethod
    def get_model_name(cls, config: dict) -> str:
695
        cfg = Config.model_validate(config)
696
        worker_name = WORKER_COMPONENT_NAMES["sglang"].decode_worker_k8s_name
697
698
699
700
701
702
703
704
705
706
        worker_service = cfg.spec.services[worker_name]
        if (
            not worker_service.extraPodSpec
            or not worker_service.extraPodSpec.mainContainer
        ):
            logger.warning(
                f"Worker service missing extraPodSpec or mainContainer, using default model name: {DEFAULT_MODEL_NAME}"
            )
            return DEFAULT_MODEL_NAME
        args = worker_service.extraPodSpec.mainContainer.args
707
708
709
710
711
712
713
714
715
716
717
718
719

        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:
720
        cfg = Config.model_validate(config)
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
        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

739
740
741
742
        args = break_arguments(args)
        try:
            idx = args.index("--http-port")
            return int(args[idx + 1])
743
        except (ValueError, IndexError):
744
745
746
747
748
749
            logger.warning(
                f"Port not found in configuration args, using default port: {DYNAMO_RUN_DEFAULT_PORT}"
            )
            return DYNAMO_RUN_DEFAULT_PORT

    @classmethod
750
751
752
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
753
754
755
756
757
758
759
        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:
760
                            return int(match.group(1)) * attention_dp_size
761
762
763
764
765
        except Exception as e:
            logger.warning(f"Failed to parse KV cache size from log file. Error: {e}")
        return 0


766
767
class TrtllmConfigModifier:
    @classmethod
768
769
770
771
772
773
774
775
776
777
778
    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"
            )

779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
        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":
            # Convert to prefill-only aggregated setup
            # Merge prefill worker config into a single worker
            if "TRTLLMPrefillWorker" in cfg.spec.services:
                # Rename prefill worker to generic worker
                cfg.spec.services["TRTLLMWorker"] = cfg.spec.services[
                    "TRTLLMPrefillWorker"
                ]
                del cfg.spec.services["TRTLLMPrefillWorker"]

            # Remove decode worker
            del cfg.spec.services["TRTLLMDecodeWorker"]

            worker_service = cfg.spec.services["TRTLLMWorker"]
            if (
                not worker_service.extraPodSpec
                or not worker_service.extraPodSpec.mainContainer
            ):
                raise ValueError(
                    "Missing extraPodSpec or mainContainer in TRTLLM worker service 'TRTLLMWorker'"
                )
            args = worker_service.extraPodSpec.mainContainer.args

            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":
            # Convert to decode-only aggregated setup
            # Use decode worker as the main worker
            if "TRTLLMDecodeWorker" in cfg.spec.services:
                # Rename decode worker to generic worker
                cfg.spec.services["TRTLLMWorker"] = cfg.spec.services[
                    "TRTLLMDecodeWorker"
                ]
                del cfg.spec.services["TRTLLMDecodeWorker"]

            # Remove prefill worker if exists
            if "TRTLLMPrefillWorker" in cfg.spec.services:
                del cfg.spec.services["TRTLLMPrefillWorker"]

            worker_service = cfg.spec.services["TRTLLMWorker"]
            if (
                not worker_service.extraPodSpec
                or not worker_service.extraPodSpec.mainContainer
            ):
                raise ValueError(
                    "Missing extraPodSpec or mainContainer in TRTLLM worker service 'TRTLLMWorker'"
                )
            args = worker_service.extraPodSpec.mainContainer.args

            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
        worker_config = cfg.spec.services["TRTLLMWorker"]
        worker_config.replicas = 1

        return cfg.model_dump()

    @classmethod
    def set_config_tp_size(cls, config: dict, tp_size: int):
        cfg = Config.model_validate(config)

        worker_service = cfg.spec.services["TRTLLMWorker"]

        # 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 = {}

        worker_service.resources.requests["gpu"] = str(tp_size)

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

        if (
            not worker_service.extraPodSpec
            or not worker_service.extraPodSpec.mainContainer
        ):
            raise ValueError(
                "Missing extraPodSpec or mainContainer in TRTLLM worker service 'TRTLLMWorker'"
            )
        args = worker_service.extraPodSpec.mainContainer.args

        # 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()

940
941
942
943
944
945
946
947
948
949
950
951
    @classmethod
    def set_config_tep_size(cls, config: dict, tep_size: int, num_gpus_per_node: int):
        raise NotImplementedError(
            "TEP (Tensor Expert Parallelism) is not implemented for TrtLLM backend"
        )

    @classmethod
    def set_config_dep_size(cls, config: dict, dep_size: int, num_gpus_per_node: int):
        raise NotImplementedError(
            "DEP (Data Expert Parallelism) is not implemented for TrtLLM backend"
        )

952
953
954
955
956
957
958
959
960
961
962
963
964
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
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
    @classmethod
    def get_model_name(cls, config: dict) -> str:
        cfg = Config.model_validate(config)
        worker_name = "TRTLLMWorker"
        worker_service = cfg.spec.services.get(worker_name)

        # Also check for disagg worker names
        if not worker_service:
            worker_name = "TRTLLMPrefillWorker"
            worker_service = cfg.spec.services.get(worker_name)
        if not worker_service:
            worker_name = "TRTLLMDecodeWorker"
            worker_service = cfg.spec.services.get(worker_name)

        if not worker_service:
            logger.warning(
                f"Worker service not found, using default model name: {DEFAULT_MODEL_NAME}"
            )
            return DEFAULT_MODEL_NAME

        if (
            not worker_service.extraPodSpec
            or not worker_service.extraPodSpec.mainContainer
        ):
            logger.warning(
                f"Worker service missing extraPodSpec or mainContainer, using default model name: {DEFAULT_MODEL_NAME}"
            )
            return DEFAULT_MODEL_NAME
        args = worker_service.extraPodSpec.mainContainer.args

        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
1010
1011
1012
    def get_kv_cache_size_from_dynamo_log(
        cls, dynamo_log_fn: str, attention_dp_size: int = 1
    ) -> int:
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
        # 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


1041
CONFIG_MODIFIERS: dict[str, type[ConfigModifierProtocol]] = {
1042
    "vllm": VllmV1ConfigModifier,
1043
    "sglang": SGLangConfigModifier,
1044
    "trtllm": TrtllmConfigModifier,
1045
}
1046
1047
1048

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