diff --git a/tests/utils_/test_utils.py b/tests/utils_/test_utils.py index 608f517f69145..658ae7e7451ae 100644 --- a/tests/utils_/test_utils.py +++ b/tests/utils_/test_utils.py @@ -23,15 +23,16 @@ from vllm_test_utils.monitor import monitor from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config from vllm.transformers_utils.detokenizer_utils import ( convert_ids_list_to_tokens) -from vllm.utils import (CacheInfo, FlexibleArgumentParser, LRUCache, - MemorySnapshot, PlaceholderModule, StoreBoolean, - bind_kv_cache, common_broadcastable_dtype, - current_stream, deprecate_kwargs, get_open_port, - get_tcp_uri, is_lossless_cast, join_host_port, - make_zmq_path, make_zmq_socket, memory_profiling, - merge_async_iterators, sha256, split_host_port, - split_zmq_path, supports_kw, swap_dict_values) +# isort: off +from vllm.utils import ( + CacheInfo, FlexibleArgumentParser, LRUCache, MemorySnapshot, + PlaceholderModule, bind_kv_cache, common_broadcastable_dtype, + current_stream, deprecate_kwargs, get_open_port, get_tcp_uri, + is_lossless_cast, join_host_port, make_zmq_path, make_zmq_socket, + memory_profiling, merge_async_iterators, sha256, split_host_port, + split_zmq_path, supports_kw, swap_dict_values, unique_filepath) +# isort: on from ..utils import create_new_process_for_each_test, error_on_warning @@ -1032,3 +1033,15 @@ def test_load_config_file(tmp_path): # Assert that the processed arguments match the expected output assert processed_args == expected_args os.remove(str(config_file_path)) + + +def test_unique_filepath(): + temp_dir = tempfile.mkdtemp() + path_fn = lambda i: Path(temp_dir) / f"file_{i}.txt" + paths = set() + for i in range(10): + path = unique_filepath(path_fn) + path.write_text("test") + paths.add(path) + assert len(paths) == 10 + assert len(list(Path(temp_dir).glob("*.txt"))) == 10 diff --git a/vllm/utils/__init__.py b/vllm/utils/__init__.py index c502a69ea500e..ba280d6dbe4a4 100644 --- a/vllm/utils/__init__.py +++ b/vllm/utils/__init__.py @@ -45,6 +45,7 @@ from concurrent.futures import ThreadPoolExecutor from concurrent.futures.process import ProcessPoolExecutor from dataclasses import dataclass, field from functools import cache, lru_cache, partial, wraps +from pathlib import Path from types import MappingProxyType from typing import (TYPE_CHECKING, Any, Callable, Generic, Literal, NamedTuple, Optional, TextIO, TypeVar, Union, cast, overload) @@ -3536,3 +3537,23 @@ def set_env_var(key, value): del os.environ[key] else: os.environ[key] = old + + +def unique_filepath(fn: Callable[[int], Path]) -> Path: + """ + unique_filepath returns a unique path by trying + to include an integer in increasing order. + + fn should be a callable that returns a path that + includes the passed int at a fixed location. + + Note: This function has a TOCTOU race condition. + Caller should use atomic operations (e.g., open with 'x' mode) + when creating the file to ensure thread safety. + """ + i = 0 + while True: + p = fn(i) + if not p.exists(): + return p + i += 1