amzn-sagemaker-checkpointing 1.0.10__tar.gz → 1.0.12__tar.gz

amzn_sagemaker_checkpointing-1.0.12/DEVELOPING.md

@@ -0,0 +1,22 @@
+# Developing SageMakerCheckpointing
+
+This package uses the [hatch](https://hatch.pypa.io/latest/) build system.
+
+### Building
+
+A number of scripts and commands exist in `pyproject.toml` under the `scripts` configurations with more
+documentation in the comments of `pyproject.toml`. Running a script for a specific environment is simply running
+`hatch run <env_name>:<script>`. You can omit the `<env_name>` for those under the `default` environment.
+
+You need to set up hatch pluging first:
+```
+./setup-hatch.sh
+```
+
+### Available Hatch Commands
+
+- **`hatch run release`** - Runs typing checks (mypy), tests, and coverage.
+- **`hatch test --cover`** - Runs tests and coverage.
+- **`hatch typing`** - Runs mypy type checking.
+- **`hatch fmt`** - Formats code using ruff.
+- **`hatch build`** - builds both source and wheel distributions in ./build directory.

{amzn_sagemaker_checkpointing-1.0.10 → amzn_sagemaker_checkpointing-1.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amzn-sagemaker-checkpointing
-Version: 1.0.10
+Version: 1.0.12
 Summary: Amazon SageMaker Checkpointing Library
 License: Apache 2.0
 License-File: LICENSE.txt
@@ -95,12 +95,12 @@ following to your S3 bucket policy
 ```
 
 ## Installation
-### PreRequisites
+### Prerequisites
 ```bash
 pip install s3torchconnector tenacity torch boto3 botocore
 ```
 
-### Install amzn-sagemaker-checkpointing library
+### SageMaker Checkpointing Library
 ```bash
 pip install amzn-sagemaker-checkpointing
 ```

{amzn_sagemaker_checkpointing-1.0.10 → amzn_sagemaker_checkpointing-1.0.12}/README.md

@@ -82,12 +82,12 @@ following to your S3 bucket policy
 ```
 
 ## Installation
-### PreRequisites
+### Prerequisites
 ```bash
 pip install s3torchconnector tenacity torch boto3 botocore
 ```
 
-### Install amzn-sagemaker-checkpointing library
+### SageMaker Checkpointing Library
 ```bash
 pip install amzn-sagemaker-checkpointing
 ```

{amzn_sagemaker_checkpointing-1.0.10 → amzn_sagemaker_checkpointing-1.0.12}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "amzn-sagemaker-checkpointing"
-version = "1.0.10"
+version = "1.0.12"
 description = "Amazon SageMaker Checkpointing Library"
 readme = "README.md"
 license = { "text" = "Apache 2.0" }
@@ -72,6 +72,7 @@ exclude = [ "./build", ".hatch", "private" ]
 
 [tool.hatch.build]
 directory = "./build"
+exclude = ["DEVELOPING_INTERNAL.md"]
 
 [tool.hatch.env]
 requires = [ "hatch-pip-compile" ]

amzn_sagemaker_checkpointing-1.0.12/setup-hatch.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+set -e
+
+mkdir -p .hatch
+
+cat > .hatch/hatch_plugin.py << 'EOF'
+from hatch.env.collectors.plugin.interface import EnvironmentCollectorInterface
+
+class CustomEnvironmentCollector(EnvironmentCollectorInterface):
+    PLUGIN_NAME = 'custom'
+    
+    def get_initial_config(self):
+        return {}
+    
+    def finalize_config(self, config):
+        return config
+EOF
+
+echo "Hatch plugin created"

{amzn_sagemaker_checkpointing-1.0.10 → amzn_sagemaker_checkpointing-1.0.12}/src/amzn_sagemaker_checkpointing/checkpointing/filesystem/filesystem.py

@@ -19,6 +19,7 @@ import pickle
 import threading
 import time
 from dataclasses import dataclass
+from enum import Enum
 from logging import FileHandler
 from typing import Any, Union
 
@@ -46,9 +47,6 @@ from torch.futures import Future
 from amzn_sagemaker_checkpointing.config.sagemaker_checkpoint_config import (
     SageMakerCheckpointConfig,
 )
-from amzn_sagemaker_checkpointing.storage.clients.inmemory.exceptions import (
-    InMemoryServerError,
-)
 from amzn_sagemaker_checkpointing.storage.clients.inmemory.inmemory_client import (
     InMemoryCheckpointClient,
 )
@@ -80,6 +78,15 @@ class _SageMakerStorageInfo:
     offset: int
     length: int
 
+class StorageTier(Enum):
+    IN_MEMORY = 0
+    S3 = 1
+
+    def __str__(self):
+        return {
+            0: "IN_MEMORY",
+            1: "S3"
+        }[self.value]
 
 def _get_step_val(step: int, path: str | os.PathLike) -> int:
     """
@@ -791,51 +798,42 @@ class SageMakerTieredStorageReader(StorageReader):
 
     def read_metadata(self) -> Metadata:
         """
-        Retrieve and deserialize checkpoint metadata from the in-memory storage.
+        Retrieve and deserialize checkpoint metadata.
 
         Returns
         -------
         Metadata
             Metadata object containing checkpoint information.
-
-        Raises
-        ------
-        RuntimeError
-            If metadata retrieval fails.
+            (or) empty Metadata if not available
         """
-        # Use provided step or find latest available
-        if self.step is None:
-            self.step = self._find_latest_complete_step_across_tiers()
-
-        if not self.step:
-            self.logger.info(
-                f"[Rank {self.rank}] Step {self.step}: No checkpoints found"
-            )
-            return Metadata({})
-
-        # Try in-memory first (faster)
-        metadata_buffer = self._try_read_md_from_memory(self.step)
-        if metadata_buffer:
-            self.logger.info(
-                f"[Rank {self.rank}] Step {self.step}: Successfully read metadata from memory, size={len(metadata_buffer)} bytes"
-            )
-            return pickle.loads(metadata_buffer)
-
-        self.logger.info(
-            f"[Rank {self.rank}] Step {self.step}: In-memory metadata not found"
-        )
-        # Fallback to S3
-        if self.s3_base_path:
-            self.logger.info(
-                f"[Rank {self.rank}] Step {self.step}: Attempting metadata read from S3"
-            )
-            metadata_buffer = self._try_read_md_from_s3(self.step)
-            if metadata_buffer:
-                self.logger.info(
-                    f"[Rank {self.rank}] Step {self.step}: Successfully read metadata from S3, size={len(metadata_buffer)} bytes"
-                )
-                return pickle.loads(metadata_buffer)
-        return Metadata({})
+        metadata = Metadata({})
+        try:
+            if self.step is not None:
+                self.logger.info(f"[Rank {self.rank}] Step {self.step}: "
+                                 "reading metadata for configured step")
+                metadata = self._read_metadata_for_step(self.step)
+            else:
+                latest_step_all_tiers = self._get_latest_step_all_tiers()
+                for latest_step, tier in latest_step_all_tiers:
+                    if tier == StorageTier.IN_MEMORY:
+                        self.logger.info(f"[Rank {self.rank}] Attempting to read "
+                                         f"metadata from memory for {latest_step}")
+                        step_metadata = self._read_metadata_from_memory(latest_step)
+                    elif tier == StorageTier.S3:
+                        self.logger.info(f"[Rank {self.rank}] Attempting to read "
+                                         f"metadata from S3 for {latest_step}")
+                        step_metadata = self._read_metadata_from_s3(latest_step)
+                    if step_metadata is not None:
+                        metadata = step_metadata
+                        self.step = latest_step
+                        self.logger.info(f"[Rank {self.rank}] Metadata "
+                                         f"read from step {latest_step} of {tier} tier")
+                        break
+                if self.step is None:
+                    self.logger.error(f"[Rank {self.rank}] No checkpoints to read metadata")
+        except Exception as e:
+            self.logger.error(f"[Rank {self.rank}] Step {self.step}: read_metadata failed: {e}")
+        return metadata
 
     def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
         """
@@ -1083,23 +1081,6 @@ class SageMakerTieredStorageReader(StorageReader):
         """
         return True
 
-    def _find_latest_complete_step_across_tiers(self) -> int | None:
-        """Find latest step from both storage tiers."""
-        memory_step = self.client.get_latest_checkpoints(limit=1)
-        s3_step = self._find_latest_complete_step()
-        latest_step = None
-        if not memory_step:
-            latest_step = s3_step
-        elif not s3_step:
-            latest_step = memory_step[0]
-        else:
-            latest_step = max(memory_step[0], s3_step)
-        self.logger.info(
-            f"[Rank {self.rank}] Step {self.step}: Latest steps: "
-            f"memory:{memory_step}, s3:{s3_step}, across_tiers:{latest_step}"
-        )
-        return latest_step
-
     def _try_read_md_from_memory(self, step: int) -> bytes | None:
         """Try reading metadata from in-memory storage."""
         try:
@@ -1252,3 +1233,80 @@ class SageMakerTieredStorageReader(StorageReader):
                 f"[Rank {self.rank}] Failed to read item {item_index} from step {step}: {e}"
             )
             return None
+
+    def _read_metadata_from_memory(self, step) -> Metadata | None:
+        metadata = None
+        try:
+            metadata_buffer = self._try_read_md_from_memory(step)
+            if metadata_buffer:
+                self.logger.info(
+                    f"[Rank {self.rank}] Step {step}: Successfully read metadata from memory, "
+                    f"size={len(metadata_buffer)} bytes"
+                )
+                metadata = pickle.loads(metadata_buffer)
+            else:
+                self.logger.info(
+                    f"[Rank {self.rank}] Step {step}: "
+                    f"In-memory metadata not found"
+                )
+        except Exception as e:
+            self.logger.error(f"[Rank {self.rank}] Step {step}: _read_metadata_from_memory failed: {e}")
+        return metadata
+
+    def _read_metadata_from_s3(self, step) -> Metadata | None:
+        metadata = None
+        try:
+            if self.s3_base_path:
+                self.logger.info(
+                    f"[Rank {self.rank}] Step {step}: Attempting metadata read from S3"
+                )
+                metadata_buffer = self._try_read_md_from_s3(step)
+                if metadata_buffer:
+                    self.logger.info(f"[Rank {self.rank}] Step {step}: "
+                                     f"Successfully read metadata from size={len(metadata_buffer)} bytes")
+                    metadata = pickle.loads(metadata_buffer)
+                else:
+                    self.logger.info(
+                        f"[Rank {self.rank}] Step {step}: "
+                        "S3 metadata not found")
+            else:
+                self.logger.info(
+                    f"[Rank {self.rank}] Step {step}: Unable to read metadata "
+                    "as S3 path is not provided"
+                )
+        except Exception as e:
+            self.logger.error(f"[Rank {self.rank}] Step {step}: _read_metadata_from_s3 failed: {e}")
+        return metadata
+
+    def _read_metadata_for_step(self, step) -> Metadata:
+        metadata = Metadata({})
+        try:
+            in_memory_metadata = self._read_metadata_from_memory(step)
+            if in_memory_metadata is not None:
+                metadata = in_memory_metadata
+            else:
+                s3_metadata = self._read_metadata_from_s3(step)
+                if s3_metadata is not None:
+                    metadata = s3_metadata
+        except Exception as e:
+            self.logger.error(f"[Rank {self.rank}] Step {step}: _read_metadata_for_step failed: {e}")
+        return metadata
+
+    def _get_latest_step_all_tiers(self) -> list[tuple[int, StorageTier]]:
+        latest_step_all_tiers = []
+        try:
+            memory_steps = self.client.get_latest_checkpoints(limit=3)
+            if memory_steps:
+                latest_step_all_tiers = [(step, StorageTier.IN_MEMORY) for step in memory_steps]
+        except Exception as e:
+            self.logger.error(f"[Rank {self.rank}]: Failed to get memory steps: {e}")
+        try:
+            s3_step = self._find_latest_complete_step()
+            if s3_step:
+                latest_step_all_tiers.append((s3_step, StorageTier.S3))
+        except Exception as e:
+            self.logger.error(f"[Rank {self.rank}]: Failed to get S3 step: {e}")
+
+        latest_step_all_tiers.sort(key=lambda tier_step: (-tier_step[0], tier_step[1].value))
+        self.logger.info(f"[Rank {self.rank}] Latest steps across tiers: {latest_step_all_tiers}")
+        return latest_step_all_tiers

{amzn_sagemaker_checkpointing-1.0.10 → amzn_sagemaker_checkpointing-1.0.12}/src/amzn_sagemaker_checkpointing/storage/clients/inmemory/inmemory_client.py

@@ -471,17 +471,22 @@ class InMemoryCheckpointClient:
                 checksum=encode_base_64(hash_xxh3_128(data)), algorithm="xxh3_128"
             ).to_json()
         }
-        if isinstance(data, str) and os.path.exists(data):
-            with open(data, "rb") as f:
-                self._make_request(
-                    "POST",
-                    endpoint,
-                    data=f,
-                    headers=headers,
-                    timeout=timeout,
-                    retries=retries,
-                    retry_backoff=retry_backoff,
-                )
+        if isinstance(data, str):
+            try:
+                with open(data, "rb") as f:
+                    self._make_request(
+                        "POST",
+                        endpoint,
+                        data=f,
+                        headers=headers,
+                        timeout=timeout,
+                        retries=retries,
+                        retry_backoff=retry_backoff,
+                    )
+            except Exception as e:
+                error_msg = f"Error opening file: {data}"
+                self._logger.error(error_msg)
+                raise InMemoryStorageError(error_msg) from e
         else:
             self._make_request(
                 "POST",

amzn_sagemaker_checkpointing-1.0.12/tests/amzn_sagemaker_checkpointing/storage/clients/inmemory/test_delete_checkpoint.py

@@ -0,0 +1,148 @@
+from unittest.mock import Mock
+import pytest
+
+from amzn_sagemaker_checkpointing.config.in_memory_client import InMemoryClientConfig
+from amzn_sagemaker_checkpointing.storage.clients.inmemory.exceptions import InMemoryConfigError
+from utils.test_base import (
+    InMemoryCheckpointClientTest,
+    BASE_URL,
+    NAMESPACE,
+    RANK,
+    REQUEST_ERROR_CASES,
+    REQUEST_TIMEOUT,
+    WORLD_SIZE
+)
+
+
+class TestDeleteCheckpoint(InMemoryCheckpointClientTest):
+    STEP = 42
+
+    def setup_method(self):
+        super().setup_method()
+        self.checkpoint_path = f"v1/cp/checkpoints/{NAMESPACE}/{RANK}/{self.STEP}"
+
+    def test_delete_checkpoint_success(self):
+        # Arrange
+        mock_response = Mock(status_code=200)
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        self.client.delete_checkpoint(step=self.STEP)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{self.checkpoint_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+
+    def test_delete_checkpoint_with_custom_rank(self):
+        # Arrange
+        mock_response = Mock(status_code=200)
+        self.mock_session.request.return_value = mock_response
+        custom_rank = 5
+        custom_path = f"v1/cp/checkpoints/{NAMESPACE}/{custom_rank}/{self.STEP}"
+
+        # Act
+        self.client.delete_checkpoint(step=self.STEP, rank=custom_rank)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{custom_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+
+    def test_delete_checkpoint_with_metadata_index(self):
+        # Arrange
+        mock_response = Mock(status_code=200)
+        self.mock_session.request.return_value = mock_response
+        metadata_index = 0
+        metadata_rank = int(WORLD_SIZE) + metadata_index
+        metadata_path = f"v1/cp/checkpoints/{NAMESPACE}/{metadata_rank}/{self.STEP}"
+
+        # Act
+        self.client.delete_checkpoint(step=self.STEP, metadata_index=metadata_index)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{metadata_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+
+    def test_delete_checkpoint_with_custom_timeout(self):
+        # Arrange
+        mock_response = Mock(status_code=200)
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        self.client.delete_checkpoint(step=self.STEP, timeout=REQUEST_TIMEOUT)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{self.checkpoint_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=REQUEST_TIMEOUT
+        )
+
+    def test_delete_checkpoint_with_string_step(self):
+        # Arrange
+        mock_response = Mock(status_code=200)
+        self.mock_session.request.return_value = mock_response
+        step = "latest"
+        path = f"v1/cp/checkpoints/{NAMESPACE}/{RANK}/{step}"
+
+        # Act
+        self.client.delete_checkpoint(step=step)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+
+    def test_delete_checkpoint_invalid_metadata_index(self):
+        # Act & Assert
+        with pytest.raises(InMemoryConfigError) as exc_info:
+            self.client.delete_checkpoint(step=self.STEP, metadata_index=999)
+        assert "Invalid metadata_index" in str(exc_info.value)
+
+    @pytest.mark.parametrize("test_case", REQUEST_ERROR_CASES)
+    def test_delete_checkpoint_request_errors(self, test_case):
+        # Arrange
+        if "exception" in test_case["response"]:
+            self.mock_session.request.side_effect = test_case["response"]["exception"]
+        else:
+            mock_response = Mock()
+            mock_response.status_code = test_case["response"]["status_code"]
+            mock_response.text = test_case["response"]["text"]
+            self.mock_session.request.return_value = mock_response
+
+        # Act & Assert
+        self.assert_request_error(
+            test_case,
+            self.client.delete_checkpoint,
+            step=self.STEP
+        )

amzn_sagemaker_checkpointing-1.0.12/tests/amzn_sagemaker_checkpointing/storage/clients/inmemory/test_delete_namespace.py

@@ -0,0 +1,70 @@
+from unittest.mock import Mock, call
+import pytest
+
+from amzn_sagemaker_checkpointing.config.in_memory_client import InMemoryClientConfig
+from utils.test_base import (
+    InMemoryCheckpointClientTest,
+    BASE_URL,
+    NAMESPACE,
+    REQUEST_ERROR_CASES,
+    REQUEST_TIMEOUT,
+)
+
+
+class TestDeleteNamespace(InMemoryCheckpointClientTest):
+    def setup_method(self):
+        super().setup_method()
+        self.namespace_path = f"v1/cp/namespaces/{NAMESPACE}"
+
+    @pytest.mark.parametrize(
+        "params",
+        [
+            {"timeout": None},
+            {"timeout": REQUEST_TIMEOUT},
+        ],
+    )
+    def test_delete_namespace_success(self, params):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        self.client.delete_namespace(**params)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{self.namespace_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=params["timeout"] or InMemoryClientConfig.request_timeout
+        )
+
+    @pytest.mark.parametrize("test_case", REQUEST_ERROR_CASES)
+    def test_delete_namespace_request_errors(self, test_case):
+        # Arrange
+        if "exception" in test_case["response"]:
+            self.mock_session.request.side_effect = test_case["response"]["exception"]
+        else:
+            mock_response = Mock()
+            mock_response.status_code = test_case["response"]["status_code"]
+            mock_response.text = test_case["response"]["text"]
+            self.mock_session.request.return_value = mock_response
+
+        # Act & Assert
+        self.assert_http_adapter_and_retry_config()
+        self.assert_request_error(
+            test_case,
+            self.client.delete_namespace
+        )
+        self.mock_session.request.assert_called_once_with(
+            method="DELETE",
+            url=f"{BASE_URL}/{self.namespace_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )