Files
2026-07-13 13:22:34 +08:00

67 lines
1.9 KiB
Python

import os
import uuid
import pytest
from mlflow.utils.os import is_windows
from mlflow.utils.process import cache_return_value_per_process
@cache_return_value_per_process
def _gen_random_str1(v):
return str(v) + uuid.uuid4().hex
@cache_return_value_per_process
def _gen_random_str2(v):
return str(v) + uuid.uuid4().hex
@cache_return_value_per_process
def _gen_random_no_arg():
return uuid.uuid4().hex
def test_cache_return_value_per_process():
path1 = _gen_random_str1(True)
path2 = _gen_random_str1(True)
assert path1 == path2
path3 = _gen_random_str1(False)
assert path3 != path2
no_arg_path1 = _gen_random_no_arg()
no_arg_path2 = _gen_random_no_arg()
assert no_arg_path1 == no_arg_path2
with pytest.raises(
ValueError,
match="The function decorated by `cache_return_value_per_process` is not allowed to be "
"called with key-word style arguments.",
):
_gen_random_str1(v=True)
f2_path1 = _gen_random_str2(True)
f2_path2 = _gen_random_str2(False)
assert len({path1, path3, f2_path1, f2_path2}) == 4
# Skip the following block on Windows which doesn't support `os.fork`
if not is_windows():
# Use `os.fork()` to make parent and child processes share the same global variable
# `_per_process_value_cache_map`.
pid = os.fork()
if pid > 0:
# in parent process
child_pid = pid
# check child process exit with return value 0.
assert os.waitpid(child_pid, 0)[1] == 0
else:
# in forked out child process
child_path1 = _gen_random_str1(True)
child_path2 = _gen_random_str1(False)
test_pass = len({path1, path3, child_path1, child_path2}) == 4
# exit forked out child process with exit code representing testing pass or fail.
os._exit(0 if test_pass else 1)