GeezerVison

transfer_queue/utils/utils.py (1)

106-106: Consider expanding the docstring.

The docstring is quite brief and doesn't describe parameters, return behavior, or usage examples. Consider adding more detail about when and why this should be used.

Example enhancement:

-    """Prevent PyTorch from overdoing the automatic parallelism during torch.stack() operation"""
+    """
+    Temporarily limit PyTorch's thread count to prevent excessive parallelism.
+    
+    Useful during operations like torch.stack() where PyTorch's automatic 
+    parallelism may spawn too many threads on high-core-count machines.
+    
+    Args:
+        target_num_threads: Maximum threads to use. If None, auto-determined 
+                          based on physical cores (16 if ≥16 cores, else physical core count).
+    
+    Yields:
+        None
+    
+    Raises:
+        RuntimeError: If target_num_threads exceeds logical CPU cores.
+    """

transfer_queue/storage/clients/factory.py (1)

1-13: Placeholder module: add a minimal stub or remove until used

This file only contains a license header. Either remove it for now or add a TODO and a minimal Factory interface to avoid confusion during discovery.

transfer_queue/utils/zmq_utils.py (4)

78-85: Ensure JSON-serializable role in to_dict

role is an enum; return role.value (or str(role)) to make this dict safe for JSON/log serialization.

-        return {
-            "role": self.role,
+        return {
+            "role": self.role.value,
             "id": self.id,
             "ip": self.ip,
             "ports": self.ports,
         }

---

86-88: Typo: class name in str

The string says ZMQSocketInfo; class is ZMQServerInfo. Fix for clarity.

-        return f"ZMQSocketInfo(role={self.role}, id={self.id}, ip={self.ip}, ports={self.ports})"
+        return f"ZMQServerInfo(role={self.role}, id={self.id}, ip={self.ip}, ports={self.ports})"

---

112-114: Short request_id (8 hex) increases collision risk under load

8 hex chars (~32 bits) can collide in busy systems. Prefer 16–32 hex for safety.

-            request_id=str(uuid4().hex[:8]),
+            request_id=uuid4().hex[:16],  # or omit slicing

---

69-76: Align get_free_port return type with ports typing

get_free_port returns an int but is annotated as str; fix typing and minor naming nits.

-def get_free_port() -> str:
-    with socket.socket() as sock:
-        sock.bind(("", 0))
-        return sock.getsockname()[1]
+def get_free_port() -> int:
+    with socket.socket() as s:
+        s.bind(("", 0))
+        return int(s.getsockname()[1])

Also ensure ZMQServerInfo.ports uses dict[str, int] (see prior comment).

Also applies to: 134-137

transfer_queue/metadata.py (4)

98-123: union() returns self but doc promises a new object/None

Doc says “New SampleMeta (None if validation fails)”; implementation mutates and returns self, and raises on validation failure. Align doc with behavior.

-        Returns:
-            New SampleMeta with unioned fields (None if validation fails)
+        Returns:
+            SampleMeta (self) with unioned fields. Raises ValueError on validation failure.

---

153-160: field_names cache may not represent “all unique field names”

You cache field names from the first sample but the property doc says “all unique field names”. Either compute the union across samples or update the doc.

Possible fix:

-            object.__setattr__(self, "_field_names", sorted(self.samples[0].field_names))
+            union = sorted({name for s in self.samples for name in s.field_names})
+            object.__setattr__(self, "_field_names", union)

Also applies to: 171-175

---

295-299: extra_info merge semantics differ from BatchMeta.union

concat uses ChainMap (first-wins), union uses dict unpacking (last-wins). Pick one policy for consistency; recommend last-wins.

-        merged_extra_info = dict(ChainMap(*(chunk.extra_info for chunk in data)))
+        merged_extra_info = {}
+        for chunk in data:
+            merged_extra_info |= chunk.extra_info  # last-wins

---

441-455: Prefer len(tensor_dict) for batch size; minor clarity

batch_size = tensor_dict.batch_size[0] works but len(tensor_dict) is clearer and resilient.

-    batch_size = tensor_dict.batch_size[0]
+    batch_size = len(tensor_dict)

.github/workflows/python-package.yml (1)

30-31: Use canonical package name pytest-asyncio

The canonical PyPI project name is "pytest-asyncio" (dash), even though pip normalizes both forms. Switch to the official name for clarity.

-        python -m pip install flake8 pytest build pytest_asyncio
+        python -m pip install flake8 pytest build pytest-asyncio

transfer_queue/storage/managers/__init__.py (1)

15-17: Prefer lazy re‑exports to reduce import-time cost and cycle risk

Replace eager imports with on-demand attribute loading via getattr.

Apply:

+from importlib import import_module
+
-from .base import TransferQueueStorageManager
-from .factory import TransferQueueStorageManagerFactory
-from .simple_backend_manager import AsyncSimpleStorageManager
+_EXPORTS = {
+    "TransferQueueStorageManager": (".base", "TransferQueueStorageManager"),
+    "TransferQueueStorageManagerFactory": (".factory", "TransferQueueStorageManagerFactory"),
+    "AsyncSimpleStorageManager": (".simple_backend_manager", "AsyncSimpleStorageManager"),
+}
+
+def __getattr__(name: str):
+    if name in _EXPORTS:
+        mod_path, attr = _EXPORTS[name]
+        return getattr(import_module(mod_path, __package__), attr)
+    raise AttributeError(name)

transfer_queue/storage/__init__.py (1)

15-17: Consider lazy re‑exports for the storage package, mirroring managers/init.py

Helps avoid importing heavy deps at module import time.

README.md (4)

55-56: Align terminology with code (“StorageManager” + “SimpleStorageUnit”)

Replace “written to the TransferQueueStorage” with “written to the storage backend (e.g., SimpleStorageUnit via TransferQueueStorageManager)” for consistency.

---

24-25: Add alt text to images for accessibility

Provide descriptive alt attributes for all tags.

Example:

-<img src=".../1761356010763-....png" width="70%">
+<img src=".../1761356010763-....png" width="70%" alt="TransferQueue architecture overview">

Also applies to: 32-33, 61-62, 125-126, 144-146, 178-179

---

152-174: Document test extras in Quick Start

Since pyproject defines [project.optional-dependencies].test, add:

• pip install -e ".[test]"
• pytest -q

---

10-13: Review third‑party embed

The DeepWiki badge references devin.ai assets; consider replacing with a static badge or project-hosted asset to avoid external tracking/fragility.

pyproject.toml (3)

83-86: Set pytest asyncio mode to avoid deprecation warnings

Add asyncio_mode = "auto".

 [tool.pytest.ini_options]
 filterwarnings = [
     "ignore:.*PyTorch API of nested tensors.*prototype.*:UserWarning",
 ]
+asyncio_mode = "auto"

---

71-79: Avoid blanket mypy ignore_errors

Prefer targeted silences; set ignore_errors = false and use per-module overrides.

-[tool.mypy]
+[tool.mypy]
 pretty            = true
 ignore_missing_imports = true
 explicit_package_bases = true
 follow_imports = "skip"
-
-# Blanket silence
-ignore_errors = true
+ignore_errors = false

---

110-114: Expose test extras in docs and CI

Leverage the new test extra: pip install -e ".[test]".

tests/test_async_simple_storage_manager.py (6)

25-33: Avoid sys.path hacks in tests

Rely on editable install or pytest config. Suggest pip install -e ".[test]" and remove manual path injection.

-parent_dir = Path(__file__).resolve().parent.parent
-sys.path.append(str(parent_dir))
+# Prefer: install package in editable mode for tests or configure pytest to run from repo root.

---

35-89: Prefer constructing the class normally vs new

Using new risks partial initialization. Patch handshake like in other tests and call AsyncSimpleStorageManager(config).

---

198-227: Add resource cleanup to avoid lingering ZMQ contexts

Call manager.close() in a finally or use addfinalizer.

-        manager = AsyncSimpleStorageManager(config)
+        manager = AsyncSimpleStorageManager(config)
+        try:
+            # assertions...
+            assert manager.global_index_storage_unit_mapping(0) == "storage_0"
+            # ...
+        finally:
+            manager.close()

---

246-301: Also close manager in error-handling test

Avoid leaking sockets/contexts across tests.

-        manager = AsyncSimpleStorageManager(config)
+        manager = AsyncSimpleStorageManager(config)
+        try:
+            # mocked operations and assertions...
+            pass
+        finally:
+            manager.close()

---

147-150: Strengthen assertions

Assert notify_data_update arguments and that _put/_get/_clear were called with expected shard targets.

Example:

manager.notify_data_update.assert_awaited_once()
manager._put_to_single_storage_unit.assert_awaited()
calls = [c.kwargs.get("target_storage_unit") for c in manager._put_to_single_storage_unit.await_args_list]
assert set(calls) == {"storage_0", "storage_1"}

Also applies to: 151-157

---

1-223: Sanity-check: pickle over ZMQ

ZMQMessage uses pickle serialization; ensure tests cover round-trip for nontrivial payloads and consider documenting the security trade-offs.

tests/test_controller.py (3)

26-27: Avoid importing internals in tests unless necessary

Importing TQ_INIT_FIELD_NUM from transfer_queue.controller couples tests to internals. Prefer asserting behaviors via public APIs or expose a stable constant through the package root if needed.

If TQ_INIT_FIELD_NUM must be asserted, consider re-exporting it in transfer_queue/init.py.

---

97-129: Add invariant checks to strengthen metadata test

The reorder/global_indexes assertions look right. Add simple invariants to catch regressions (size and field_names unchanged by reorder).

     metadata = ray.get(
         tq_controller._get_metadata.remote(
             data_fields=data_fields,
             batch_size=global_batch_size * n_samples,
             global_step=global_step,
             mode="insert",
         )
     )
+    # Invariants
+    assert metadata.size == global_batch_size * n_samples
+    assert metadata.field_names == sorted(data_fields)
     metadata.reorder([15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0])

---

65-84: Tests rely on private controller methods—consider adding public accessors or test helper methods

Both _update_production_status (line 78) and _get_consumption_status (line 93) are internal methods that lack public alternatives. While the .remote() calls are valid (TransferQueueController is properly decorated with @ray.remote), testing through private methods risks breakage during future refactors. Either add public wrapper methods for testing, create dedicated test helper methods in the controller, or add comments explaining why private access is necessary.

transfer_queue/storage/managers/factory.py (1)

25-41: Harden factory: normalize keys, guard duplicates, and default to class name

Make manager_type case-insensitive, prevent accidental overwrite, and use the class name by default during registration.

@@
-    def register(cls, manager_type: str, manager_cls: type[TransferQueueStorageManager]):
+    def register(cls, manager_type: str | None, manager_cls: type[TransferQueueStorageManager]):
         if not issubclass(manager_cls, TransferQueueStorageManager):
             raise TypeError(
                 f"manager_cls {getattr(manager_cls, '__name__', repr(manager_cls))} must be "
                 f"a subclass of TransferQueueStorageManager"
             )
-        cls._registry[manager_type] = manager_cls
+        key = (manager_type or manager_cls.__name__).lower()
+        if key in cls._registry and cls._registry[key] is not manager_cls:
+            raise ValueError(f"Duplicate registration for manager_type '{key}'")
+        cls._registry[key] = manager_cls
@@
-    def create(cls, manager_type: str, config: dict[str, Any]) -> TransferQueueStorageManager:
-        if manager_type not in cls._registry:
+    def create(cls, manager_type: str, config: dict[str, Any]) -> TransferQueueStorageManager:
+        key = manager_type.lower()
+        if key not in cls._registry:
             raise ValueError(
-                f"Unknown manager_type: {manager_type}. Supported managers include: {list(cls._registry.keys())}"
+                f"Unknown manager_type: {manager_type}. Supported: {list(cls._registry.keys())}"
             )
-        return cls._registry[manager_type](config)
+        return cls._registry[key](config)
@@
-TransferQueueStorageManagerFactory.register("AsyncSimpleStorageManager", AsyncSimpleStorageManager)
+TransferQueueStorageManagerFactory.register(None, AsyncSimpleStorageManager)

Also applies to: 43-44

tests/test_simple_storage_unit.py (3)

34-51: Set SNDTIMEO and client identity for stability and debuggability

Avoid indefinite send on backpressure and add an identity to ease ROUTER-side tracing.

 class MockStorageClient:
@@
-    def __init__(self, storage_put_get_address):
+    def __init__(self, storage_put_get_address, identity: str | None = None):
         self.context = zmq.Context()
-        self.socket = self.context.socket(zmq.DEALER)
-        self.socket.setsockopt(zmq.RCVTIMEO, 5000)  # 5 second timeout
+        self.socket = self.context.socket(zmq.DEALER)
+        if identity:
+            self.socket.setsockopt(zmq.IDENTITY, identity.encode())
+        self.socket.setsockopt(zmq.RCVTIMEO, 5000)  # 5s
+        self.socket.setsockopt(zmq.SNDTIMEO, 5000)  # 5s
         self.socket.connect(storage_put_get_address)

---

92-99: Replace fixed sleep with readiness probe to reduce flakiness

Actively probe the ROUTER by issuing a harmless CLEAR with empty indexes; retry a few times instead of sleeping.

-    put_get_address = zmq_info.to_addr("put_get_socket")
-    time.sleep(1)  # Wait for socket to be ready
+    put_get_address = zmq_info.to_addr("put_get_socket")
+    # Probe readiness: send a no-op CLEAR ([]) with short timeout
+    ready = False
+    for attempt in range(10):
+        try:
+            client = MockStorageClient(put_get_address, identity=f"probe_{attempt}")
+            resp = client.send_clear("probe", [])
+            # Accept either CLEAR_DATA_RESPONSE or CLEAR_DATA_ERROR (server up)
+            if resp.request_type in {ZMQRequestType.CLEAR_DATA_RESPONSE, ZMQRequestType.CLEAR_DATA_ERROR}:
+                ready = True
+                client.close()
+                break
+        except Exception:
+            time.sleep(0.1)
+    assert ready, "StorageUnit put_get_socket not ready in time"

---

202-256: Gate performance test to avoid CI flakiness

Mark as slow and allow opt-in via env var to keep PR signal clean.

-def test_performance_basic(storage_setup):
+@pytest.mark.slow
+def test_performance_basic(storage_setup):
@@
-    assert avg_put_latency < 1500, f"Avg PUT latency {avg_put_latency}ms exceeds threshold"
-    assert avg_get_latency < 1500, f"Avg GET latency {avg_get_latency}ms exceeds threshold"
+    if os.getenv("TQ_RUN_PERF_TESTS", "0") == "1":
+        assert avg_put_latency < 1500, f"Avg PUT latency {avg_put_latency}ms exceeds threshold"
+        assert avg_get_latency < 1500, f"Avg GET latency {avg_get_latency}ms exceeds threshold"

transfer_queue/storage/managers/base.py (2)

65-76: Set DEALER socket timeouts to avoid accidental indefinite blocks

Add conservative SNDTIMEO/RCVTIMEO on controller_handshake_socket and data_status_update_socket.

             self.controller_handshake_socket = create_zmq_socket(
                 self.zmq_context,
                 zmq.DEALER,
                 identity=f"{self.storage_manager_id}-controller_handshake_socket-{uuid4().hex[:8]}".encode(),
             )
+            self.controller_handshake_socket.setsockopt(zmq.RCVTIMEO, 10_000)
+            self.controller_handshake_socket.setsockopt(zmq.SNDTIMEO, 10_000)
             self.data_status_update_socket = create_zmq_socket(
                 self.zmq_context,
                 zmq.DEALER,
                 identity=f"{self.storage_manager_id}-data_status_update_socket-{uuid4().hex[:8]}".encode(),
             )
+            self.data_status_update_socket.setsockopt(zmq.RCVTIMEO, 10_000)
+            self.data_status_update_socket.setsockopt(zmq.SNDTIMEO, 10_000)
             self.data_status_update_socket.connect(self.controller_info.to_addr("data_status_update_socket"))

---

218-230: Don’t wait for ACK after emitting NOTIFY_DATA_UPDATE_ERROR

If sending the update fails, you emit an error message and still wait for an ACK, which will likely never arrive. Return early.

         except Exception as e:
             request_msg = ZMQMessage.create(
                 request_type=ZMQRequestType.NOTIFY_DATA_UPDATE_ERROR,
                 sender_id=self.storage_manager_id,
                 body={
                     "message": f"Failed to notify data status update information from "
                     f"storage manager id #{self.storage_manager_id}, "
                     f"detail error message: {str(e)}"
                 },
             ).serialize()
 
             self.data_status_update_socket.send(request_msg)
+            return  # Avoid waiting for ACK after error
 
         # Make sure controller successfully receives data status update information.

Also applies to: 252-257

transfer_queue/storage/managers/simple_backend_manager.py (1)

322-325: Reduce log verbosity for GET responses

Logging entire response bodies at info can be noisy and expensive; use debug or log minimal identifiers.

-            logger.info(
-                f"[{self.storage_manager_id}]: get data response from storage unit "
-                f"{target_storage_unit}: {response_msg}"
-            )
+            logger.debug(
+                "[%s]: GET response from %s (type=%s, size=%s)",
+                self.storage_manager_id,
+                target_storage_unit,
+                getattr(response_msg, "request_type", None),
+                len(getattr(response_msg.body, "get", lambda *_: [])("data", {})) if hasattr(response_msg, "body") else None,
+            )

transfer_queue/storage/simple_backend.py (3)

178-185: Add graceful shutdown; set socket LINGER and close Context.

Provide close() to avoid resource leaks and hanging ROUTER sockets on actor teardown.

@@ class SimpleStorageUnit:
     def _init_zmq_socket(self) -> None:
@@
         self.put_get_socket = create_zmq_socket(self.zmq_context, zmq.ROUTER)
         self.put_get_socket.bind(self.zmq_server_info.to_addr("put_get_socket"))
+
+    def close(self) -> None:
+        try:
+            if hasattr(self, "put_get_socket") and not self.put_get_socket.closed:
+                self.put_get_socket.setsockopt(zmq.LINGER, 0)
+                self.put_get_socket.close()
+        finally:
+            if hasattr(self, "zmq_context"):
+                self.zmq_context.term()
@@
     def get_zmq_server_info(self) -> ZMQServerInfo:

Also applies to: 195-197, 345-352

---

178-183: Port type inconsistency.

ZMQServerInfo.ports is typed dict[str, str], but get_free_port returns int; you pass ints. Either cast to str here or update the annotation to int.

-            ports={"put_get_socket": get_free_port()},
+            ports={"put_get_socket": str(get_free_port())},

---

382-389: Deterministic field order for reproducibility.

get_field_names returns an arbitrary set order. Sort to stabilize downstream behavior and tests.

-        return list(all_fields)
+        return sorted(all_fields)

transfer_queue/controller.py (1)

512-521: Ports typing consistency.

As with storage, cast port ints to str or change ZMQServerInfo annotation.

-                "handshake_socket": self._handshake_socket_port,
-                "request_handle_socket": self._request_handle_socket_port,
-                "data_status_update_socket": self._data_status_update_socket_port,
+                "handshake_socket": str(self._handshake_socket_port),
+                "request_handle_socket": str(self._request_handle_socket_port),
+                "data_status_update_socket": str(self._data_status_update_socket_port),

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

transfer_queue/storage/simple_backend.py (4)

173-173: UUID collision risk with truncated hex.

Using only 8 hex characters from UUID (uuid4().hex[:8]) creates collision risk in deployments with many storage units (birthday paradox: ~1% chance at ~300 units).

Use the full UUID hex or at least 16 characters:

-        self.storage_unit_id = f"TQ_STORAGE_UNIT_{uuid4().hex[:8]}"
+        self.storage_unit_id = f"TQ_STORAGE_UNIT_{uuid4().hex}"

---

86-103: Validate local_indexes and handle None slots defensively.

The get_data method doesn't validate that local_indexes are within bounds before accessing self.field_data[field][local_indexes[0]] or using itemgetter. This can raise IndexError for out-of-range indexes. Additionally, if storage slots contain None (e.g., after a partial clear or before data is written), itemgetter will return None values that can break stacking operations.

Apply this diff to add bounds checking and filter out None values:

 def get_data(self, fields: list[str], local_indexes: list[int]) -> TensorDict[str, list]:
     """
     Get data from storage unit according to given fields and local_indexes.
     """
     result: dict[str, list] = {}
+    
+    # Validate indexes once upfront
+    for idx in local_indexes:
+        if idx < 0 or idx >= self.storage_size:
+            raise ValueError(
+                f"StorageUnitData get_data operation receive invalid local_index: {idx} beyond "
+                f"storage_size: {self.storage_size}"
+            )
 
     for field in fields:
         # Validate field name
         if field not in self.field_data:
             raise ValueError(
                 f"StorageUnitData get_data operation receive invalid field: {field} beyond {self.field_data.keys()}"
             )
 
         if len(local_indexes) == 1:
             gathered_item = self.field_data[field][local_indexes[0]]
+            if gathered_item is None:
+                raise ValueError(f"Field '{field}' at local_index {local_indexes[0]} is empty/None")
             if not isinstance(gathered_item, torch.Tensor):
                 result[field] = NonTensorStack(gathered_item)
             else:
                 result[field] = gathered_item.unsqueeze(0)
         else:
             gathered_items = list(itemgetter(*local_indexes)(self.field_data[field]))
+            
+            # Check for None values
+            if any(x is None for x in gathered_items):
+                none_indexes = [idx for idx, x in zip(local_indexes, gathered_items) if x is None]
+                raise ValueError(f"Field '{field}' contains empty slots at indexes: {none_indexes}")
 
             if gathered_items:
                 all_tensors = all(isinstance(x, torch.Tensor) for x in gathered_items)
                 if all_tensors:
                     result[field] = torch.nested.as_nested_tensor(gathered_items)
                 else:
                     result[field] = NonTensorStack(*gathered_items)
 
     return TensorDict(result)

---

115-127: Prevent length mismatch and avoid enumerate-coupling in put_data.

When len(values) ≠ len(local_indexes), the current code using enumerate(local_indexes) and indexing values[i] will either raise IndexError or write incorrect data. The bounds check happens per-index but after the mismatch has already caused issues.

Apply this diff to validate lengths upfront and use zip:

 def put_data(self, field_data: TensorDict[str, Any], local_indexes: list[int]) -> None:
     """
     Put or update data into storage unit according to given field_data and local_indexes.
     """
     extracted_data = dict(field_data)
 
     for f, values in extracted_data.items():
         if f not in self.field_data:
             self.field_data[f] = [None] * self.storage_size
+        
+        if len(values) != len(local_indexes):
+            raise ValueError(
+                f"put_data length mismatch for field '{f}': "
+                f"{len(values)} values for {len(local_indexes)} indexes"
+            )
 
-        for i, idx in enumerate(local_indexes):
+        for idx, value in zip(local_indexes, values):
             if idx < 0 or idx >= self.storage_size:
                 raise ValueError(
                     f"StorageUnitData put_data operation receive invalid local_index: {idx} beyond "
                     f"storage_size: {self.storage_size}"
                 )
 
-            self.field_data[f][idx] = values[i]
+            self.field_data[f][idx] = value

---

216-246: Unsafe network deserialization via pickle.

Deserializing pickled messages from sockets (line 217: ZMQMessage.deserialize(serialized_msg)) enables arbitrary code execution if an attacker can connect to the storage unit's ZMQ endpoint. This is a known Python security issue with pickle over the network.

Consider these mitigations:

1. Preferred: Replace pickle with a safe serialization format (msgpack, JSON + schema validation)
2. If pickle unavoidable:
   • Restrict bind addresses to trusted networks only
   • Add ZMQ CURVE authentication
   • Implement message authentication (HMAC)
   • Document the security model and deployment restrictions

Also see: Python pickle security documentation

transfer_queue/storage/managers/simple_backend_manager.py (2)

120-128: Missing port validation when building address.

Line 126 builds the address using server_info.ports.get(socket_name) without verifying the port exists. If the port is missing, this produces tcp://127.0.0.1:None, leading to obscure connection errors.

Add explicit validation:

             server_info = self.storage_unit_infos.get(server_key)
 
             if not server_info:
                 raise RuntimeError(f"Server {server_key} not found in registered servers")
 
             context = zmq.asyncio.Context()
-            address = f"tcp://{server_info.ip}:{server_info.ports.get(socket_name)}"
+            port = server_info.ports.get(socket_name)
+            if port is None:
+                raise ValueError(
+                    f"StorageUnit {server_info.id} missing port '{socket_name}' in ports={server_info.ports}"
+                )
+            address = f"tcp://{server_info.ip}:{port}"
             identity = f"{self.storage_manager_id}_to_{server_info.id}_{uuid4().hex[:8]}".encode()
             sock = create_zmq_socket(context, zmq.DEALER, identity=identity)

---

211-221: Handle empty field_data to avoid NonTensorStack construction errors.

If transfer_data["field_data"][field] is an empty list, NonTensorStack(*[]) is called with no arguments, which may raise an error depending on the NonTensorStack implementation.

Add defensive handling for empty field data:

         tensordict_data = TensorDict(
             {
                 field: (
-                    torch.nested.as_nested_tensor(transfer_data["field_data"][field])
-                    if transfer_data["field_data"][field]
-                    and all(isinstance(x, torch.Tensor) for x in transfer_data["field_data"][field])
-                    else NonTensorStack(*transfer_data["field_data"][field])
+                    self._to_stackable(transfer_data["field_data"][field])
                 )
                 for field in transfer_data["field_data"]
             }
         )
+    
+    def _to_stackable(self, field_data_list):
+        """Convert field data to stackable format, handling empty lists."""
+        if not field_data_list:
+            return NonTensorStack()  # Handle empty case
+        if all(isinstance(x, torch.Tensor) for x in field_data_list):
+            return torch.nested.as_nested_tensor(field_data_list)
+        return NonTensorStack(*field_data_list)

Add this helper method to the AsyncSimpleStorageManager class to improve readability and handle edge cases.

transfer_queue/storage/simple_backend.py (1)

260-262: Clarify field_data parameter type in docstring.

The comment at line 260 states field_data should be a TensorDict, which is correct. However, past comments noted the format description was inconsistent with usage.

The implementation is correct. For clarity, you could enhance the docstring:

         """
         Handle put request, add or update data into storage unit.
 
         Args:
-            data_parts: ZMQMessage from client.
+            data_parts: ZMQMessage from client containing:
+                        - local_indexes: list[int] - target storage positions
+                        - data: TensorDict - field data to store
 
         Returns:
             Put data success response ZMQMessage.
         """

transfer_queue/storage/managers/simple_backend_manager.py (1)

287-303: Review nested comprehension complexity.

The nested conditional in lines 289-299 for building tensor_data is complex. While functionally correct, it could benefit from extraction into a helper method for readability.

Consider extracting to a helper:

def _stack_field_data(self, field: str, v: list) -> torch.Tensor | NonTensorStack:
    """Stack field data with appropriate nesting."""
    if not v:
        return NonTensorStack()
    
    if not all(isinstance(item, torch.Tensor) for item in v):
        return NonTensorStack(*v)
    
    # All tensors - check if same shape
    if all(item.shape == v[0].shape for item in v):
        return torch.stack(torch.nested.as_nested_tensor(v).unbind())
    else:
        return torch.nested.as_nested_tensor(v)

Then simplify the dict comprehension:

tensor_data = {
    field: self._stack_field_data(field, v)
    for field, v in ordered_data.items()
}

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

transfer_queue/metadata.py (3)

89-92: Docstring still contradicts implementation.

The docstring claims default values will be used for missing dtype/shape/production_status, but the implementation (line 93) simply merges the provided fields dictionary without any defaulting logic.

---

298-298: Consider list comprehension for better readability.

As previously noted, a list comprehension would be more direct and readable than creating an intermediate iterator.

---

300-300: Breaking change: ChainMap reverses precedence order.

The current implementation uses dict(ChainMap(*(chunk.extra_info for chunk in data))), which gives precedence to the first chunk's extra_info values when keys overlap.

If the original code used update() in a loop, later chunks would overwrite earlier ones (last-wins semantics). This is a breaking behavioral change.

To restore last-wins semantics, apply this diff:

-    merged_extra_info = dict(ChainMap(*(chunk.extra_info for chunk in data)))
+    merged_extra_info = {}
+    for chunk in data:
+        merged_extra_info.update(chunk.extra_info)

Or if first-wins is intentional, document this explicitly in the method docstring.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro