From f85416147d1f7177514c96d2fe4f35ad864ff8df Mon Sep 17 00:00:00 2001 From: aviruthen <91846056+aviruthen@users.noreply.github.com> Date: Wed, 28 Oct 2026 13:19:03 -0700 Subject: [PATCH 1/2] fix: normalize tar member paths when extracting Bedrock model artifacts --- .../sagemaker/serve/bedrock_model_builder.py | 29 ++++- .../tests/unit/test_bedrock_model_builder.py | 106 ++++++++++++++++++ 2 files changed, 134 insertions(+), 1 deletion(-) diff --git a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py index 0c1693b2f6..28b3115ffa 100644 --- a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py @@ -16,6 +16,7 @@ import json import os +import re import time import logging from datetime import datetime, timezone @@ -95,6 +96,28 @@ def _log_bedrock_api_call(api_name: str, params: Dict[str, Any], response: Dict[ print(f"[BedrockModelBuilder] API call logged to: {filepath}") +def _tar_member_relative_key(name: str) -> Optional[str]: + """Normalize a tar member name into a key relative to the extraction prefix. + + Leading ``./`` segments and empty segments are removed so that archives + created with ``tar -C .`` map to the same keys as archives created + from the directory contents. + + Args: + name: The tar member name. + + Returns: + The normalized relative key, or None if the name is absolute, contains + parent directory references, or does not name a file. + """ + if not name or name.startswith(("/", "\\")) or re.match(r"^[A-Za-z]:", name): + return None + if ".." in re.split(r"[\\/]", name): + return None + relative = "/".join(part for part in name.split("/") if part not in ("", ".")) + return relative or None + + class BedrockModelBuilder: """Builder class for deploying models to Amazon Bedrock. @@ -1080,10 +1103,14 @@ def _extract_tar_gz_to_s3(self, tar_gz_uri: str) -> str: for member in tar: if not member.isfile(): continue + relative_key = _tar_member_relative_key(member.name) + if relative_key is None: + logger.warning("Skipping archive member with unsupported path: %s", member.name) + continue f = tar.extractfile(member) if f is None: continue - dest_key = extract_prefix + member.name + dest_key = extract_prefix + relative_key size_mb = member.size / (1024 * 1024) extracted_count += 1 logger.info("Extracting [%d]: %s (%.1f MB)", extracted_count, member.name, size_mb) diff --git a/sagemaker-serve/tests/unit/test_bedrock_model_builder.py b/sagemaker-serve/tests/unit/test_bedrock_model_builder.py index f1e8a3d573..6677dc1ebf 100644 --- a/sagemaker-serve/tests/unit/test_bedrock_model_builder.py +++ b/sagemaker-serve/tests/unit/test_bedrock_model_builder.py @@ -1625,3 +1625,109 @@ def test_oss_reuse_preserves_user_tags(self): assert user_tag in tags source_tag = {"key": "sagemaker.amazonaws.com/model-source", "value": "s3://bucket/path/"} assert source_tag in tags + + +# ── _extract_tar_gz_to_s3 ─────────────────────────────────────────────────── + + +def _tar_gz_bytes(members): + """Build an in-memory tar.gz from (name, content) pairs.""" + import io + import tarfile + + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + for name, content in members: + info = tarfile.TarInfo(name=name) + info.size = len(content) + tar.addfile(info, io.BytesIO(content)) + return buf.getvalue() + + +def _extract_with_members(members): + """Run _extract_tar_gz_to_s3 against a fake S3 client and return put_object keys.""" + import io + + s3_client = Mock() + s3_client.list_objects_v2.return_value = {"KeyCount": 0} + s3_client.get_object.return_value = {"Body": io.BytesIO(_tar_gz_bytes(members))} + b = object.__new__(BedrockModelBuilder) + b.boto_session = Mock() + b.boto_session.client.return_value = s3_client + + uri = b._extract_tar_gz_to_s3("s3://bucket/jobs/run/output/model.tar.gz") + + assert uri == "s3://bucket/jobs/run/output/extracted/" + return {c.kwargs["Key"]: c.kwargs["Body"] for c in s3_client.put_object.call_args_list} + + +class TestTarMemberRelativeKey: + @pytest.mark.parametrize( + "name,expected", + [ + ("config.json", "config.json"), + ("checkpoints/hf/config.json", "checkpoints/hf/config.json"), + ("./config.json", "config.json"), + ("./checkpoints//hf/./config.json", "checkpoints/hf/config.json"), + ("model..safetensors", "model..safetensors"), + ], + ) + def test_normalizes_relative_names(self, name, expected): + from sagemaker.serve.bedrock_model_builder import _tar_member_relative_key + + assert _tar_member_relative_key(name) == expected + + @pytest.mark.parametrize( + "name", + [ + "", + ".", + "./", + "/config.json", + "\\config.json", + "C:/config.json", + "../config.json", + "a/../../config.json", + "a\\..\\config.json", + ], + ) + def test_rejects_non_relative_names(self, name): + from sagemaker.serve.bedrock_model_builder import _tar_member_relative_key + + assert _tar_member_relative_key(name) is None + + +class TestExtractTarGzToS3: + def test_uploads_members_under_extract_prefix(self): + keys = _extract_with_members( + [("config.json", b"{}"), ("./checkpoints/hf/model.safetensors", b"weights")] + ) + + assert keys == { + "jobs/run/output/extracted/config.json": b"{}", + "jobs/run/output/extracted/checkpoints/hf/model.safetensors": b"weights", + } + + def test_skips_members_with_unsupported_paths(self): + keys = _extract_with_members( + [("config.json", b"{}"), ("../other/file.txt", b"x"), ("/abs/file.txt", b"y")] + ) + + assert keys == {"jobs/run/output/extracted/config.json": b"{}"} + + def test_raises_when_no_members_extracted(self): + with pytest.raises(RuntimeError, match="No files found"): + _extract_with_members([("../file.txt", b"x")]) + + def test_skips_extraction_when_prefix_exists(self): + s3_client = Mock() + s3_client.list_objects_v2.return_value = {"KeyCount": 1} + b = object.__new__(BedrockModelBuilder) + b.boto_session = Mock() + b.boto_session.client.return_value = s3_client + + uri = b._extract_tar_gz_to_s3("s3://bucket/jobs/model.tar.gz") + + assert uri == "s3://bucket/jobs/extracted/" + s3_client.get_object.assert_not_called() + s3_client.put_object.assert_not_called() From 1ad6824d596e9dd9e8fbc7132e1e55bd29229963 Mon Sep 17 00:00:00 2001 From: aviruthen <91846056+aviruthen@users.noreply.github.com> Date: Thu, 29 Oct 2026 16:32:57 -0700 Subject: [PATCH 2/2] fix: report skipped archive members when extracting Bedrock model artifacts --- .../sagemaker/serve/bedrock_model_builder.py | 24 +++++++- .../tests/unit/test_bedrock_model_builder.py | 58 +++++++++++++------ 2 files changed, 63 insertions(+), 19 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py index 28b3115ffa..38adea2505 100644 --- a/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py +++ b/sagemaker-serve/src/sagemaker/serve/bedrock_model_builder.py @@ -101,14 +101,21 @@ def _tar_member_relative_key(name: str) -> Optional[str]: Leading ``./`` segments and empty segments are removed so that archives created with ``tar -C .`` map to the same keys as archives created - from the directory contents. + from the directory contents. Names that normalize to the same key (for + example ``./a`` and ``a``) map to one key, and the later member wins, as + it would when extracting the archive to a directory. + + Only ``/`` is treated as a separator when building the key, so a backslash + in a name is kept as a literal character. Names that are rooted or contain + a ``..`` segment under either separator are not supported, so that keys + stay relative under both POSIX and Windows path conventions. Args: name: The tar member name. Returns: The normalized relative key, or None if the name is absolute, contains - parent directory references, or does not name a file. + a ``..`` segment, or is empty after normalization. """ if not name or name.startswith(("/", "\\")) or re.match(r"^[A-Za-z]:", name): return None @@ -1099,12 +1106,14 @@ def _extract_tar_gz_to_s3(self, tar_gz_uri: str) -> str: stream.seekable = lambda: False extracted_count = 0 + skipped_count = 0 with tarfile.open(fileobj=stream, mode="r|gz") as tar: for member in tar: if not member.isfile(): continue relative_key = _tar_member_relative_key(member.name) if relative_key is None: + skipped_count += 1 logger.warning("Skipping archive member with unsupported path: %s", member.name) continue f = tar.extractfile(member) @@ -1117,7 +1126,18 @@ def _extract_tar_gz_to_s3(self, tar_gz_uri: str) -> str: s3_client.put_object(Bucket=bucket, Key=dest_key, Body=f.read()) if extracted_count == 0: + if skipped_count: + raise RuntimeError( + f"No files extracted from {tar_gz_uri}: {skipped_count} archive member(s) " + "were skipped because their paths are absolute or contain '..' segments." + ) raise RuntimeError(f"No files found in {tar_gz_uri}.") + if skipped_count: + logger.warning( + "Skipped %d archive member(s) with unsupported paths in %s.", + skipped_count, + tar_gz_uri, + ) logger.info("Extracted %d files to %s", extracted_count, extract_uri) return extract_uri diff --git a/sagemaker-serve/tests/unit/test_bedrock_model_builder.py b/sagemaker-serve/tests/unit/test_bedrock_model_builder.py index 6677dc1ebf..e7775a6aca 100644 --- a/sagemaker-serve/tests/unit/test_bedrock_model_builder.py +++ b/sagemaker-serve/tests/unit/test_bedrock_model_builder.py @@ -12,7 +12,10 @@ # language governing permissions and limitations under the License. """Unit tests for BedrockModelBuilder.""" +import io import json +import tarfile + import pytest from unittest.mock import Mock, patch from botocore.exceptions import ClientError @@ -20,7 +23,11 @@ from sagemaker.core.resources import TrainingJob from sagemaker.core.shapes import ModelArtifacts, OutputDataConfig from sagemaker.core.utils.utils import Unassigned -from sagemaker.serve.bedrock_model_builder import BedrockModelBuilder, _is_nova_model +from sagemaker.serve.bedrock_model_builder import ( + BedrockModelBuilder, + _is_nova_model, + _tar_member_relative_key, +) MODULE = "sagemaker.serve.bedrock_model_builder" @@ -1631,23 +1638,25 @@ def test_oss_reuse_preserves_user_tags(self): def _tar_gz_bytes(members): - """Build an in-memory tar.gz from (name, content) pairs.""" - import io - import tarfile + """Build an in-memory tar.gz from (name, content) pairs. + A content of None adds a directory entry instead of a file. + """ buf = io.BytesIO() with tarfile.open(fileobj=buf, mode="w:gz") as tar: for name, content in members: info = tarfile.TarInfo(name=name) - info.size = len(content) - tar.addfile(info, io.BytesIO(content)) + if content is None: + info.type = tarfile.DIRTYPE + tar.addfile(info) + else: + info.size = len(content) + tar.addfile(info, io.BytesIO(content)) return buf.getvalue() def _extract_with_members(members): """Run _extract_tar_gz_to_s3 against a fake S3 client and return put_object keys.""" - import io - s3_client = Mock() s3_client.list_objects_v2.return_value = {"KeyCount": 0} s3_client.get_object.return_value = {"Body": io.BytesIO(_tar_gz_bytes(members))} @@ -1670,11 +1679,10 @@ class TestTarMemberRelativeKey: ("./config.json", "config.json"), ("./checkpoints//hf/./config.json", "checkpoints/hf/config.json"), ("model..safetensors", "model..safetensors"), + ("a\\b.json", "a\\b.json"), ], ) def test_normalizes_relative_names(self, name, expected): - from sagemaker.serve.bedrock_model_builder import _tar_member_relative_key - assert _tar_member_relative_key(name) == expected @pytest.mark.parametrize( @@ -1691,9 +1699,7 @@ def test_normalizes_relative_names(self, name, expected): "a\\..\\config.json", ], ) - def test_rejects_non_relative_names(self, name): - from sagemaker.serve.bedrock_model_builder import _tar_member_relative_key - + def test_rejects_unsupported_names(self, name): assert _tar_member_relative_key(name) is None @@ -1708,16 +1714,34 @@ def test_uploads_members_under_extract_prefix(self): "jobs/run/output/extracted/checkpoints/hf/model.safetensors": b"weights", } - def test_skips_members_with_unsupported_paths(self): + def test_skips_directory_members(self): keys = _extract_with_members( - [("config.json", b"{}"), ("../other/file.txt", b"x"), ("/abs/file.txt", b"y")] + [("./", None), ("./checkpoints", None), ("./checkpoints/config.json", b"{}")] ) + assert keys == {"jobs/run/output/extracted/checkpoints/config.json": b"{}"} + + def test_names_normalizing_to_same_key_keep_last_member(self): + keys = _extract_with_members([("./config.json", b"first"), ("config.json", b"second")]) + + assert keys == {"jobs/run/output/extracted/config.json": b"second"} + + def test_skips_members_with_unsupported_paths(self, caplog): + with caplog.at_level("WARNING"): + keys = _extract_with_members( + [("config.json", b"{}"), ("../other/file.txt", b"x"), ("/abs/file.txt", b"y")] + ) + assert keys == {"jobs/run/output/extracted/config.json": b"{}"} + assert "Skipped 2 archive member(s)" in caplog.text + + def test_raises_with_skipped_count_when_all_members_skipped(self): + with pytest.raises(RuntimeError, match="2 archive member\\(s\\) were skipped"): + _extract_with_members([("../file.txt", b"x"), ("/abs/file.txt", b"y")]) - def test_raises_when_no_members_extracted(self): + def test_raises_when_archive_has_no_files(self): with pytest.raises(RuntimeError, match="No files found"): - _extract_with_members([("../file.txt", b"x")]) + _extract_with_members([("./", None)]) def test_skips_extraction_when_prefix_exists(self): s3_client = Mock()