feat(serve): add opt-in model source tag-based resource reuse

sagemaker-core/src/sagemaker/core/training/utils.py

@@ -13,8 +13,12 @@
 """Training utilities."""
 from __future__ import absolute_import
 
+import io
+import json
 import os
+import tarfile
 from typing import Any, Literal
+from urllib.parse import urlparse
 from sagemaker.core.utils.utils import Unassigned
 
 
@@ -75,3 +79,147 @@ def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] =
         return path.endswith("/")
 
     return path_type == "Any"
+
+
+_MANIFEST_CHECKPOINT_KEY = "checkpoint_s3_bucket"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+    """Build the manifest.json S3 URI for a Nova training job.
+
+    Args:
+        s3_output_path: The training job's ``output_data_config.s3_output_path``.
+        training_job_name: The training job name.
+
+    Returns:
+        Fully-qualified S3 URI to the job's manifest.json.
+    """
+    output_path = s3_output_path.rstrip("/")
+    return f"{output_path}/{training_job_name}/output/output/manifest.json"
+
+
+def build_nova_output_tar_gz_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+    """Build the output.tar.gz S3 URI for a Nova training job.
+
+    Args:
+        s3_output_path: The training job's ``output_data_config.s3_output_path``.
+        training_job_name: The training job name.
+
+    Returns:
+        Fully-qualified S3 URI to the job's output.tar.gz.
+    """
+    output_path = s3_output_path.rstrip("/")
+    return f"{output_path}/{training_job_name}/output/output.tar.gz"
+
+
+def _split_s3_uri(s3_uri: str) -> tuple:
+    """Split an S3 URI into (bucket, key)."""
+    parsed = urlparse(s3_uri)
+    return parsed.netloc, parsed.path.lstrip("/")
+
+
+def read_nova_checkpoint_uri_from_manifest(s3_client, s3_uri: str) -> str:
+    """Read the checkpoint URI from a raw manifest.json object in S3.
+
+    Args:
+        s3_client: A boto3 S3 client.
+        s3_uri: S3 URI of the manifest.json object.
+
+    Returns:
+        The ``checkpoint_s3_bucket`` value from the manifest.
+
+    Raises:
+        ValueError: If the object is missing, unparseable, or lacks the key.
+    """
+    bucket, key = _split_s3_uri(s3_uri)
+    try:
+        response = s3_client.get_object(Bucket=bucket, Key=key)
+        manifest = json.loads(response["Body"].read().decode("utf-8"))
+    except s3_client.exceptions.NoSuchKey:
+        raise ValueError(f"manifest.json not found at s3://{bucket}/{key}")
+    except json.JSONDecodeError as e:
+        raise ValueError(f"Failed to parse manifest.json: {e}")
+
+    checkpoint_uri = manifest.get(_MANIFEST_CHECKPOINT_KEY)
+    if not checkpoint_uri:
+        raise ValueError(
+            f"'{_MANIFEST_CHECKPOINT_KEY}' not found in manifest.json. "
+            f"Available keys: {list(manifest.keys())}"
+        )
+    return checkpoint_uri
+
+
+def _read_checkpoint_uri_from_tar_gz(s3_client, s3_uri: str) -> str:
+    """Read the checkpoint URI from a manifest.json inside an output.tar.gz in S3.
+
+    Args:
+        s3_client: A boto3 S3 client.
+        s3_uri: S3 URI of the output.tar.gz object.
+
+    Returns:
+        The ``checkpoint_s3_bucket`` value from the embedded manifest.
+
+    Raises:
+        ValueError: If the archive or manifest is missing or lacks the key.
+    """
+    bucket, key = _split_s3_uri(s3_uri)
+    response = s3_client.get_object(Bucket=bucket, Key=key)
+    body = response["Body"].read()
+    with tarfile.open(fileobj=io.BytesIO(body), mode="r:gz") as tar:
+        for member in tar.getmembers():
+            if not member.name.endswith("manifest.json"):
+                continue
+            extracted = tar.extractfile(member)
+            if extracted is None:
+                continue
+            manifest = json.loads(extracted.read().decode("utf-8"))
+            checkpoint_uri = manifest.get(_MANIFEST_CHECKPOINT_KEY)
+            if checkpoint_uri:
+                return checkpoint_uri
+
+    raise ValueError(
+        f"'{_MANIFEST_CHECKPOINT_KEY}' not found in manifest.json within "
+        f"s3://{bucket}/{key}"
+    )
+
+
+def resolve_nova_checkpoint_uri(
+    s3_client,
+    s3_output_path: str,
+    training_job_name: str,
+) -> str:
+    """Resolve the Nova checkpoint (escrow) URI from a training job's output.
+
+    Reads ``checkpoint_s3_bucket`` from the job's manifest.json. The manifest is
+    first looked up as a raw object, and if that fails, it falls back to the copy
+    packaged inside ``output.tar.gz``.
+
+    Args:
+        s3_client: A boto3 S3 client.
+        s3_output_path: The training job's ``output_data_config.s3_output_path``.
+        training_job_name: The training job name.
+
+    Returns:
+        The checkpoint URI recorded in the manifest.
+
+    Raises:
+        ValueError: If the checkpoint URI cannot be resolved from either source.
+    """
+    manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+    try:
+        return read_nova_checkpoint_uri_from_manifest(s3_client, manifest_uri)
+    except Exception as manifest_error:
+        # The raw manifest.json is the primary source (e.g. HyperPod). Only fall
+        # back to the copy inside output.tar.gz (e.g. some SMTJ jobs) if it exists.
+        # If the fallback also fails, surface BOTH errors so the raw-manifest
+        # failure isn't masked by a misleading "not found in tar.gz" message.
+        tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+        try:
+            return _read_checkpoint_uri_from_tar_gz(s3_client, tar_gz_uri)
+        except Exception as tar_gz_error:
+            raise ValueError(
+                "Could not resolve the Nova checkpoint URI. "
+                f"Reading manifest.json at {manifest_uri} failed: {manifest_error}. "
+                f"Falling back to output.tar.gz at {tar_gz_uri} also failed: "
+                f"{tar_gz_error}."
+            ) from manifest_error

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,146 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file 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.
+"""Unit tests for Nova manifest/checkpoint helpers in training/utils.py."""
+import io
+import json
+import tarfile
+
+import pytest
+from unittest.mock import Mock
+
+from sagemaker.core.training.utils import (
+    build_nova_manifest_s3_uri,
+    build_nova_output_tar_gz_s3_uri,
+    read_nova_checkpoint_uri_from_manifest,
+    resolve_nova_checkpoint_uri,
+)
+
+CHECKPOINT_URI = "s3://bucket/ckpt/step_100"
+
+
+def test_build_nova_manifest_s3_uri():
+    result = build_nova_manifest_s3_uri("s3://bucket/output/", "my-job")
+    assert result == "s3://bucket/output/my-job/output/output/manifest.json"
+
+
+def test_build_nova_manifest_s3_uri_strips_trailing_slash():
+    assert build_nova_manifest_s3_uri(
+        "s3://bucket/output//", "my-job"
+    ) == "s3://bucket/output/my-job/output/output/manifest.json"
+
+
+def test_build_nova_output_tar_gz_s3_uri():
+    result = build_nova_output_tar_gz_s3_uri("s3://bucket/output", "my-job")
+    assert result == "s3://bucket/output/my-job/output/output.tar.gz"
+
+
+def _s3_client_returning(body_bytes):
+    client = Mock()
+    client.exceptions = Mock()
+    client.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {})
+    body = Mock()
+    body.read.return_value = body_bytes
+    client.get_object.return_value = {"Body": body}
+    return client
+
+
+def test_read_manifest_returns_checkpoint_uri():
+    client = _s3_client_returning(
+        json.dumps({"checkpoint_s3_bucket": CHECKPOINT_URI}).encode("utf-8")
+    )
+    result = read_nova_checkpoint_uri_from_manifest(
+        client, "s3://bucket/output/my-job/output/output/manifest.json"
+    )
+    assert result == CHECKPOINT_URI
+    client.get_object.assert_called_once_with(
+        Bucket="bucket", Key="output/my-job/output/output/manifest.json"
+    )
+
+
+def test_read_manifest_missing_key_raises():
+    client = _s3_client_returning(json.dumps({"other": "value"}).encode("utf-8"))
+    with pytest.raises(ValueError, match="checkpoint_s3_bucket"):
+        read_nova_checkpoint_uri_from_manifest(client, "s3://bucket/manifest.json")
+
+
+def test_read_manifest_not_found_raises():
+    client = Mock()
+    client.exceptions = Mock()
+    client.exceptions.NoSuchKey = type("NoSuchKey", (Exception,), {})
+    client.get_object.side_effect = client.exceptions.NoSuchKey()
+    with pytest.raises(ValueError, match="manifest.json not found"):
+        read_nova_checkpoint_uri_from_manifest(client, "s3://bucket/manifest.json")
+
+
+def test_read_manifest_invalid_json_raises():
+    client = _s3_client_returning(b"not-json")
+    with pytest.raises(ValueError, match="Failed to parse manifest.json"):
+        read_nova_checkpoint_uri_from_manifest(client, "s3://bucket/manifest.json")
+
+
+def _make_tar_gz_with_manifest(manifest_dict):
+    content = json.dumps(manifest_dict).encode("utf-8")
+    buf = io.BytesIO()
+    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
+        info = tarfile.TarInfo(name="manifest.json")
+        info.size = len(content)
+        tar.addfile(info, io.BytesIO(content))
+    return buf.getvalue()
+
+
+def test_resolve_checkpoint_uri_from_raw_manifest():
+    client = _s3_client_returning(
+        json.dumps({"checkpoint_s3_bucket": CHECKPOINT_URI}).encode("utf-8")
+    )
+    result = resolve_nova_checkpoint_uri(client, "s3://bucket/output/", "my-job")
+    assert result == CHECKPOINT_URI
+
+
+def test_resolve_checkpoint_uri_falls_back_to_tar_gz():
+    client = Mock()
+    client.exceptions = Mock()
+    no_such_key = type("NoSuchKey", (Exception,), {})
+    client.exceptions.NoSuchKey = no_such_key
+    tar_bytes = _make_tar_gz_with_manifest({"checkpoint_s3_bucket": CHECKPOINT_URI})
+
+    def get_object(Bucket, Key):
+        if Key.endswith("manifest.json"):
+            raise no_such_key()
+        body = Mock()
+        body.read.return_value = tar_bytes
+        return {"Body": body}
+
+    client.get_object.side_effect = get_object
+
+    result = resolve_nova_checkpoint_uri(client, "s3://bucket/output/", "my-job")
+    assert result == CHECKPOINT_URI
+
+
+def test_resolve_checkpoint_uri_raises_when_both_sources_fail():
+    client = Mock()
+    client.exceptions = Mock()
+    no_such_key = type("NoSuchKey", (Exception,), {})
+    client.exceptions.NoSuchKey = no_such_key
+    tar_bytes = _make_tar_gz_with_manifest({"other": "value"})
+
+    def get_object(Bucket, Key):
+        if Key.endswith("manifest.json"):
+            raise no_such_key()
+        body = Mock()
+        body.read.return_value = tar_bytes
+        return {"Body": body}
+
+    client.get_object.side_effect = get_object
+
+    with pytest.raises(ValueError, match="checkpoint_s3_bucket"):
+        resolve_nova_checkpoint_uri(client, "s3://bucket/output/", "my-job")