amzn-sagemaker-checkpointing 1.0.11__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.11 → amzn_sagemaker_checkpointing-1.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amzn-sagemaker-checkpointing
-Version: 1.0.11
+Version: 1.0.12
 Summary: Amazon SageMaker Checkpointing Library
 License: Apache 2.0
 License-File: LICENSE.txt

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

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "amzn-sagemaker-checkpointing"
-version = "1.0.11"
+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.11 → 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
+        )

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

@@ -0,0 +1,217 @@
+from unittest.mock import Mock, call, patch
+import pytest
+
+from amzn_sagemaker_checkpointing.config.in_memory_client import InMemoryClientConfig
+from amzn_sagemaker_checkpointing.storage.clients.inmemory.exceptions import (
+    InMemoryConfigError,
+    InMemoryStorageError
+)
+from utils.test_base import (
+    InMemoryCheckpointClientTest,
+    BASE_URL,
+    NAMESPACE,
+    RANK,
+    REQUEST_ERROR_CASES,
+    REQUEST_TIMEOUT,
+    WORLD_SIZE
+)
+
+
+class TestGetCheckpoint(InMemoryCheckpointClientTest):
+    # Constants
+    STEP = 42
+    TEST_CONTENT = b"test checkpoint data"
+    MOCK_CHECKSUM = "mock-checksum"
+    MOCK_ENCODED_CHECKSUM = "bW9jay1jaGVja3N1bQ=="  # base64 of mock-checksum
+    CONTENT_LENGTH = len(TEST_CONTENT)
+    RESPONSE_HEADERS = {
+        "Content-Length": str(CONTENT_LENGTH),
+        "Shard-Meta": '{"checksum": "bW9jay1jaGVja3N1bQ==", "algorithm": "xxh3_128"}'
+    }
+
+    def setup_method(self):
+        super().setup_method()
+        self.checkpoint_path = f"v1/cp/checkpoints/{NAMESPACE}/{RANK}/{self.STEP}"
+
+        # Mock checksum functions
+        self.mock_decode_base_64 = patch(
+            "amzn_sagemaker_checkpointing.storage.clients.inmemory.inmemory_client.decode_base_64"
+        ).start()
+        self.mock_hash_xxh3_128 = patch(
+            "amzn_sagemaker_checkpointing.storage.clients.inmemory.inmemory_client.hash_xxh3_128"
+        ).start()
+        self.mock_decode_base_64.return_value = self.MOCK_CHECKSUM
+        self.mock_hash_xxh3_128.return_value = self.MOCK_CHECKSUM
+
+    def test_get_checkpoint_success(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.content = self.TEST_CONTENT
+        mock_response.headers = self.RESPONSE_HEADERS
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        result = self.client.get_checkpoint(step=self.STEP)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_hash_xxh3_128.assert_called_once_with(self.TEST_CONTENT)
+        self.mock_decode_base_64.assert_called_once_with(self.MOCK_ENCODED_CHECKSUM)
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{self.checkpoint_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+        assert result == self.TEST_CONTENT
+
+    def test_get_checkpoint_with_custom_rank(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.content = self.TEST_CONTENT
+        mock_response.headers = self.RESPONSE_HEADERS
+        self.mock_session.request.return_value = mock_response
+        custom_rank = 5
+        custom_path = f"v1/cp/checkpoints/{NAMESPACE}/{custom_rank}/{self.STEP}"
+
+        # Act
+        result = self.client.get_checkpoint(step=self.STEP, rank=custom_rank)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{custom_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+        assert result == self.TEST_CONTENT
+
+    def test_get_checkpoint_with_metadata_index(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.content = self.TEST_CONTENT
+        mock_response.headers = self.RESPONSE_HEADERS
+        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
+        result = self.client.get_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="GET",
+            url=f"{BASE_URL}/{metadata_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+        assert result == self.TEST_CONTENT
+
+    def test_get_checkpoint_with_custom_timeout(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.content = self.TEST_CONTENT
+        mock_response.headers = self.RESPONSE_HEADERS
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        result = self.client.get_checkpoint(step=self.STEP, timeout=REQUEST_TIMEOUT)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{self.checkpoint_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=REQUEST_TIMEOUT
+        )
+        assert result == self.TEST_CONTENT
+
+    def test_get_checkpoint_not_found(self):
+        # Arrange
+        mock_response = Mock(status_code=404)
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        result = self.client.get_checkpoint(step=self.STEP)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{self.checkpoint_path}",
+            params=None,
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+        assert result is None
+
+    def test_get_checkpoint_checksum_mismatch(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.content = self.TEST_CONTENT
+        mock_response.headers = self.RESPONSE_HEADERS
+        self.mock_session.request.return_value = mock_response
+        self.mock_hash_xxh3_128.return_value = "different-checksum"
+
+        # Act & Assert
+        with pytest.raises(InMemoryStorageError) as exc_info:
+            self.client.get_checkpoint(step=self.STEP)
+        assert "Checksum mismatch in response" in str(exc_info.value)
+
+    def test_get_checkpoint_content_length_mismatch(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.content = self.TEST_CONTENT
+        mock_response.headers = {
+            "Content-Length": str(self.CONTENT_LENGTH + 1),  # Wrong content length
+            "Shard-Meta": '{"checksum": "bW9jay1jaGVja3N1bQ==", "algorithm": "xxh3_128"}'
+        }
+        self.mock_session.request.return_value = mock_response
+
+        # Act & Assert
+        with pytest.raises(InMemoryStorageError) as exc_info:
+            self.client.get_checkpoint(step=self.STEP)
+        assert "Content length mismatch" in str(exc_info.value)
+
+    def test_get_checkpoint_invalid_metadata_index(self):
+        # Act & Assert
+        with pytest.raises(InMemoryConfigError) as exc_info:
+            self.client.get_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_get_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.get_checkpoint,
+            step=self.STEP
+        )

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

@@ -0,0 +1,116 @@
+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 TestGetLatestCheckpoints(InMemoryCheckpointClientTest):
+    def setup_method(self):
+        super().setup_method()
+        self.checkpoints_path = f"v1/cp/checkpoints/{NAMESPACE}/latest"
+
+    @pytest.mark.parametrize(
+        "params, api_steps, expected_steps",
+        [
+            # Default limit, ordered steps
+            (
+                {"limit": 5, "timeout": None},
+                ["10", "8", "6", "4", "2"],
+                [10, 8, 6, 4, 2]
+            ),
+            # Custom limit
+            (
+                {"limit": 2, "timeout": None},
+                ["10", "8"],
+                [10, 8]
+            ),
+            # Custom timeout
+            (
+                {"limit": 5, "timeout": REQUEST_TIMEOUT},
+                ["10", "8", "6", "4", "2"],
+                [10, 8, 6, 4, 2]
+            ),
+            # Unordered steps
+            (
+                {"limit": 5, "timeout": None},
+                ["4", "8", "2", "10", "6"],
+                [10, 8, 6, 4, 2]
+            ),
+        ],
+    )
+    def test_get_latest_checkpoints_success(self, params, api_steps, expected_steps):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.return_value = api_steps
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        result = self.client.get_latest_checkpoints(**params)
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{self.checkpoints_path}",
+            params={"limit": params["limit"]},
+            data=None,
+            headers=None,
+            timeout=params["timeout"] or InMemoryClientConfig.request_timeout
+        )
+        assert result == expected_steps
+
+    @pytest.mark.parametrize("test_case", REQUEST_ERROR_CASES)
+    def test_get_latest_checkpoints_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
+        result = self.client.get_latest_checkpoints()
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{self.checkpoints_path}",
+            params={"limit": 5},
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+        assert result == []
+
+    def test_get_latest_checkpoints_json_decode_error(self):
+        # Arrange
+        mock_response = Mock()
+        mock_response.status_code = 200
+        mock_response.json.side_effect = ValueError("Invalid JSON")
+        self.mock_session.request.return_value = mock_response
+
+        # Act
+        result = self.client.get_latest_checkpoints()
+
+        # Assert
+        self.assert_http_adapter_and_retry_config()
+        self.mock_session.request.assert_called_once_with(
+            method="GET",
+            url=f"{BASE_URL}/{self.checkpoints_path}",
+            params={"limit": 5},
+            data=None,
+            headers=None,
+            timeout=InMemoryClientConfig.request_timeout
+        )
+        assert result == []