diff --git a/truss/base/truss_config.py b/truss/base/truss_config.py index 53f7c7f70..db602ddad 100644 --- a/truss/base/truss_config.py +++ b/truss/base/truss_config.py @@ -579,6 +579,130 @@ def _validate_unique_mount_locations(self) -> "Weights": return self +_BDN_PREFIX = "bdn://" +_BDN_VOLUME_SOURCE_REGEX = re.compile( + r"bdn://(?P[^/:@\x00]+)/(?P[^/:@\x00]+)" + r"(?::(?P[^/:@\x00]+)|@(?:b3:)?(?P[0-9a-fA-F]+))?" +) +_MIN_BDN_DIGEST_PREFIX_LENGTH = 12 +_MAX_BDN_DIGEST_LENGTH = 64 + + +def _validate_bdn_identifier(kind: str, value: str) -> None: + if not value: + raise ValueError(f"BDN {kind} must not be empty") + if len(value) > 256: + raise ValueError(f"BDN {kind} must be at most 256 characters") + if "/" in value or ".." in value or "\0" in value: + raise ValueError(f"Invalid BDN {kind}: {value!r}") + + +def _normalize_bdn_mount_path(value: str) -> str: + if "\0" in value: + raise ValueError("Volume mount must not contain null bytes") + mount_path = pathlib.PurePosixPath(value) + if not mount_path.is_absolute(): + raise ValueError( + f"Volume mount must be an absolute path (start with /), got: {value}" + ) + if mount_path == pathlib.PurePosixPath("/"): + raise ValueError("Volume mount must not be the filesystem root") + if ".." in mount_path.parts: + raise ValueError(f"Volume mount must not contain parent traversal: {value}") + return str(mount_path) + + +class BDNVolumeMount(custom_types.ConfigModel): + """An existing BDN volume mounted into a model container. + + BDN vocabulary, read off a reference like `bdn://weights/llama-8b:prod`: + + - A *namespace* (`weights`) groups volumes within your organization, and is + the unit that access grants and storage are scoped to. Names are + lowercase alphanumeric plus hyphens, at least two characters, and may not + begin with a digit; `namespaces` and `resolve` are reserved. + - A *volume* (`llama-8b`) is one versioned collection of files. Every + published version is immutable and identified by its content digest. + - A *tag* (`prod`) is a mutable, case-sensitive name pointing at one + version, repointed as newer versions are published. A reference carrying + neither tag nor digest resolves to the volume's head, its latest version. + + ``` + bdn: + mounts: + - source: bdn://weights/llama-8b:prod + path: /models/llama + ``` + """ + + source: Annotated[str, pydantic.StringConstraints(min_length=1)] = pydantic.Field( + ..., + description="BDN volume reference to mount (for example, bdn://weights/llama-8b:prod).", + ) + path: Annotated[str, pydantic.StringConstraints(min_length=1)] = pydantic.Field( + ..., description="Absolute path where the volume will be mounted at runtime." + ) + + @pydantic.field_validator("source") + @classmethod + def _validate_source(cls, value: str) -> str: + if not value.startswith(_BDN_PREFIX): + raise ValueError(f"Volume source must use the bdn:// scheme, got: {value}") + + match = _BDN_VOLUME_SOURCE_REGEX.fullmatch(value) + if match is None: + raise ValueError( + f"Invalid BDN volume source: '{value}'. " + "Expected format: bdn://namespace/volume[:tag|@digest]" + ) + + _validate_bdn_identifier("namespace", match.group("namespace")) + _validate_bdn_identifier("volume", match.group("volume")) + if tag := match.group("tag"): + _validate_bdn_identifier("tag", tag) + if digest := match.group("digest"): + if ( + not _MIN_BDN_DIGEST_PREFIX_LENGTH + <= len(digest) + <= _MAX_BDN_DIGEST_LENGTH + ): + raise ValueError( + "BDN digest must contain between " + f"{_MIN_BDN_DIGEST_PREFIX_LENGTH} and " + f"{_MAX_BDN_DIGEST_LENGTH} hexadecimal characters" + ) + return value + + @pydantic.field_validator("path") + @classmethod + def _validate_path(cls, value: str) -> str: + return _normalize_bdn_mount_path(value) + + +class BDNConfig(custom_types.ConfigModel): + """Configuration for mounting BDN volumes.""" + + mounts: list[BDNVolumeMount] = pydantic.Field( + default_factory=list, + description="Existing BDN volumes to mount when the model starts.", + ) + + @pydantic.field_validator("mounts") + @classmethod + def _validate_unique_mount_paths( + cls, mounts: list[BDNVolumeMount] + ) -> list[BDNVolumeMount]: + mount_paths: set[str] = set() + for volume_mount in mounts: + if volume_mount.path in mount_paths: + raise ValueError( + f"Duplicate volume mount path '{volume_mount.path}' - " + "each volume must have a unique mount path." + ) + mount_paths.add(volume_mount.path) + return mounts + + class AutoscalingMetric(pydantic.BaseModel): name: str target: float @@ -1393,6 +1517,9 @@ class TrussConfig(custom_types.ConfigModel): default_factory=lambda: Weights([]), description="Configure Baseten Delivery Network (BDN) for model weight delivery with multi-tier caching.", ) + bdn: BDNConfig = pydantic.Field( + default_factory=BDNConfig, description="Configure BDN volume mounts." + ) trt_llm: Optional[trt_llm_config.TRTLLMConfiguration] = pydantic.Field( default=None, description="TensorRT-LLM configuration for optimized LLM inference.", diff --git a/truss/config.schema.json b/truss/config.schema.json index cccb35d9a..2b2c262ed 100644 --- a/truss/config.schema.json +++ b/truss/config.schema.json @@ -46,6 +46,46 @@ "title": "AutoscalingMetric", "type": "object" }, + "BDNConfig": { + "additionalProperties": true, + "description": "Configuration for mounting BDN volumes.", + "properties": { + "mounts": { + "description": "Existing BDN volumes to mount when the model starts.", + "items": { + "$ref": "#/$defs/BDNVolumeMount" + }, + "title": "Mounts", + "type": "array" + } + }, + "title": "BDNConfig", + "type": "object" + }, + "BDNVolumeMount": { + "additionalProperties": true, + "description": "An existing BDN volume mounted into a model container.\n\nBDN vocabulary, read off a reference like `bdn://weights/llama-8b:prod`:\n\n- A *namespace* (`weights`) groups volumes within your organization, and is\n the unit that access grants and storage are scoped to. Names are\n lowercase alphanumeric plus hyphens, at least two characters, and may not\n begin with a digit; `namespaces` and `resolve` are reserved.\n- A *volume* (`llama-8b`) is one versioned collection of files. Every\n published version is immutable and identified by its content digest.\n- A *tag* (`prod`) is a mutable, case-sensitive name pointing at one\n version, repointed as newer versions are published. A reference carrying\n neither tag nor digest resolves to the volume's head, its latest version.\n\n```\nbdn:\n mounts:\n - source: bdn://weights/llama-8b:prod\n path: /models/llama\n```", + "properties": { + "source": { + "description": "BDN volume reference to mount (for example, bdn://weights/llama-8b:prod).", + "minLength": 1, + "title": "Source", + "type": "string" + }, + "path": { + "description": "Absolute path where the volume will be mounted at runtime.", + "minLength": 1, + "title": "Path", + "type": "string" + } + }, + "required": [ + "source", + "path" + ], + "title": "BDNVolumeMount", + "type": "object" + }, "BISLLM": { "additionalProperties": true, "description": "Configuration options for BIS LLM deployments.", @@ -2259,6 +2299,10 @@ "$ref": "#/$defs/Weights", "description": "Configure Baseten Delivery Network (BDN) for model weight delivery with multi-tier caching." }, + "bdn": { + "$ref": "#/$defs/BDNConfig", + "description": "Configure BDN volume mounts." + }, "trt_llm": { "anyOf": [ { diff --git a/truss/tests/test_config.py b/truss/tests/test_config.py index 9924aaa53..971c76dd8 100644 --- a/truss/tests/test_config.py +++ b/truss/tests/test_config.py @@ -15,6 +15,8 @@ Accelerator, AcceleratorSpec, BaseImage, + BDNConfig, + BDNVolumeMount, Build, CacheInternal, CheckpointList, @@ -2023,6 +2025,133 @@ def test_weights_serialization_roundtrip(self, tmp_path): assert config_new.weights.sources[0].allow_patterns == ["*.safetensors"] +class TestTrussConfigVolumeMounts: + def test_bdn_mounts_from_yaml(self, tmp_path): + yaml_content = """ + bdn: + mounts: + - source: bdn://weights/some-model:mytag + path: /models/some-model + """ + config_path = tmp_path / "config.yaml" + config_path.write_text(yaml_content) + + config = TrussConfig.from_yaml(config_path) + + assert config.bdn.mounts == [ + BDNVolumeMount( + source="bdn://weights/some-model:mytag", path="/models/some-model" + ) + ] + + def test_bdn_mounts_serialization_roundtrip(self, tmp_path): + config = TrussConfig( + bdn=BDNConfig( + mounts=[ + BDNVolumeMount( + source="bdn://weights/some-model:mytag", + path="/models/some-model", + ) + ] + ) + ) + config_path = tmp_path / "config.yaml" + config.write_to_yaml_file(config_path, verbose=False) + + serialized = yaml.safe_load(config_path.read_text()) + assert serialized["bdn"] == { + "mounts": [ + { + "source": "bdn://weights/some-model:mytag", + "path": "/models/some-model", + } + ] + } + parsed_config = TrussConfig.from_yaml(config_path) + assert parsed_config.bdn == config.bdn + + @pytest.mark.parametrize( + "source", + [ + "weights/llama:prod", + "ftp://weights/llama", + "bdn://foo", + "bdn:///llama:prod", + "bdn://weights/:prod", + "bdn://weights/llama:", + "bdn://weights/models/llama:prod", + "bdn://weights/llama:prod:extra", + "bdn://weights/llama@not-a-digest", + "bdn://weights/llama@abcdef01234", + "bdn://weights/llama@b3:abcdef01234", + f"bdn://weights/llama@{'a' * 65}", + f"bdn://weights/llama@b3:{'a' * 65}", + "bdn://weights/llama@b3:", + ], + ) + def test_volume_source_rejects_invalid_reference(self, source): + with pytest.raises(pydantic.ValidationError): + BDNVolumeMount(source=source, path="/models/llama") + + @pytest.mark.parametrize( + "source", + [ + "bdn://weights/llama", + "bdn://weights/llama:prod", + "bdn://weights/llama@abcdef012345", + "bdn://weights/llama@b3:ABCDEF012345", + f"bdn://weights/llama@{'a' * 64}", + ], + ) + def test_volume_source_accepts_supported_references(self, source): + volume_mount = BDNVolumeMount(source=source, path="/models/llama") + + assert volume_mount.source == source + + def test_volume_mount_requires_absolute_path(self): + with pytest.raises(pydantic.ValidationError, match="absolute path"): + BDNVolumeMount(source="bdn://weights/llama:prod", path="models/llama") + + def test_volume_mount_normalizes_path(self): + volume_mount = BDNVolumeMount( + source="bdn://weights/llama:prod", path="/models/./llama/" + ) + + assert volume_mount.path == "/models/llama" + + def test_volume_mount_paths_must_be_unique(self): + with pytest.raises( + pydantic.ValidationError, match="Duplicate volume mount path" + ): + BDNConfig( + mounts=[ + BDNVolumeMount(source="bdn://weights/llama:prod", path="/models"), + BDNVolumeMount( + source="bdn://weights/mistral:prod", path="/models/" + ), + ] + ) + + @pytest.mark.parametrize( + "source", + [ + "hf://Qwen/Qwen3-Omni-30B-A3B-Instruct", + "hf://meta-llama/Llama-2-7b@main", + "s3://bucket/path", + "gs://bucket/path", + "azure://account/container/path", + "r2://account_id.bucket/path", + "cw://bucket/path", + "https://example.com/model.bin", + ], + ) + def test_volume_mount_rejects_external_source(self, source): + with pytest.raises( + pydantic.ValidationError, match="must use the bdn:// scheme" + ): + BDNVolumeMount(source=source, path="/models/external") + + class TestCheckpointListNoMixing: """CheckpointList rejects mixing training-job and loops checkpoint sources."""