Resolve comments.

Author: Leahlijuan

src/ml_flashpoint/checkpoint_object_manager/buffer_metadata.py

@@ -21,10 +21,10 @@
 class BufferMetadataType(ctypes.LittleEndianStructure):
     """Metadata block stored at the beginning of the BufferIO buffer."""
 
-    _pack_ = 1
+    _pack_ = 1  # Ensure tight packing for cross-platform consistency
     _fields_ = [
         ("len_written_data", ctypes.c_uint64),
-        # 8 bytes for MAGIC_BYTES to identify the file format version
+        # 8 bytes for checkpoint format signature to identify the file format version
         ("format_signature", ctypes.c_char * 8),
         # Pad the rest of the structure to reach METADATA_SIZE
         (

src/ml_flashpoint/core/checkpoint_loader.py

@@ -42,8 +42,8 @@
 from ml_flashpoint.core.defaults import (
     COMMON_STATE_FNAME,
     DIRTY_MARKER_SUFFIX,
-    FORMAT_SIGNATURE,
     GLOBAL_RANK_PATTERN,
+    CheckpointFormat,
     default_metadata_object_name,
 )
 from ml_flashpoint.core.mlf_logging import get_logger
@@ -301,7 +301,7 @@ def read_data(
 
         with self._checkpoint_object_manager.get_buffer(checkpoint_object_id) as stream:
             use_optimized_loader = False
-            if stream.format_signature == FORMAT_SIGNATURE:
+            if stream.format_signature == CheckpointFormat.MLF_FORMAT:
                 use_optimized_loader = True
                 _LOGGER.debug("Using optimized loader for '%s'", checkpoint_object_id.data)
 

src/ml_flashpoint/core/checkpoint_saver.py

@@ -625,9 +625,9 @@ def _write_to_buffer_from_queue_worker(
                 ) as buffer_io_writer:
                     # Set the format signature
                     if use_optimized_write:
-                        buffer_io_writer.set_format_signature(CheckpointFormat.MLF_FORMAT.value)
+                        buffer_io_writer.set_format_signature(CheckpointFormat.MLF_FORMAT)
                     else:
-                        buffer_io_writer.set_format_signature(CheckpointFormat.TORCH_SAVE.value)
+                        buffer_io_writer.set_format_signature(CheckpointFormat.TORCH_SAVE)
 
                     # First write tensors.
                     for tensor_item, tensor in tensor_tuples:

src/ml_flashpoint/core/defaults.py

@@ -27,9 +27,6 @@ class CheckpointFormat(bytes, Enum):
     MLF_FORMAT = b"MLF_TENS"
 
 
-FORMAT_SIGNATURE = CheckpointFormat.MLF_FORMAT.value
-
-
 def default_metadata_object_name() -> str:
     """Returns the default object name for metadata files (i.e. filename).
 

tests/checkpoint_object_manager/test_buffer_io.py

@@ -24,7 +24,7 @@
 from ml_flashpoint.checkpoint_object_manager.buffer_io import METADATA_SIZE, BufferIO
 from ml_flashpoint.checkpoint_object_manager.buffer_metadata import BufferMetadataType
 from ml_flashpoint.checkpoint_object_manager.buffer_object.buffer_object_ext import BufferObject
-from ml_flashpoint.core.defaults import FORMAT_SIGNATURE
+from ml_flashpoint.core.defaults import CheckpointFormat
 
 
 # --- Mocks and Fixtures ---
@@ -994,7 +994,7 @@ def test_buffer_io_does_not_set_signature_automatically(self, temp_dir_path):
             # BufferIO should NOT auto-set signature.
             sig = bio.format_signature
 
-            assert sig != FORMAT_SIGNATURE, f"BufferIO SHOULD NOT auto-set signature. Got {sig!r}"
+            assert sig != CheckpointFormat.MLF_FORMAT, f"BufferIO SHOULD NOT auto-set signature. Got {sig!r}"
 
     def test_set_format_signature(self, temp_dir_path):
         """Test that set_format_signature correctly updates the metadata."""

tests/core/test_checkpoint_saver.py

@@ -39,7 +39,7 @@
 from ml_flashpoint.checkpoint_object_manager.object_manager import object_manager_ext
 from ml_flashpoint.core.checkpoint_id_types import CheckpointContainerId, CheckpointObjectId
 from ml_flashpoint.core.checkpoint_saver import DefaultMLFlashpointCheckpointSaver, WriteItemResolver
-from ml_flashpoint.core.defaults import FORMAT_SIGNATURE
+from ml_flashpoint.core.defaults import CheckpointFormat
 from ml_flashpoint.replication.replication_manager import ReplicationManager
 
 
@@ -202,7 +202,7 @@ def test_write_data_with_optimized_save_false(self, chkpt_object_manager, replic
         assert buffer_io is not None
         with buffer_io:
             magic = buffer_io.read(8)
-            assert magic != FORMAT_SIGNATURE
+            assert magic != CheckpointFormat.MLF_FORMAT
 
             # Reset position for loading
             buffer_io.seek(0)