diff --git a/python/ray/_private/test_utils.py b/python/ray/_private/test_utils.py index 50812015c5b6c..3149fb9b7943c 100644 --- a/python/ray/_private/test_utils.py +++ b/python/ray/_private/test_utils.py @@ -19,6 +19,7 @@ import traceback from collections import defaultdict from contextlib import contextmanager, redirect_stderr, redirect_stdout + from typing import Any, Callable, Dict, List, Optional, Tuple import uuid from dataclasses import dataclass diff --git a/python/ray/tests/test_queue.py b/python/ray/tests/test_queue.py index e4db15c279493..6c9b97e76b8ae 100644 --- a/python/ray/tests/test_queue.py +++ b/python/ray/tests/test_queue.py @@ -3,7 +3,7 @@ import ray from ray.exceptions import GetTimeoutError, RayActorError -from ray.util.queue import Queue, Empty, Full +from ray.util.queue import Empty, Full, Queue from ray._private.test_utils import wait_for_condition, BatchQueue diff --git a/python/ray/util/queue.py b/python/ray/util/queue.py index b97bd80d13713..b714bfb7f7f5a 100644 --- a/python/ray/util/queue.py +++ b/python/ray/util/queue.py @@ -1,4 +1,5 @@ import asyncio +import queue from typing import Optional, Any, List, Dict from collections.abc import Iterable @@ -7,12 +8,12 @@ @PublicAPI(stability="beta") -class Empty(Exception): +class Empty(queue.Empty): pass @PublicAPI(stability="beta") -class Full(Exception): +class Full(queue.Full): pass