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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: amzn-sagemaker-checkpointing
-Version: 1.0.10
+Version: 1.0.11
 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.11}/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.11}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "amzn-sagemaker-checkpointing"
-version = "1.0.10"
+version = "1.0.11"
 description = "Amazon SageMaker Checkpointing Library"
 readme = "README.md"
 license = { "text" = "Apache 2.0" }

{amzn_sagemaker_checkpointing-1.0.10 → amzn_sagemaker_checkpointing-1.0.11}/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