[Misc] fix unique_filepath (vllm-project#25732)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Signed-off-by: xuebwang-amd <xuebwang@amd.com>

Author: 2 people

tests/utils_/test_utils.py

@@ -23,15 +23,16 @@
 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

vllm/utils/__init__.py

@@ -45,6 +45,7 @@
 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