From 617b333afc8d22bc63204fc7044698566f150105 Mon Sep 17 00:00:00 2001 From: ji-huazhong Date: Sun, 29 Mar 2026 18:14:18 +0800 Subject: [PATCH 1/2] Remove sys.path hacks; rely on installed package layout Drop manual sys.path.append to the repo root in tests, recipe demo, and scripts. Imports from transfer_queue now assume the package is installed (e.g. pip install -e .) or the interpreter path is set correctly. Clean up noqa: E402 on affected imports where applicable. Signed-off-by: ji-huazhong --- recipe/simple_use_case/single_controller_demo.py | 5 ----- scripts/performance_test/perftest.py | 7 +------ .../performance_test/ray_perftest_baseline.py | 5 ----- scripts/put_benchmark.py | 15 +++++---------- tests/e2e/test_e2e_lifecycle_consistency.py | 5 ----- tests/e2e/test_kv_interface_e2e.py | 8 +------- tests/test_async_simple_storage_manager.py | 14 ++++---------- tests/test_client.py | 16 ++++------------ tests/test_controller.py | 7 +------ tests/test_controller_data_partitions.py | 6 ------ tests/test_kv_storage_manager.py | 10 ++-------- tests/test_metadata.py | 9 +-------- tests/test_ray_p2p.py | 15 +++++---------- tests/test_samplers.py | 14 ++++---------- tests/test_serial_utils_on_cpu.py | 8 +------- tests/test_simple_storage_unit.py | 10 ++-------- tests/test_yuanrong_client_zero_copy.py | 11 ++--------- transfer_queue/storage/clients/factory.py | 1 - 18 files changed, 33 insertions(+), 133 deletions(-) diff --git a/recipe/simple_use_case/single_controller_demo.py b/recipe/simple_use_case/single_controller_demo.py index 870e6f14..5592ce2c 100644 --- a/recipe/simple_use_case/single_controller_demo.py +++ b/recipe/simple_use_case/single_controller_demo.py @@ -18,12 +18,10 @@ import logging import os import random -import sys import time import uuid from dataclasses import dataclass, field from importlib import resources -from pathlib import Path import ray import torch @@ -35,9 +33,6 @@ import transfer_queue as tq from transfer_queue import KVBatchMeta -parent_dir = Path(__file__).resolve().parent.parent.parent -sys.path.append(str(parent_dir)) - logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) diff --git a/scripts/performance_test/perftest.py b/scripts/performance_test/perftest.py index 95eb45ba..7abe1ddb 100644 --- a/scripts/performance_test/perftest.py +++ b/scripts/performance_test/perftest.py @@ -18,9 +18,7 @@ import csv import logging import os -import sys import time -from pathlib import Path from typing import Any import ray @@ -28,10 +26,7 @@ from omegaconf import OmegaConf from tensordict import NonTensorStack, TensorDict -parent_dir = Path(__file__).resolve().parent.parent.parent -sys.path.append(str(parent_dir)) - -import transfer_queue as tq # noqa: E402 +import transfer_queue as tq logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) diff --git a/scripts/performance_test/ray_perftest_baseline.py b/scripts/performance_test/ray_perftest_baseline.py index fe40788b..1542b628 100644 --- a/scripts/performance_test/ray_perftest_baseline.py +++ b/scripts/performance_test/ray_perftest_baseline.py @@ -18,18 +18,13 @@ import csv import logging import os -import sys import time -from pathlib import Path from typing import Any import ray import torch from tensordict import NonTensorStack, TensorDict -parent_dir = Path(__file__).resolve().parent.parent.parent -sys.path.append(str(parent_dir)) - 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 index 6b2afb59..005572c2 100644 --- a/scripts/put_benchmark.py +++ b/scripts/put_benchmark.py @@ -19,9 +19,7 @@ import logging import math import os -import sys import time -from pathlib import Path import numpy as np import ray @@ -30,14 +28,11 @@ from tensordict import TensorDict from tensordict.utils import LinkedList -parent_dir = Path(__file__).resolve().parent.parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue import TransferQueueClient # noqa: E402 -from transfer_queue.controller import TransferQueueController # noqa: E402 -from transfer_queue.storage.simple_backend import SimpleStorageUnit # noqa: E402 -from transfer_queue.utils.common import get_placement_group # noqa: E402 -from transfer_queue.utils.zmq_utils import process_zmq_server_info # noqa: E402 +from transfer_queue import TransferQueueClient +from transfer_queue.controller import TransferQueueController +from transfer_queue.storage.simple_backend import SimpleStorageUnit +from transfer_queue.utils.common import get_placement_group +from transfer_queue.utils.zmq_utils import process_zmq_server_info logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) diff --git a/tests/e2e/test_e2e_lifecycle_consistency.py b/tests/e2e/test_e2e_lifecycle_consistency.py index c2717860..218ecba8 100644 --- a/tests/e2e/test_e2e_lifecycle_consistency.py +++ b/tests/e2e/test_e2e_lifecycle_consistency.py @@ -16,7 +16,6 @@ import os import sys import time -from pathlib import Path import numpy as np import pytest @@ -26,10 +25,6 @@ from tensordict import TensorDict from tensordict.tensorclass import NonTensorData -# Setup paths (transfer_queue is not pip-installed) -parent_dir = Path(__file__).resolve().parent.parent.parent -sys.path.append(str(parent_dir)) - # Module-level default fields to avoid repeated generation DEFAULT_FIELDS = [ "tensor_f32", diff --git a/tests/e2e/test_kv_interface_e2e.py b/tests/e2e/test_kv_interface_e2e.py index 6ace0992..fb12edbb 100644 --- a/tests/e2e/test_kv_interface_e2e.py +++ b/tests/e2e/test_kv_interface_e2e.py @@ -22,8 +22,6 @@ import asyncio import os -import sys -from pathlib import Path import pytest import ray @@ -31,11 +29,7 @@ from omegaconf import OmegaConf from tensordict import TensorDict -# Add parent directory to path -parent_dir = Path(__file__).resolve().parent.parent.parent -sys.path.append(str(parent_dir)) - -import transfer_queue as tq # noqa: E402 +import transfer_queue as tq class TQAPIWrapper: diff --git a/tests/test_async_simple_storage_manager.py b/tests/test_async_simple_storage_manager.py index 2b16d500..bbf6d4b3 100644 --- a/tests/test_async_simple_storage_manager.py +++ b/tests/test_async_simple_storage_manager.py @@ -13,8 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys -from pathlib import Path from unittest.mock import AsyncMock, Mock, patch import numpy as np @@ -24,14 +22,10 @@ import zmq from tensordict import NonTensorStack, TensorDict -# Setup path -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.metadata import BatchMeta # noqa: E402 -from transfer_queue.storage import AsyncSimpleStorageManager # noqa: E402 -from transfer_queue.utils.enum_utils import TransferQueueRole # noqa: E402 -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo # noqa: E402 +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage import AsyncSimpleStorageManager +from transfer_queue.utils.enum_utils import TransferQueueRole +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType, ZMQServerInfo @pytest_asyncio.fixture diff --git a/tests/test_client.py b/tests/test_client.py index b9d10397..a911511d 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -13,9 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys import time -from pathlib import Path from threading import Thread from unittest.mock import patch @@ -24,16 +22,10 @@ import zmq from tensordict import NonTensorStack, TensorDict -# Import your classes here -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue import TransferQueueClient # noqa: E402 -from transfer_queue.metadata import ( # noqa: E402 - BatchMeta, -) -from transfer_queue.utils.enum_utils import TransferQueueRole # noqa: E402 -from transfer_queue.utils.zmq_utils import ( # noqa: E402 +from transfer_queue import TransferQueueClient +from transfer_queue.metadata import BatchMeta +from transfer_queue.utils.enum_utils import TransferQueueRole +from transfer_queue.utils.zmq_utils import ( ZMQMessage, ZMQRequestType, ZMQServerInfo, diff --git a/tests/test_controller.py b/tests/test_controller.py index 2e559600..707497cd 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -14,22 +14,17 @@ # limitations under the License. import logging -import sys -from pathlib import Path import pytest import ray import torch -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) +from transfer_queue.controller import TestTransferQueueController # Set up logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) -from transfer_queue.controller import TransferQueueController # noqa: E402 - @pytest.fixture(scope="function") def ray_setup(): diff --git a/tests/test_controller_data_partitions.py b/tests/test_controller_data_partitions.py index 02bacbb5..8f14517a 100644 --- a/tests/test_controller_data_partitions.py +++ b/tests/test_controller_data_partitions.py @@ -15,13 +15,7 @@ import logging import os -import sys import time -from pathlib import Path - -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - # Set up logging logging.basicConfig(level=logging.INFO) diff --git a/tests/test_kv_storage_manager.py b/tests/test_kv_storage_manager.py index 2d1dfcc0..bfe9a3bf 100644 --- a/tests/test_kv_storage_manager.py +++ b/tests/test_kv_storage_manager.py @@ -13,20 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. import asyncio -import sys -from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch import pytest import torch from tensordict import TensorDict -# Setup path -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.metadata import BatchMeta # noqa: E402 -from transfer_queue.storage.managers.base import KVStorageManager # noqa: E402 +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers.base import KVStorageManager def get_meta(data, global_indexes=None): diff --git a/tests/test_metadata.py b/tests/test_metadata.py index 8691bad6..20ace48e 100644 --- a/tests/test_metadata.py +++ b/tests/test_metadata.py @@ -15,18 +15,11 @@ """Unit tests for TransferQueue metadata module - Columnar BatchMeta + KVBatchMeta.""" -import sys -from pathlib import Path - import numpy as np import pytest import torch -# Setup path -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.metadata import BatchMeta, KVBatchMeta # noqa: E402 +from transfer_queue.metadata import BatchMeta, KVBatchMeta # ============================================================================== # Columnar BatchMeta Tests diff --git a/tests/test_ray_p2p.py b/tests/test_ray_p2p.py index e958b84c..b9f0b835 100644 --- a/tests/test_ray_p2p.py +++ b/tests/test_ray_p2p.py @@ -13,9 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys import time -from pathlib import Path import numpy as np import ray @@ -23,14 +21,11 @@ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy from tensordict import TensorDict -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.client import TransferQueueClient # noqa: E402 -from transfer_queue.metadata import BatchMeta # noqa: E402 -from transfer_queue.storage.managers.base import KVStorageManager # noqa: E402 -from transfer_queue.storage.managers.factory import TransferQueueStorageManagerFactory # noqa: E402 -from transfer_queue.utils.zmq_utils import ZMQServerInfo # noqa: E402 +from transfer_queue.client import TransferQueueClient +from transfer_queue.metadata import BatchMeta +from transfer_queue.storage.managers.base import KVStorageManager +from transfer_queue.storage.managers.factory import TransferQueueStorageManagerFactory +from transfer_queue.utils.zmq_utils import ZMQServerInfo TEST_CONFIGS: list[tuple[tuple[int, int], torch.dtype]] = [ ((5000, 5000), torch.float32), diff --git a/tests/test_samplers.py b/tests/test_samplers.py index da7aae1a..f72f345a 100644 --- a/tests/test_samplers.py +++ b/tests/test_samplers.py @@ -15,20 +15,14 @@ """Unit tests for TransferQueue samplers.""" -import sys -from pathlib import Path from typing import Any import pytest -# Setup path -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.sampler import BaseSampler # noqa: E402 -from transfer_queue.sampler.grpo_group_n_sampler import GRPOGroupNSampler # noqa: E402 -from transfer_queue.sampler.rank_aware_sampler import RankAwareSampler # noqa: E402 -from transfer_queue.sampler.sequential_sampler import SequentialSampler # noqa: E402 +from transfer_queue.sampler import BaseSampler +from transfer_queue.sampler.grpo_group_n_sampler import GRPOGroupNSampler +from transfer_queue.sampler.rank_aware_sampler import RankAwareSampler +from transfer_queue.sampler.sequential_sampler import SequentialSampler class TestBaseSampler: diff --git a/tests/test_serial_utils_on_cpu.py b/tests/test_serial_utils_on_cpu.py index 9a0ec45a..ffa08d75 100644 --- a/tests/test_serial_utils_on_cpu.py +++ b/tests/test_serial_utils_on_cpu.py @@ -13,19 +13,13 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys -from pathlib import Path import numpy as np import pytest import torch from tensordict import TensorDict -# Import your classes here -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.utils.serial_utils import MsgpackDecoder, MsgpackEncoder # noqa: E402 +from transfer_queue.utils.serial_utils import MsgpackDecoder, MsgpackEncoder @pytest.mark.parametrize( diff --git a/tests/test_simple_storage_unit.py b/tests/test_simple_storage_unit.py index 01ef2027..c553a1a7 100644 --- a/tests/test_simple_storage_unit.py +++ b/tests/test_simple_storage_unit.py @@ -13,9 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys import time -from pathlib import Path import pytest import ray @@ -23,12 +21,8 @@ import torch import zmq -# Setup path -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) - -from transfer_queue.storage.simple_backend import SimpleStorageUnit # noqa: E402 -from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType # noqa: E402 +from transfer_queue.storage.simple_backend import SimpleStorageUnit +from transfer_queue.utils.zmq_utils import ZMQMessage, ZMQRequestType class MockStorageClient: diff --git a/tests/test_yuanrong_client_zero_copy.py b/tests/test_yuanrong_client_zero_copy.py index 423b1c7b..f3da2716 100644 --- a/tests/test_yuanrong_client_zero_copy.py +++ b/tests/test_yuanrong_client_zero_copy.py @@ -13,22 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -import sys -from pathlib import Path from unittest.mock import MagicMock import numpy as np import pytest import torch -pytest.importorskip("yr") - -parent_dir = Path(__file__).resolve().parent.parent -sys.path.append(str(parent_dir)) +from transfer_queue.storage.clients.yuanrong_client import GeneralKVClientAdapter -from transfer_queue.storage.clients.yuanrong_client import ( # noqa: E402 - GeneralKVClientAdapter, -) +pytest.importorskip("yr") class MockBuffer: diff --git a/transfer_queue/storage/clients/factory.py b/transfer_queue/storage/clients/factory.py index 1d9bda55..73979bd5 100644 --- a/transfer_queue/storage/clients/factory.py +++ b/transfer_queue/storage/clients/factory.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. - from transfer_queue.storage.clients.base import TransferQueueStorageKVClient From b1d5d4e866a3d51a9d42c8af566fe00034c7cd73 Mon Sep 17 00:00:00 2001 From: ji-huazhong Date: Sun, 29 Mar 2026 18:30:22 +0800 Subject: [PATCH 2/2] Fix test_controller import for TransferQueueController Signed-off-by: ji-huazhong --- tests/test_controller.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_controller.py b/tests/test_controller.py index 707497cd..d45c54eb 100644 --- a/tests/test_controller.py +++ b/tests/test_controller.py @@ -19,7 +19,7 @@ import ray import torch -from transfer_queue.controller import TestTransferQueueController +from transfer_queue.controller import TransferQueueController # Set up logging logging.basicConfig(level=logging.INFO)