feat(dcp): Add thread_count parameter to S3StorageWriter (#370)

* The thread_count parameter is used and defaulted to 1 within
FileSystemWriter (parent class). This change is to expose this
parameter to our users through our S3StorageWriter.

We currently use it already within e2e filesystem integration 
tests and DCP benchmarks. We have extended our unit 
tests and added docstrings to give better visibility to this 
parameter.

Author: MSH4R1F

CHANGELOG.md

@@ -5,6 +5,9 @@
 * Override S3Writer closed property and block writes after close (#360)
 * Fix SequentialS3Reader seek beyond EOF to clamp position to object size (#362)
 
+### Other changes
+* Added thread_count parameter to S3StorageWriter 
+
 ## v1.4.3 (July 25, 2025)
 
 ### New features

README.md

@@ -162,7 +162,11 @@ REGION = "us-east-1"
 model = torchvision.models.resnet18()
 
 # Save distributed checkpoint to S3
-s3_storage_writer = S3StorageWriter(region=REGION, path=CHECKPOINT_URI)
+s3_storage_writer = S3StorageWriter(
+    region=REGION, 
+    path=CHECKPOINT_URI,
+    thread_count=8, # optional; number of IO threads to use to write
+) 
 DCP.save(
     state_dict=model.state_dict(),
     storage_writer=s3_storage_writer,

s3torchconnector/src/s3torchconnector/dcp/s3_file_system.py

@@ -271,6 +271,7 @@ def __init__(
         path: str,
         s3client_config: Optional[S3ClientConfig] = None,
         prefix_strategy: Optional[S3PrefixStrategyBase] = None,
+        thread_count: int = 1,
         **kwargs,
     ) -> None:
         """
@@ -282,11 +283,13 @@ def __init__(
             s3client_config (Optional[S3ClientConfig]): Optional S3ClientConfig with parameters for S3 client.
             prefix_strategy (Optional[S3PrefixStrategyBase]): Optional strategy for generating S3 prefixes to
                 optimize checkpoint organization and prevent throttling.
+            thread_count (int): Number of IO threads to use to write. Defaults to 1 (Pytorch Default)
             kwargs (dict): Keyword arguments to pass to the parent :class:`FileSystemWriter`.
         """
         super().__init__(
             path=path,
             sync_files=False,  # FIXME: setting this to True makes the run to fail (L#333: `os.fsync(stream.fileno())`)
+            thread_count=thread_count,
             **kwargs,
         )
         self.fs = S3FileSystem(region, s3client_config=s3client_config)  # type: ignore

s3torchconnector/tst/unit/dcp/test_s3_storage_writer.py

@@ -0,0 +1,23 @@
+#  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#  // SPDX-License-Identifier: BSD
+
+import pytest
+from s3torchconnector.dcp import S3StorageWriter
+
+TEST_REGION = "eu-east-1"
+TEST_BUCKET = "test-bucket"
+TEST_KEY = "test-key.txt"
+TEST_PATH = f"s3://{TEST_BUCKET}/{TEST_KEY}"
+
+
+@pytest.mark.parametrize("thread_count", [1, 2, 4, 8, 16])
+def test_s3storage_writer_thread_count(thread_count):
+    storage_writer = S3StorageWriter(
+        region=TEST_REGION, path=TEST_PATH, thread_count=thread_count
+    )
+    assert storage_writer.thread_count == thread_count
+
+
+def test_s3storage_writer_thread_count_defaults_to_one():
+    storage_writer = S3StorageWriter(region=TEST_REGION, path=TEST_PATH)
+    assert storage_writer.thread_count == 1