truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/__init__.py

@@ -1,9 +1,12 @@
-# flake8: noqa F401
-
+import warnings
 from pathlib import Path
 
+from pydantic import PydanticDeprecatedSince20
 from single_source import get_version
 
+# Suppress Pydantic V1 warnings, because we have to use it for backwards compat.
+warnings.filterwarnings("ignore", category=PydanticDeprecatedSince20)
+
 __version__ = get_version(__name__, Path(__file__).parent.parent)
 
 
@@ -11,4 +14,8 @@ def version():
     return __version__
 
 
-from truss.build import from_directory, init, kill_all, load
+from truss.api import login, push, whoami
+from truss.base import truss_config
+from truss.truss_handle.build import load  # TODO: Refactor all usages and remove.
+
+__all__ = ["push", "login", "load", "whoami", "truss_config"]

truss/api/__init__.py

@@ -0,0 +1,123 @@
+from typing import TYPE_CHECKING, Optional, Type, cast
+
+if TYPE_CHECKING:
+    from rich import progress
+
+from truss.api import definitions
+from truss.remote.baseten.service import BasetenService
+from truss.remote.remote_factory import RemoteFactory
+from truss.remote.truss_remote import RemoteConfig
+from truss.truss_handle.build import load
+
+
+def login(api_key: str):
+    """
+    Logs user into Baseten account. Persists information to ~/.trussrc file,
+    so only needs to be invoked once.
+    Args:
+        api_key: Baseten API Key
+    """
+    remote_url = "https://app.baseten.co"
+    remote_config = RemoteConfig(
+        name="baseten",
+        configs={
+            "remote_provider": "baseten",
+            "api_key": api_key,
+            "remote_url": remote_url,
+        },
+    )
+    RemoteFactory.update_remote_config(remote_config)
+
+
+def whoami(remote: Optional[str] = None):
+    """
+    Returns account information for the current user.
+    """
+    if not remote:
+        available_remotes = RemoteFactory.get_available_config_names()
+        if len(available_remotes) == 1:
+            remote = available_remotes[0]
+        elif len(available_remotes) == 0:
+            raise ValueError(
+                "Please authenticate via truss.login and pass it as an argument."
+            )
+        else:
+            raise ValueError(
+                "Multiple remotes found. Please pass the remote as an argument."
+            )
+
+    remote_provider = RemoteFactory.create(remote=remote)
+    return remote_provider.whoami()
+
+
+def push(
+    target_directory: str,
+    remote: Optional[str] = None,
+    model_name: Optional[str] = None,
+    publish: bool = False,
+    promote: bool = False,
+    preserve_previous_production_deployment: bool = False,
+    trusted: bool = False,
+    deployment_name: Optional[str] = None,
+    environment: Optional[str] = None,
+    progress_bar: Optional[Type["progress.Progress"]] = None,
+) -> definitions.ModelDeployment:
+    """
+    Pushes a Truss to Baseten.
+
+    Args:
+        target_directory: Directory of Truss to push.
+        remote: Name of the remote in .trussrc to patch changes to.
+        model_name: The name of the model, if different from the one in the config.yaml.
+        publish: Push the truss as a published deployment. If no production deployment exists,
+            promote the truss to production after deploy completes.
+        promote: Push the truss as a published deployment. Even if a production deployment exists,
+            promote the truss to production after deploy completes.
+        preserve_previous_production_deployment: Preserve the previous production deployment’s autoscaling
+            setting. When not specified, the previous production deployment will be updated to allow it to
+            scale to zero. Can only be use in combination with `promote` option.
+        trusted: Give Truss access to secrets on remote host.
+        deployment_name: Name of the deployment created by the push. Can only be
+            used in combination with `publish` or `promote`. Deployment name must
+            only contain alphanumeric, ’.’, ’-’ or ’_’ characters.
+        environment: Name of stable environment on baseten.
+        progress_bar: Optional `rich.progress.Progress` if output is desired.
+
+    Returns:
+        The newly created ModelDeployment.
+    """
+
+    if not remote:
+        available_remotes = RemoteFactory.get_available_config_names()
+        if len(available_remotes) == 1:
+            remote = available_remotes[0]
+        elif len(available_remotes) == 0:
+            raise ValueError(
+                "Please authenticate via truss.login and pass it as an argument."
+            )
+        else:
+            raise ValueError(
+                "Multiple remotes found. Please pass the remote as an argument."
+            )
+
+    remote_provider = RemoteFactory.create(remote=remote)
+    tr = load(target_directory)
+    model_name = model_name or tr.spec.config.model_name
+    if not model_name:
+        raise ValueError(
+            "No model name provided. Please specify a model name in config.yaml."
+        )
+
+    service = remote_provider.push(
+        tr,
+        model_name=model_name,
+        publish=publish,
+        trusted=trusted,
+        promote=promote,
+        preserve_previous_prod_deployment=preserve_previous_production_deployment,
+        deployment_name=deployment_name,
+        environment=environment,
+        progress_bar=progress_bar,
+    )  # type: ignore
+
+    return definitions.ModelDeployment(cast(BasetenService, service))

truss/api/definitions.py

@@ -0,0 +1,51 @@
+import time
+
+import pydantic
+
+from truss.remote.baseten import service
+from truss.remote.baseten.core import ACTIVE_STATUS, DEPLOYING_STATUSES
+
+
+class ModelDeployment:
+    model_config = pydantic.ConfigDict(protected_namespaces=())
+
+    model_id: str
+    model_deployment_id: str
+    _baseten_service: service.BasetenService
+
+    def __init__(self, service: service.BasetenService):
+        self.model_id = service._model_id
+        self.model_deployment_id = service._model_version_id
+        self._baseten_service = service
+
+    def wait_for_active(self, timeout_seconds: int = 600) -> bool:
+        """
+        Waits for the deployment to be active.
+
+        Args:
+            timeout_seconds: The maximum time to wait for the deployment to be active.
+
+        Returns:
+            The status of the deployment.
+        """
+        start_time = time.time()
+        for deployment_status in self._baseten_service.poll_deployment_status():
+            if (
+                timeout_seconds is not None
+                and time.time() - start_time > timeout_seconds
+            ):
+                raise TimeoutError("Deployment timed out.")
+
+            if deployment_status == ACTIVE_STATUS:
+                return True
+
+            if deployment_status not in DEPLOYING_STATUSES:
+                raise Exception(f"Deployment failed with status: {deployment_status}")
+
+        raise RuntimeError("Error polling deployment status.")
+
+    def __repr__(self):
+        return (
+            f"ModelDeployment(model_id={self.model_id}, "
+            f"model_deployment_id={self.model_deployment_id})"
+        )

truss/base/constants.py

@@ -0,0 +1,116 @@
+import pathlib
+from typing import Set
+
+SKLEARN = "sklearn"
+TENSORFLOW = "tensorflow"
+KERAS = "keras"
+XGBOOST = "xgboost"
+PYTORCH = "pytorch"
+CUSTOM = "custom"
+HUGGINGFACE_TRANSFORMER = "huggingface_transformer"
+LIGHTGBM = "lightgbm"
+
+_TRUSS_ROOT = pathlib.Path(__file__).parent.parent.resolve()
+
+TEMPLATES_DIR = _TRUSS_ROOT / "templates"
+TRADITIONAL_CUSTOM_TEMPLATE_DIR = TEMPLATES_DIR / "custom"
+PYTHON_DX_CUSTOM_TEMPLATE_DIR = TEMPLATES_DIR / "custom_python_dx"
+DOCKER_SERVER_TEMPLATES_DIR = TEMPLATES_DIR / "docker_server"
+SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "server"
+TRITON_SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "triton"
+TRTLLM_TRUSS_DIR: pathlib.Path = TEMPLATES_DIR / "trtllm-briton"
+SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME = "shared"
+SHARED_SERVING_AND_TRAINING_CODE_DIR: pathlib.Path = (
+    TEMPLATES_DIR / SHARED_SERVING_AND_TRAINING_CODE_DIR_NAME
+)
+CONTROL_SERVER_CODE_DIR: pathlib.Path = TEMPLATES_DIR / "control"
+CHAINS_CODE_DIR: pathlib.Path = _TRUSS_ROOT.parent / "truss-chains" / "truss_chains"
+
+SUPPORTED_PYTHON_VERSIONS = {"3.8", "3.9", "3.10", "3.11"}
+MAX_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE = "3.12"
+MIN_SUPPORTED_PYTHON_VERSION_IN_CUSTOM_BASE_IMAGE = "3.8"
+
+TRTLLM_PREDICT_CONCURRENCY = 512
+BEI_TRTLLM_CLIENT_BATCH_SIZE = 128
+BEI_MAX_CONCURRENCY_TARGET_REQUESTS = 2048
+BEI_REQUIRED_MAX_NUM_TOKENS = 16384
+
+TRTLLM_MIN_MEMORY_REQUEST_GI = 16
+HF_MODELS_API_URL = "https://huggingface.co/api/models"
+HF_ACCESS_TOKEN_KEY = "hf_access_token"
+TRUSSLESS_MAX_PAYLOAD_SIZE = "64M"
+# Alias for TEMPLATES_DIR
+SERVING_DIR: pathlib.Path = TEMPLATES_DIR
+
+REQUIREMENTS_TXT_FILENAME = "requirements.txt"
+USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME = "user_requirements.txt"
+BASE_SERVER_REQUIREMENTS_TXT_FILENAME = "base_server_requirements.txt"
+SERVER_REQUIREMENTS_TXT_FILENAME = "server_requirements.txt"
+SYSTEM_PACKAGES_TXT_FILENAME = "system_packages.txt"
+
+FILENAME_CONSTANTS_MAP = {
+    "config_requirements_filename": REQUIREMENTS_TXT_FILENAME,
+    "user_supplied_requirements_filename": USER_SUPPLIED_REQUIREMENTS_TXT_FILENAME,
+    "base_server_requirements_filename": BASE_SERVER_REQUIREMENTS_TXT_FILENAME,
+    "server_requirements_filename": SERVER_REQUIREMENTS_TXT_FILENAME,
+    "system_packages_filename": SYSTEM_PACKAGES_TXT_FILENAME,
+}
+
+SERVER_DOCKERFILE_TEMPLATE_NAME = "server.Dockerfile.jinja"
+MODEL_DOCKERFILE_NAME = "Dockerfile"
+
+README_TEMPLATE_NAME = "README.md.jinja"
+MODEL_README_NAME = "README.md"
+
+CONFIG_FILE = "config.yaml"
+DOCKERFILE = "Dockerfile"
+# Used to indicate whether to associate a container with Truss
+TRUSS = "truss"
+# Used to create unique identifier based on last time truss was updated
+TRUSS_MODIFIED_TIME = "truss_modified_time"
+# Path of the Truss used to identify which Truss is being referred
+TRUSS_DIR = "truss_dir"
+TRUSS_HASH = "truss_hash"
+
+HUGGINGFACE_TRANSFORMER_MODULE_NAME: Set[str] = set({})
+
+# list from https://scikit-learn.org/stable/developers/advanced_installation.html
+SKLEARN_REQ_MODULE_NAMES: Set[str] = {
+    "numpy",
+    "scipy",
+    "joblib",
+    "scikit-learn",
+    "threadpoolctl",
+}
+
+XGBOOST_REQ_MODULE_NAMES: Set[str] = {"xgboost"}
+
+# list from https://www.tensorflow.org/install/pip
+# if problematic, lets look to https://www.tensorflow.org/install/source
+TENSORFLOW_REQ_MODULE_NAMES: Set[str] = {"tensorflow"}
+
+LIGHTGBM_REQ_MODULE_NAMES: Set[str] = {"lightgbm"}
+
+# list from https://pytorch.org/get-started/locally/
+PYTORCH_REQ_MODULE_NAMES: Set[str] = {"torch", "torchvision", "torchaudio"}
+
+MLFLOW_REQ_MODULE_NAMES: Set[str] = {"mlflow"}
+
+INFERENCE_SERVER_PORT = 8080
+
+HTTP_PUBLIC_BLOB_BACKEND = "http_public"
+
+REGISTRY_BUILD_SECRET_PREFIX = "DOCKER_REGISTRY_"
+
+TRTLLM_SPEC_DEC_TARGET_MODEL_NAME = "target"
+TRTLLM_SPEC_DEC_DRAFT_MODEL_NAME = "draft"
+TRTLLM_BASE_IMAGE = "baseten/briton-server:v0.16.0-5be7b58"
+TRTLLM_PYTHON_EXECUTABLE = "/usr/local/briton/venv/bin/python"
+BASE_TRTLLM_REQUIREMENTS = ["briton==0.4.2"]
+BEI_TRTLLM_BASE_IMAGE = "baseten/bei:0.0.17@sha256:9c3577f6ec672d6da5aca18e9c0ebdddd65ed80c8858e757fbde7e9cf48de01d"
+
+BEI_TRTLLM_PYTHON_EXECUTABLE = "/usr/bin/python3"
+
+OPENAI_COMPATIBLE_TAG = "openai-compatible"
+
+PRODUCTION_ENVIRONMENT_NAME = "production"

truss/base/custom_types.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+
+# TODO(marius/TaT): kill this.
+class ModelFrameworkType(Enum):
+    SKLEARN = "sklearn"
+    TENSORFLOW = "tensorflow"
+    KERAS = "keras"
+    PYTORCH = "pytorch"
+    HUGGINGFACE_TRANSFORMER = "huggingface_transformer"
+    XGBOOST = "xgboost"
+    LIGHTGBM = "lightgbm"
+    MLFLOW = "mlflow"
+    CUSTOM = "custom"
+
+
+@dataclass
+class Example:
+    name: str
+    input: Any
+
+    @staticmethod
+    def from_dict(example_dict):
+        return Example(name=example_dict["name"], input=example_dict["input"])
+
+    def to_dict(self) -> dict:
+        return {"name": self.name, "input": self.input}

truss/{errors.py → base/errors.py}

@@ -38,3 +38,7 @@ class ContainerNotFoundError(Error):
 
 class ContainerAPINoResponseError(Error):
     pass
+
+
+class RemoteNetworkError(Exception):
+    pass

truss/base/trt_llm_config.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import warnings
+from enum import Enum
+from typing import Any, Optional
+
+from huggingface_hub.errors import HFValidationError
+from huggingface_hub.utils import validate_repo_id
+from pydantic import BaseModel, PydanticDeprecatedSince20, model_validator, validator
+
+from truss.base.constants import BEI_REQUIRED_MAX_NUM_TOKENS
+
+logger = logging.getLogger(__name__)
+# Suppress Pydantic V1 warnings, because we have to use it for backwards compat.
+warnings.filterwarnings("ignore", category=PydanticDeprecatedSince20)
+
+ENGINE_BUILDER_TRUSS_RUNTIME_MIGRATION = (
+    os.environ.get("ENGINE_BUILDER_TRUSS_RUNTIME_MIGRATION", "False") == "True"
+)
+
+
+class TrussTRTLLMModel(str, Enum):
+    LLAMA = "llama"
+    MISTRAL = "mistral"
+    DEEPSEEK = "deepseek"
+    WHISPER = "whisper"
+    QWEN = "qwen"
+    ENCODER = "encoder"
+    PALMYRA = "palmyra"
+
+
+class TrussTRTLLMQuantizationType(str, Enum):
+    NO_QUANT = "no_quant"
+    WEIGHTS_ONLY_INT8 = "weights_int8"
+    WEIGHTS_KV_INT8 = "weights_kv_int8"
+    WEIGHTS_ONLY_INT4 = "weights_int4"
+    WEIGHTS_INT4_KV_INT8 = "weights_int4_kv_int8"
+    SMOOTH_QUANT = "smooth_quant"
+    FP8 = "fp8"
+    FP8_KV = "fp8_kv"
+
+
+class TrussTRTLLMPluginConfiguration(BaseModel):
+    paged_kv_cache: bool = True
+    gemm_plugin: str = "auto"
+    use_paged_context_fmha: bool = True
+    use_fp8_context_fmha: bool = False
+
+
+class CheckpointSource(str, Enum):
+    HF = "HF"
+    GCS = "GCS"
+    LOCAL = "LOCAL"
+    # REMOTE_URL is useful when the checkpoint lives on remote storage accessible via HTTP (e.g a presigned URL)
+    REMOTE_URL = "REMOTE_URL"
+
+
+class CheckpointRepository(BaseModel):
+    source: CheckpointSource
+    repo: str
+    revision: Optional[str] = None
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        if self.source == CheckpointSource.HF:
+            self._validate_hf_repo_id()
+
+    def _validate_hf_repo_id(self):
+        try:
+            validate_repo_id(self.repo)
+        except HFValidationError as e:
+            raise ValueError(
+                f"HuggingFace repository validation failed: {str(e)}"
+            ) from e
+
+
+class TrussTRTLLMBatchSchedulerPolicy(str, Enum):
+    MAX_UTILIZATION = "max_utilization"
+    GUARANTEED_NO_EVICT = "guaranteed_no_evict"
+
+
+class TrussSpecDecMode(str, Enum):
+    DRAFT_EXTERNAL = "DRAFT_TOKENS_EXTERNAL"
+
+
+class TrussTRTLLMRuntimeConfiguration(BaseModel):
+    kv_cache_free_gpu_mem_fraction: float = 0.9
+    kv_cache_host_memory_bytes: Optional[int] = None
+    enable_chunked_context: bool = True
+    batch_scheduler_policy: TrussTRTLLMBatchSchedulerPolicy = (
+        TrussTRTLLMBatchSchedulerPolicy.GUARANTEED_NO_EVICT
+    )
+    request_default_max_tokens: Optional[int] = None
+    total_token_limit: int = 500000
+
+
+class TrussTRTLLMBuildConfiguration(BaseModel):
+    base_model: TrussTRTLLMModel
+    max_seq_len: int
+    max_batch_size: int = 256
+    max_num_tokens: int = 8192
+    max_beam_width: int = 1
+    max_prompt_embedding_table_size: int = 0
+    checkpoint_repository: CheckpointRepository
+    gather_all_token_logits: bool = False
+    strongly_typed: bool = False
+    quantization_type: TrussTRTLLMQuantizationType = (
+        TrussTRTLLMQuantizationType.NO_QUANT
+    )
+    tensor_parallel_count: int = 1
+    pipeline_parallel_count: int = 1
+    plugin_configuration: TrussTRTLLMPluginConfiguration = (
+        TrussTRTLLMPluginConfiguration()
+    )
+    num_builder_gpus: Optional[int] = None
+    speculator: Optional[TrussSpeculatorConfiguration] = None
+
+    class Config:
+        extra = "forbid"
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self._validate_kv_cache_flags()
+        self._validate_speculator_config()
+        self._bei_specfic_migration()
+
+    @validator("max_beam_width")
+    def check_max_beam_width(cls, v: int):
+        if isinstance(v, int):
+            if v != 1:
+                raise ValueError(
+                    "max_beam_width greater than 1 is not currently supported"
+                )
+        return v
+
+    def _bei_specfic_migration(self):
+        """performs embedding specfic optimizations (no kv-cache, high batch size)"""
+        if self.base_model == TrussTRTLLMModel.ENCODER:
+            # Encoder specific settings
+            logger.info(
+                f"Your setting of `build.max_seq_len={self.max_seq_len}` is not used and "
+                "automatically inferred from the model repo config.json -> `max_position_embeddings`"
+            )
+
+            if self.max_num_tokens < BEI_REQUIRED_MAX_NUM_TOKENS:
+                logger.warning(
+                    f"build.max_num_tokens={self.max_num_tokens}, upgrading to {BEI_REQUIRED_MAX_NUM_TOKENS}"
+                )
+                self.max_num_tokens = BEI_REQUIRED_MAX_NUM_TOKENS
+            self.plugin_configuration.paged_kv_cache = False
+            self.plugin_configuration.use_paged_context_fmha = False
+
+            if "_kv" in self.quantization_type.value:
+                raise ValueError(
+                    "encoder does not have a kv-cache, therefore a kv specfic datatype is not valid"
+                    f"you selected build.quantization_type {self.quantization_type}"
+                )
+
+    def _validate_kv_cache_flags(self):
+        if not self.plugin_configuration.paged_kv_cache and (
+            self.plugin_configuration.use_paged_context_fmha
+            or self.plugin_configuration.use_fp8_context_fmha
+        ):
+            raise ValueError(
+                "Using paged context fmha or fp8 context fmha requires requires paged kv cache"
+            )
+        if (
+            self.plugin_configuration.use_fp8_context_fmha
+            and not self.plugin_configuration.use_paged_context_fmha
+        ):
+            raise ValueError("Using fp8 context fmha requires paged context fmha")
+        if (
+            self.plugin_configuration.use_fp8_context_fmha
+            and not self.quantization_type == TrussTRTLLMQuantizationType.FP8_KV
+        ):
+            raise ValueError("Using fp8 context fmha requires fp8 kv cache dtype")
+        return self
+
+    def _validate_speculator_config(self):
+        if self.speculator:
+            if self.base_model is TrussTRTLLMModel.WHISPER:
+                raise ValueError("Speculative decoding for Whisper is not supported.")
+            if not all(
+                [
+                    self.plugin_configuration.use_paged_context_fmha,
+                    self.plugin_configuration.paged_kv_cache,
+                ]
+            ):
+                raise ValueError(
+                    "KV cache block reuse must be enabled for speculative decoding target model."
+                )
+            if self.speculator.build:
+                if (
+                    self.tensor_parallel_count
+                    != self.speculator.build.tensor_parallel_count
+                ):
+                    raise ValueError(
+                        "Speculative decoding requires the same tensor parallelism for target and draft models."
+                    )
+
+    @property
+    def max_draft_len(self) -> Optional[int]:
+        if self.speculator:
+            return self.speculator.num_draft_tokens
+        return None
+
+
+class TrussSpeculatorConfiguration(BaseModel):
+    speculative_decoding_mode: TrussSpecDecMode = TrussSpecDecMode.DRAFT_EXTERNAL
+    num_draft_tokens: int
+    checkpoint_repository: Optional[CheckpointRepository] = None
+    runtime: TrussTRTLLMRuntimeConfiguration = TrussTRTLLMRuntimeConfiguration()
+    build: Optional[TrussTRTLLMBuildConfiguration] = None
+
+    def __init__(self, **data):
+        super().__init__(**data)
+        self._validate_checkpoint()
+
+    def _validate_checkpoint(self):
+        if not (bool(self.checkpoint_repository) ^ bool(self.build)):
+            raise ValueError(
+                "Speculative decoding requires exactly one of checkpoint_repository or build to be configured."
+            )
+
+    @property
+    def resolved_checkpoint_repository(self) -> CheckpointRepository:
+        if self.build:
+            return self.build.checkpoint_repository
+        elif self.checkpoint_repository:
+            return self.checkpoint_repository
+        else:
+            raise ValueError(
+                "Speculative decoding requires exactly one of checkpoint_repository or build to be configured."
+            )
+
+
+class TRTLLMConfiguration(BaseModel):
+    runtime: TrussTRTLLMRuntimeConfiguration = TrussTRTLLMRuntimeConfiguration()
+    build: TrussTRTLLMBuildConfiguration
+
+    @model_validator(mode="before")
+    @classmethod
+    def migrate_runtime_fields(cls, data: Any) -> Any:
+        extra_runtime_fields = {}
+        valid_build_fields = {}
+        if isinstance(data.get("build"), dict):
+            for key, value in data.get("build").items():
+                if key in TrussTRTLLMBuildConfiguration.__annotations__:
+                    valid_build_fields[key] = value
+                else:
+                    if key in TrussTRTLLMRuntimeConfiguration.__annotations__:
+                        logger.warning(f"Found runtime.{key}: {value} in build config")
+                        extra_runtime_fields[key] = value
+            if extra_runtime_fields:
+                logger.warning(
+                    f"Found extra fields {list(extra_runtime_fields.keys())} in build configuration, unspecified runtime fields will be configured using these values."
+                    " This configuration of deprecated fields is scheduled for removal, please upgrade to the latest truss version and update configs according to https://docs.baseten.co/performance/engine-builder-config."
+                )
+                if data.get("runtime"):
+                    data.get("runtime").update(
+                        {
+                            k: v
+                            for k, v in extra_runtime_fields.items()
+                            if k not in data.get("runtime")
+                        }
+                    )
+                else:
+                    data.update(
+                        {"runtime": {k: v for k, v in extra_runtime_fields.items()}}
+                    )
+            data.update({"build": valid_build_fields})
+            return data
+        return data
+
+    @model_validator(mode="after")
+    def after(self: "TRTLLMConfiguration") -> "TRTLLMConfiguration":
+        # check if there is an error wrt. runtime.enable_chunked_context
+        if (
+            self.runtime.enable_chunked_context
+            and (self.build.base_model != TrussTRTLLMModel.ENCODER)
+            and not (
+                self.build.plugin_configuration.use_paged_context_fmha
+                and self.build.plugin_configuration.paged_kv_cache
+            )
+        ):
+            if ENGINE_BUILDER_TRUSS_RUNTIME_MIGRATION:
+                logger.warning(
+                    "If trt_llm.runtime.enable_chunked_context is True, then trt_llm.build.plugin_configuration.use_paged_context_fmha and trt_llm.build.plugin_configuration.paged_kv_cache should be True. "
+                    "Setting trt_llm.build.plugin_configuration.use_paged_context_fmha and trt_llm.build.plugin_configuration.paged_kv_cache to True."
+                )
+                self.build.plugin_configuration.use_paged_context_fmha = True
+                self.build.plugin_configuration.paged_kv_cache = True
+            else:
+                raise ValueError(
+                    "If runtime.enable_chunked_context is True, then build.plugin_configuration.use_paged_context_fmha and build.plugin_configuration.paged_kv_cache should be True"
+                )
+
+        return self
+
+    @property
+    def requires_build(self):
+        return self.build is not None
+
+    # TODO(Abu): Replace this with model_dump(json=True)
+    # when pydantic v2 is used here
+    def to_json_dict(self, verbose=True):
+        return json.loads(self.json(exclude_unset=not verbose))