GeezerVison

recipe/simple_use_case/sync_demo.py (4)

94-94: Comment translation implemented as previously suggested.

This change implements the translation suggested in a previous review.

---

109-109: Comment translation implemented as previously suggested.

This change implements the translation suggested in a previous review.

---

118-118: Comment translation implemented as previously suggested.

This change implements the translation suggested in a previous review.

---

126-126: Comment translation implemented as previously suggested.

This change implements the translation suggested in a previous review.

transfer_queue/controller.py (1)

339-356: Critical: itemgetter fails with empty list and None checks are ineffective.

Two critical issues persist from previous review:

1. Line 339-340: itemgetter(*global_indices) raises TypeError when global_indices is empty
2. Line 353-355: None checks occur after tuple conversion, making them ineffective. After lines 342-345, dtype_value and shape_value are always tuples, so the None check passes even when containing None elements.

Apply this fix:

     def _update_field_metadata(
         self,
         global_indices: list[int],
         field_names: list[str],
         dtypes: Optional[dict[int, dict[str, Any]]] = None,
         shapes: Optional[dict[int, dict[str, Any]]] = None,
     ):
         """Update field dtype and shape metadata."""
-        dtype_value = itemgetter(*global_indices)(dtypes) if dtypes else None
-        shape_value = itemgetter(*global_indices)(shapes) if shapes else None
+        # Handle empty global_indices
+        if not global_indices:
+            return
+        
+        dtype_value = itemgetter(*global_indices)(dtypes) if dtypes and global_indices else None
+        shape_value = itemgetter(*global_indices)(shapes) if shapes and global_indices else None
 
         if not isinstance(dtype_value, tuple):
             dtype_value = (dtype_value,)
         if not isinstance(shape_value, tuple):
             shape_value = (shape_value,)
 
         for i, global_idx in enumerate(global_indices):
             if global_idx not in self.field_dtypes:
                 self.field_dtypes[global_idx] = {}
             if global_idx not in self.field_shapes:
                 self.field_shapes[global_idx] = {}
 
-            if dtype_value is not None:
+            if dtype_value is not None and dtype_value[i] is not None:
                 self.field_dtypes[global_idx].update(dtype_value[i])
-            if shape_value is not None:
+            if shape_value is not None and shape_value[i] is not None:
                 self.field_shapes[global_idx].update(shape_value[i])

tests/test_async_put.py (3)

46-59: Consider making client_id unique per worker instance.

All AsyncPutWorker instances share the hardcoded client_id "AsyncvLLMServer", which may cause confusion in logs when multiple workers operate concurrently.

Consider passing a unique identifier to each worker:

 @ray.remote(num_cpus=1)
 class AsyncPutWorker:
-    def __init__(self, config, data_system_controller_info):
+    def __init__(self, config, data_system_controller_info, worker_id=None):
         self.config = config
+        client_id = f"AsyncPutWorker_{worker_id}" if worker_id is not None else "AsyncPutWorker"
         self.data_system_client = AsyncTransferQueueClient(
-            client_id="AsyncvLLMServer",
+            client_id=client_id,
             controller_info=data_system_controller_info,
         )

Then update the manager to pass worker_id:

for i in range(num_workers):
    self.async_put_workers.append(AsyncPutWorker.remote(config, data_system_controller_info, worker_id=i))

---

157-178: LGTM! Test validates async_put workflow effectively.

The test correctly validates data integrity by comparing sorted flattened tensors, which accounts for potential ordering inconsistencies from distributed writes.

Consider adding a comment explaining why sorting is necessary:

     assert data_fetch.shape == data_origin.shape
     assert data_fetch["input_ids"].shape == data_origin["input_ids"].shape
+    # Compare sorted tensors to handle potential ordering differences from distributed async_put
     assert torch.equal(
         torch.sort(data_fetch["input_ids"].flatten())[0], torch.sort(data_origin["input_ids"].flatten())[0]
     )

---

41-43: Consider moving Ray initialization into test function for better isolation.

Module-level ray.init() may affect test isolation if other tests in the suite have different Ray configuration requirements.

Consider using a fixture or setUp method:

@pytest.fixture(scope="module")
def ray_cluster():
    os.environ["RAY_DEDUP_LOGS"] = "0"
    os.environ["RAY_DEBUG"] = "1"
    ray.init()
    yield
    ray.shutdown()

def test_async_put(ray_cluster):
    # test implementation
    ...

However, if this pattern is consistent with other tests in the project, the current approach is acceptable.

transfer_queue/controller.py (1)

183-183: Inconsistent initialization: production_status is never None but still typed Optional.

The production_status is now pre-initialized to a tensor, so it will never be None. However, the type hint remains Optional[torch.Tensor] and several methods still perform None checks.

Consider one of the following:

Option 1: Remove Optional and None checks

-    production_status: Optional[torch.Tensor] = torch.zeros(TQ_INIT_SAMPLE_NUM, TQ_INIT_FIELD_NUM, dtype=torch.int8)
+    production_status: torch.Tensor = field(default_factory=lambda: torch.zeros(TQ_INIT_SAMPLE_NUM, TQ_INIT_FIELD_NUM, dtype=torch.int8))

Then remove None checks in methods like ensure_fields_capacity (line 256), scan_data_status (line 414), etc.

Option 2: Keep Optional for flexibility

If None should be supported, initialize as None and handle the first-use initialization properly.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro