[chore] Optimize utils structure and README

0oshowero0 · ascend-robot · commit c2e514a0bb26 · 2026-01-29T12:31:39.000+08:00
Co-authored-by: 0oshowero0&lt;o0shower0o@outlook.com&gt;



# message auto-generated for no-merge-commit merge:
!20 merge chore into main

[chore] Optimize utils structure and README

Created-by: hanzhenyu8
Commit-by: 0oshowero0
Merged-by: ascend-robot
Description: 1. Refactor utils/utils.py -&gt; utils/common.py and utils/enum_utils.py
2. Add back github workflow UT for export TQ_ZERO_COPY_SERIALIZATION=True
3. Update README

See merge request: Ascend/TransferQueue!20
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
@@ -43,6 +43,11 @@ jobs:
       run: |
         python -m build --wheel
         pip install dist/*.whl --force-reinstall
-    - name: Test with pytest
+    - name: Test with pytest (TQ_ZERO_COPY_SERIALIZATION=False)
       run: |
+        pytest
+    - name: Test with pytest (TQ_ZERO_COPY_SERIALIZATION=True)
+      run: |
+        ray stop --force
+        export TQ_ZERO_COPY_SERIALIZATION=True
         pytest
diff --git a/README.md b/README.md
@@ -25,7 +25,7 @@ TransferQueue is a high-performance data storage and transfer module with panora
 TransferQueue offers **fine-grained, sub-sample-level** data management and **load-balancing** (on the way) capabilities, serving as a data gateway that decouples explicit data dependencies across computational tasks. This enables a divide-and-conquer approach, significantly simplifies the algorithm controller design.
 
 <p align="center">
-  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/main_func.png?raw=true" width="70%">
+  <img src="https://github.com/TransferQueue/community_doc/blob/main/docs/main_func.png?raw=true" width="100%">
 </p>
 
 <h2 id="updates">🔄 Updates</h2>
diff --git a/recipe/simple_use_case/async_demo.py b/recipe/simple_use_case/async_demo.py
@@ -36,7 +36,7 @@
     TransferQueueController,
     process_zmq_server_info,
 )
-from transfer_queue.utils.utils import get_placement_group  # noqa: E402
+from transfer_queue.utils.common import get_placement_group  # noqa: E402
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
diff --git a/recipe/simple_use_case/sync_demo.py b/recipe/simple_use_case/sync_demo.py
@@ -34,7 +34,7 @@
     TransferQueueController,
     process_zmq_server_info,
 )
-from transfer_queue.utils.utils import get_placement_group  # noqa: E402
+from transfer_queue.utils.common import get_placement_group  # noqa: E402
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
diff --git a/scripts/performance_test.py b/scripts/performance_test.py
@@ -37,7 +37,7 @@
     TransferQueueController,
     process_zmq_server_info,
 )
-from transfer_queue.utils.utils import get_placement_group  # noqa: E402
+from transfer_queue.utils.common import get_placement_group  # noqa: E402
 
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
diff --git a/scripts/put_benchmark.py b/scripts/put_benchmark.py
@@ -39,7 +39,7 @@
     TransferQueueController,
     process_zmq_server_info,
 )
-from transfer_queue.utils.utils import get_placement_group  # noqa: E402
+from transfer_queue.utils.common import get_placement_group  # noqa: E402
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py
@@ -29,7 +29,7 @@
 
 from transfer_queue.metadata import BatchMeta, FieldMeta, SampleMeta  # noqa: E402
 from transfer_queue.storage import AsyncSimpleStorageManager  # noqa: E402
-from transfer_queue.utils.utils import TransferQueueRole  # noqa: E402
+from transfer_queue.utils.enum_utils import TransferQueueRole  # noqa: E402
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo  # noqa: E402
 
 
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -34,7 +34,7 @@
     FieldMeta,
     SampleMeta,
 )
-from transfer_queue.utils.utils import TransferQueueRole  # noqa: E402
+from transfer_queue.utils.enum_utils import TransferQueueRole  # noqa: E402
 from transfer_queue.utils.zmq_utils import (  # noqa: E402
     ZMQMessage,
     ZMQRequestType,
diff --git a/tests/test_controller.py b/tests/test_controller.py
@@ -29,7 +29,7 @@
 logger = logging.getLogger(__name__)
 
 from transfer_queue import TransferQueueController  # noqa: E402
-from transfer_queue.utils.utils import ProductionStatus  # noqa: E402
+from transfer_queue.utils.enum_utils import ProductionStatus  # noqa: E402
 
 
 @pytest.fixture(scope="function")
diff --git a/tests/test_metadata.py b/tests/test_metadata.py
@@ -28,7 +28,7 @@
 sys.path.append(str(parent_dir))
 
 from transfer_queue.metadata import BatchMeta, FieldMeta, SampleMeta  # noqa: E402
-from transfer_queue.utils.utils import ProductionStatus  # noqa: E402
+from transfer_queue.utils.enum_utils import ProductionStatus  # noqa: E402
 
 
 class TestFieldMeta:
diff --git a/transfer_queue/__init__.py b/transfer_queue/__init__.py
@@ -16,7 +16,6 @@
 import os
 
 from .client import (
-    AsyncTransferQueueClient,
     TransferQueueClient,
     process_zmq_server_info,
 )
@@ -28,11 +27,10 @@
 from .sampler.rank_aware_sampler import RankAwareSampler
 from .sampler.sequential_sampler import SequentialSampler
 from .storage import SimpleStorageUnit
-from .utils.utils import get_placement_group
+from .utils.common import get_placement_group
 from .utils.zmq_utils import ZMQServerInfo
 
 __all__ = [
-    "AsyncTransferQueueClient",
     "TransferQueueClient",
     "StreamingDataset",
     "StreamingDataLoader",
diff --git a/transfer_queue/client.py b/transfer_queue/client.py
@@ -37,7 +37,7 @@
     TransferQueueStorageManager,
     TransferQueueStorageManagerFactory,
 )
-from transfer_queue.utils.utils import limit_pytorch_auto_parallel_threads
+from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
     ZMQRequestType,
diff --git a/transfer_queue/controller.py b/transfer_queue/controller.py
@@ -37,11 +37,8 @@
     SampleMeta,
 )
 from transfer_queue.sampler import BaseSampler, SequentialSampler
+from transfer_queue.utils.enum_utils import ProductionStatus, TransferQueueRole
 from transfer_queue.utils.perf_utils import IntervalPerfMonitor
-from transfer_queue.utils.utils import (
-    ProductionStatus,
-    TransferQueueRole,
-)
 from transfer_queue.utils.zmq_utils import (
     ZMQMessage,
     ZMQRequestType,
diff --git a/transfer_queue/metadata.py b/transfer_queue/metadata.py
@@ -27,7 +27,7 @@
 from tensordict import TensorDict
 from tensordict.tensorclass import NonTensorData, NonTensorStack
 
-from transfer_queue.utils.utils import ProductionStatus
+from transfer_queue.utils.enum_utils import ProductionStatus
 
 logger = logging.getLogger(__name__)
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
diff --git a/transfer_queue/storage/managers/simple_backend_manager.py b/transfer_queue/storage/managers/simple_backend_manager.py
@@ -30,7 +30,7 @@
 from transfer_queue.storage.managers.base import TransferQueueStorageManager
 from transfer_queue.storage.managers.factory import TransferQueueStorageManagerFactory
 from transfer_queue.storage.simple_backend import StorageMetaGroup
-from transfer_queue.utils.utils import get_env_bool
+from transfer_queue.utils.common import get_env_bool
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket
 
 logger = logging.getLogger(__name__)
diff --git a/transfer_queue/storage/simple_backend.py b/transfer_queue/storage/simple_backend.py
@@ -27,8 +27,9 @@
 from ray.util import get_node_ip_address
 
 from transfer_queue.metadata import SampleMeta
+from transfer_queue.utils.common import limit_pytorch_auto_parallel_threads
+from transfer_queue.utils.enum_utils import TransferQueueRole
 from transfer_queue.utils.perf_utils import IntervalPerfMonitor
-from transfer_queue.utils.utils import TransferQueueRole, limit_pytorch_auto_parallel_threads
 from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo, create_zmq_socket, get_free_port
 
 logger = logging.getLogger(__name__)
diff --git a/transfer_queue/utils/common.py b/transfer_queue/utils/common.py
@@ -16,7 +16,6 @@
 import logging
 import os
 from contextlib import contextmanager
-from enum import Enum
 from typing import Optional
 
 import psutil
@@ -35,35 +34,6 @@
     logger.addHandler(handler)
 
 
-class ExplicitEnum(str, Enum):
-    """
-    Enum with more explicit error message for missing values.
-    """
-
-    @classmethod
-    def _missing_(cls, value):
-        raise ValueError(
-            f"{value} is not a valid {cls.__name__}, please select one of {list(cls._value2member_map_.keys())}"
-        )
-
-
-class TransferQueueRole(ExplicitEnum):
-    """Available Roles of TransferQueue."""
-
-    CONTROLLER = "TransferQueueController"
-    STORAGE = "TransferQueueStorage"
-    CLIENT = "TransferQueueClient"
-
-
-class ProductionStatus(ExplicitEnum):
-    """
-    Data Production Status.
-    """
-
-    NOT_PRODUCED = 0
-    READY_FOR_CONSUME = 1
-
-
 def get_placement_group(num_ray_actors: int, num_cpus_per_actor: int = 1):
     """
     Create a placement group with SPREAD strategy for Ray actors.
@@ -81,44 +51,6 @@ def get_placement_group(num_ray_actors: int, num_cpus_per_actor: int = 1):
     return placement_group
 
 
-def sequential_sampler(
-    ready_for_consume_idx: list[int],
-    batch_size: int,
-    get_n_samples: bool,
-    n_samples_per_prompt: int,
-) -> list[int]:
-    """
-    Sequentially samples a batch of indices from global indexes ready_for_consume_idx.
-
-    Args:
-        ready_for_consume_idx: A sorted list of available indices for sampling.
-            - When get_n_samples=True:
-                Expected to be grouped by prompts, e.g.,
-                [0,1,2,3, 8,9,10,11, 12,13,14,15] (3 groups of 4 samples each)
-            - When get_n_samples=False:
-                Can be any ordered list, e.g., [0,3,5,6,7,8]
-        batch_size: Total number of samples to return
-        get_n_samples: Flag indicating the sampling mode
-        n_samples_per_prompt: Number of samples per prompt (used when get_n_samples=True)
-
-    Returns:
-        list[int]: Sequentially sampled indices of length batch_size
-    """
-    if get_n_samples:
-        assert len(ready_for_consume_idx) % n_samples_per_prompt == 0
-        assert batch_size % n_samples_per_prompt == 0
-        batch_size_n_samples = batch_size // n_samples_per_prompt
-
-        group_ready_for_consume_idx = torch.tensor(ready_for_consume_idx, dtype=torch.int).view(
-            -1, n_samples_per_prompt
-        )
-
-        sampled_indexes = group_ready_for_consume_idx[list(range(batch_size_n_samples))].flatten().tolist()
-    else:
-        sampled_indexes = [int(ready_for_consume_idx[i]) for i in range(batch_size)]
-    return sampled_indexes
-
-
 @contextmanager
 def limit_pytorch_auto_parallel_threads(target_num_threads: Optional[int] = None, info: str = ""):
     """Prevent PyTorch from overdoing the automatic parallelism during torch.stack() operation"""
diff --git a/transfer_queue/utils/enum_utils.py b/transfer_queue/utils/enum_utils.py
@@ -0,0 +1,56 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import os
+from enum import Enum
+
+logger = logging.getLogger(__name__)
+logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
+
+# Ensure logger has a handler
+if not logger.hasHandlers():
+    handler = logging.StreamHandler()
+    handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(name)s - %(message)s"))
+    logger.addHandler(handler)
+
+
+class ExplicitEnum(str, Enum):
+    """
+    Enum with more explicit error message for missing values.
+    """
+
+    @classmethod
+    def _missing_(cls, value):
+        raise ValueError(
+            f"{value} is not a valid {cls.__name__}, please select one of {list(cls._value2member_map_.keys())}"
+        )
+
+
+class TransferQueueRole(ExplicitEnum):
+    """Available Roles of TransferQueue."""
+
+    CONTROLLER = "TransferQueueController"
+    STORAGE = "TransferQueueStorage"
+    CLIENT = "TransferQueueClient"
+
+
+class ProductionStatus(ExplicitEnum):
+    """
+    Data Production Status.
+    """
+
+    NOT_PRODUCED = 0
+    READY_FOR_CONSUME = 1
diff --git a/transfer_queue/utils/zmq_utils.py b/transfer_queue/utils/zmq_utils.py
@@ -25,12 +25,11 @@
 import psutil
 import zmq
 
-from transfer_queue.utils.serial_utils import _decoder, _encoder
-from transfer_queue.utils.utils import (
-    ExplicitEnum,
-    TransferQueueRole,
+from transfer_queue.utils.common import (
     get_env_bool,
 )
+from transfer_queue.utils.enum_utils import ExplicitEnum, TransferQueueRole
+from transfer_queue.utils.serial_utils import _decoder, _encoder
 
 logger = logging.getLogger(__name__)
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
diff --git a/tutorial/02_metadata_concepts.py b/tutorial/02_metadata_concepts.py
@@ -51,7 +51,7 @@
     process_zmq_server_info,
 )
 from transfer_queue.metadata import BatchMeta, FieldMeta, SampleMeta  # noqa: E402
-from transfer_queue.utils.utils import ProductionStatus  # noqa: E402
+from transfer_queue.utils.enum_utils import ProductionStatus  # noqa: E402
 
 # Configure Ray
 os.environ["RAY_DEDUP_LOGS"] = "0"

Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@`
`36`	`36`	`TransferQueueController,`
`37`	`37`	`process_zmq_server_info,`
`38`	`38`	`)`
`39`		`-from transfer_queue.utils.utils import get_placement_group # noqa: E402`
	`39`	`+from transfer_queue.utils.common import get_placement_group # noqa: E402`
`40`	`40`
`41`	`41`	`logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")`
`42`	`42`	`logger = logging.getLogger(__name__)`
Original file line number	Diff line number	Diff line change
`@@ -34,7 +34,7 @@`
`34`	`34`	`TransferQueueController,`
`35`	`35`	`process_zmq_server_info,`
`36`	`36`	`)`
`37`		`-from transfer_queue.utils.utils import get_placement_group # noqa: E402`
	`37`	`+from transfer_queue.utils.common import get_placement_group # noqa: E402`
`38`	`38`
`39`	`39`	`logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")`
`40`	`40`	`logger = logging.getLogger(__name__)`
Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@`
`37`	`37`	`TransferQueueController,`
`38`	`38`	`process_zmq_server_info,`
`39`	`39`	`)`
`40`		`-from transfer_queue.utils.utils import get_placement_group # noqa: E402`
	`40`	`+from transfer_queue.utils.common import get_placement_group # noqa: E402`
`41`	`41`
`42`	`42`	`logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")`
`43`	`43`	`logger = logging.getLogger(__name__)`
Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@`
`39`	`39`	`TransferQueueController,`
`40`	`40`	`process_zmq_server_info,`
`41`	`41`	`)`
`42`		`-from transfer_queue.utils.utils import get_placement_group # noqa: E402`
	`42`	`+from transfer_queue.utils.common import get_placement_group # noqa: E402`
`43`	`43`
`44`	`44`	`logging.basicConfig(level=logging.INFO)`
`45`	`45`	`logger = logging.getLogger(__name__)`