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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")

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,180 @@ 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_hyperpod_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the HyperPod manifest.json S3 URI for a Nova training job.
+
+ HyperPod jobs write the manifest directly under the job directory:
+ ``<s3_output_path>/<training_job_name>/manifest.json``.
+
+ 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}/manifest.json"
+
+
+def build_nova_manifest_s3_uri(s3_output_path: str, training_job_name: str) -> str:
+ """Build the serverless manifest.json S3 URI for a Nova training job.
+
+ Serverless jobs write the manifest under a nested output directory:
+ ``<s3_output_path>/<training_job_name>/output/output/manifest.json``.
+
+ 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 any known
+ output layout.
+ """
+ # Nova jobs write their manifest to different locations depending on the
+ # training platform:
+ # HyperPod: <output>/<job>/manifest.json
+ # Serverless: <output>/<job>/output/output/manifest.json
+ # Serverful: <output>/<job>/output/output.tar.gz (manifest is inside)
+ # Try each in turn and surface every failure if none resolve, so the real
+ # cause is not masked by a misleading message from the last attempt.
+ hyperpod_manifest_uri = build_nova_hyperpod_manifest_s3_uri(
+ s3_output_path, training_job_name
+ )
+ serverless_manifest_uri = build_nova_manifest_s3_uri(s3_output_path, training_job_name)
+ tar_gz_uri = build_nova_output_tar_gz_s3_uri(s3_output_path, training_job_name)
+
+ attempts = [
+ ("HyperPod manifest.json", hyperpod_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverless manifest.json", serverless_manifest_uri, read_nova_checkpoint_uri_from_manifest),
+ ("serverful output.tar.gz", tar_gz_uri, _read_checkpoint_uri_from_tar_gz),
+ ]
+
+ errors = []
+ for label, uri, reader in attempts:
+ try:
+ return reader(s3_client, uri)
+ except Exception as error: # noqa: PERF203 - each attempt may fail independently
+ errors.append(f"{label} at {uri} failed: {error}")
+
+ raise ValueError(
+ "Could not resolve the Nova checkpoint URI from any known output layout. "
+ + " ".join(errors)
+ )

sagemaker-core/tests/unit/test_training_utils.py

@@ -0,0 +1,198 @@
+# 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_hyperpod_manifest_s3_uri,
+ 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_hyperpod_manifest_s3_uri():
+ result = build_nova_hyperpod_manifest_s3_uri("s3://bucket/output/", "my-job")
+ assert result == "s3://bucket/output/my-job/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_from_hyperpod_layout():
+ """HyperPod writes the manifest at <output>/<job>/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ hyperpod_key = "output/my-job/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != hyperpod_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_from_serverless_layout():
+ """Serverless writes the manifest at <output>/<job>/output/output/manifest.json."""
+ client = Mock()
+ client.exceptions = Mock()
+ no_such_key = type("NoSuchKey", (Exception,), {})
+ client.exceptions.NoSuchKey = no_such_key
+ serverless_key = "output/my-job/output/output/manifest.json"
+
+ def get_object(Bucket, Key):
+ if Key != serverless_key:
+ raise no_such_key()
+ body = Mock()
+ body.read.return_value = json.dumps(
+ {"checkpoint_s3_bucket": CHECKPOINT_URI}
+ ).encode("utf-8")
+ 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_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")