From fdd499102ccb8431f81a0c4d336d09822e0479bb Mon Sep 17 00:00:00 2001 From: morluto <76467478+morluto@users.noreply.github.com> Date: Thu, 10 Sep 2026 16:51:25 +0800 Subject: [PATCH] fix(cache): unlink symlinks during temporary build cleanup --- include/deep_jit/utils/filesystem.hpp | 8 +- tests/test_cuda_proj/main.cpp | 17 ++ .../scripts/symlink_failing_post_hook.py | 12 ++ tests/test_filesystem.py | 148 ++++++++++++++++++ 4 files changed, 182 insertions(+), 3 deletions(-) create mode 100644 tests/test_cuda_proj/scripts/symlink_failing_post_hook.py create mode 100644 tests/test_filesystem.py diff --git a/include/deep_jit/utils/filesystem.hpp b/include/deep_jit/utils/filesystem.hpp index d48e5b3..719f9a9 100644 --- a/include/deep_jit/utils/filesystem.hpp +++ b/include/deep_jit/utils/filesystem.hpp @@ -87,11 +87,13 @@ inline void write_file_sync(const std::filesystem::path& path, const std::string // same parent directory. inline void safe_remove_all(const std::filesystem::path& path) { std::error_code error_code; - if (not std::filesystem::exists(path, error_code) or error_code) + const auto status = std::filesystem::symlink_status(path, error_code); + if (error_code or not std::filesystem::exists(status)) return; - // A single file - if (not std::filesystem::is_directory(path, error_code) or error_code) { + // Inspect the entry itself: unlink symlinks, including dangling ones, + // without traversing a directory outside this tree. + if (not std::filesystem::is_directory(status)) { std::filesystem::remove(path, error_code); return; } diff --git a/tests/test_cuda_proj/main.cpp b/tests/test_cuda_proj/main.cpp index 4baf23e..01711e5 100644 --- a/tests/test_cuda_proj/main.cpp +++ b/tests/test_cuda_proj/main.cpp @@ -2988,6 +2988,22 @@ void test_compiler_failure_cleanup(Runtime& runtime, const fs::path& cache_root) check_tmp_is_empty(cache_root); } +void test_symlink_failure_cleanup(Runtime& runtime, const fs::path& cache_root) { + const auto outside = cache_root / "symlink_cleanup_outside"; + deep_jit::make_dirs(outside); + const auto sentinel = outside / "sentinel"; + deep_jit::write_file_sync(sentinel, "preserve outside contents"); + set_env("DEEP_JIT_TEST_CLEANUP_OUTSIDE", outside.string()); + const CompilerOptions options {.post_hook = "scripts/symlink_failing_post_hook.py"}; + expect_failure( + [&] { runtime.compile_without_load("symlink_failure", get_template_source(96), options); }, + "command failed with exit code 7"); + unset_env("DEEP_JIT_TEST_CLEANUP_OUTSIDE"); + DJ_HOST_ASSERT(fs::is_regular_file(sentinel), "failed hook cleanup deleted an outside file"); + DJ_HOST_ASSERT(deep_jit::read(sentinel) == "preserve outside contents"); + check_tmp_is_empty(cache_root); +} + void test_backend_output_validation(Runtime& runtime, const fs::path& cache_root) { const auto include_dir = get_test_cuda_project_dir() / "include_original"; const auto source = get_template_source(94); @@ -3314,6 +3330,7 @@ void run_tests(pybind11::module_ module) { run_test("kernel count", [&] { test_kernel_count(*runtime); }); run_test("PTXAS checks", [&] { test_ptxas_checks(cache_root); }); run_test("compiler failure cleanup", [&] { test_compiler_failure_cleanup(*runtime, cache_root); }); + run_test("symlink failure cleanup", [&] { test_symlink_failure_cleanup(*runtime, cache_root); }); run_test("backend output validation", [&] { test_backend_output_validation(*runtime, cache_root); }); run_test("dump artifacts and launch overhead", [&] { test_dump_and_launch_overhead(*runtime); }); run_test("dump options on cache hit", [&] { test_dump_options_on_cache_hit(cache_root); }); diff --git a/tests/test_cuda_proj/scripts/symlink_failing_post_hook.py b/tests/test_cuda_proj/scripts/symlink_failing_post_hook.py new file mode 100644 index 0000000..891f93a --- /dev/null +++ b/tests/test_cuda_proj/scripts/symlink_failing_post_hook.py @@ -0,0 +1,12 @@ +import os +import sys +from pathlib import Path + + +cubin_path = Path(sys.argv[1]) +assert cubin_path.is_file() +outside = Path(os.environ['DEEP_JIT_TEST_CLEANUP_OUTSIDE']) +(cubin_path.parent / 'directory_link').symlink_to(outside, target_is_directory=True) +(cubin_path.parent / 'file_link').symlink_to(outside / 'sentinel') +(cubin_path.parent / 'dangling_link').symlink_to(outside / 'missing') +raise SystemExit(7) diff --git a/tests/test_filesystem.py b/tests/test_filesystem.py new file mode 100644 index 0000000..d9cb04e --- /dev/null +++ b/tests/test_filesystem.py @@ -0,0 +1,148 @@ +"""CPU-only cleanup regressions; run with python tests/test_filesystem.py.""" + +import tempfile +from pathlib import Path + +from torch.utils.cpp_extension import load_inline + + +ROOT = Path(__file__).resolve().parents[1] + +CPP_SOURCE = r''' +#include +#include +#include + +namespace fs = std::filesystem; + +// Exercise the real runtime's exception path without a GPU or external compiler. +struct FailingBackend { + struct Device {}; + struct Kernel {}; + struct CompilerOptions { + static CompilerOptions default_options(const deep_jit::Env&, Device&) { return {}; } + }; + struct LaunchOptions { + static LaunchOptions default_options(const deep_jit::Env&) { return {}; } + }; + struct CompilerInfo { + std::string get_hash() const { return "cleanup-test"; } + } compiler_info; + + explicit FailingBackend(const deep_jit::Env&) {} + + void compile(const std::string&, const fs::path& path, const deep_jit::Env&, + const deep_jit::Config& config, const CompilerOptions&) const { + deep_jit::write_file_sync(path / "partial", "incomplete build"); + fs::create_directory_symlink(config.python_library_root / "outside", path / "link"); + fs::create_symlink(config.python_library_root / "missing", path / "dangling"); + throw std::runtime_error("synthetic compiler failure"); + } +}; + +void remove_tree(const std::string& path) { + deep_jit::safe_remove_all(path); +} + +void fail_compilation(const std::string& root) { + deep_jit::Runtime runtime(deep_jit::Config(root, "CLEANUP_TEST")); + runtime.compile("cleanup", "source", "testkey", {}); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def("remove_tree", &remove_tree); + module.def("fail_compilation", &fail_compilation); +} +''' + + +def check_cleanup(module, root: Path, kind: str, nested: bool) -> None: + outside = root / 'outside' + outside.mkdir() + sentinel = outside / 'sentinel' + sentinel.write_text('keep me') + target = root / 'cleanup' + link = target + if nested: + target.mkdir() + (target / 'partial').write_text('partial build') + link = target / 'link' + if kind == 'directory': + link.symlink_to(outside, target_is_directory=True) + elif kind == 'file': + link.symlink_to(sentinel) + else: + link.symlink_to(root / 'missing') + + module.remove_tree(str(target)) + assert sentinel.is_file(), f'{kind}, nested={nested}: outside sentinel deleted' + assert sentinel.read_text() == 'keep me' + assert not target.exists() and not target.is_symlink(), f'{kind}, nested={nested}: cleanup left entries' + module.remove_tree(str(target)) # Repeated cleanup of an absent path is harmless. + + +def main() -> None: + import os + + # Keep build products and disposable test trees in the checkout, not /tmp. + with tempfile.TemporaryDirectory(prefix='.filesystem-test-', dir=ROOT) as directory: + root = Path(directory) + build = root / 'build' + build.mkdir() + module = load_inline( + name='deep_jit_filesystem_test', + cpp_sources=CPP_SOURCE, + extra_cflags=['-std=c++20', '-O0'], + extra_include_paths=[str(ROOT / 'include')], + build_directory=str(build), + with_cuda=False, + verbose=True, + ) + failures = [] + for kind in ('directory', 'file', 'dangling'): + for nested in (False, True): + case = root / f'{kind}-{nested}' + case.mkdir() + try: + check_cleanup(module, case, kind, nested) + print(f'PASS: {case.name}') + except AssertionError as error: + failures.append(str(error)) + + ordinary = root / 'ordinary' + (ordinary / 'nested').mkdir(parents=True) + (ordinary / 'nested' / 'file').write_text('ordinary') + module.remove_tree(str(ordinary)) + assert not ordinary.exists() + print('PASS: ordinary tree') + + runtime_root = root / 'runtime' + (runtime_root / 'outside').mkdir(parents=True) + sentinel = runtime_root / 'outside' / 'sentinel' + sentinel.write_text('keep me') + env_name = 'CLEANUP_TEST_JIT_CACHE_DIR' + previous = os.environ.get(env_name) + os.environ[env_name] = str(runtime_root) + try: + try: + module.fail_compilation(str(runtime_root)) + except RuntimeError as error: + assert str(error) == 'synthetic compiler failure' + else: + raise AssertionError('compilation did not fail') + finally: + if previous is None: + del os.environ[env_name] + else: + os.environ[env_name] = previous + if not sentinel.is_file() or sentinel.read_text() != 'keep me': + failures.append('compilation failure deleted outside sentinel') + if list((runtime_root / 'tmp').iterdir()): + failures.append('compilation failure left a temporary build') + if not failures: + print('PASS: compilation failure cleanup') + assert not failures, '\n'.join(failures) + + +if __name__ == '__main__': + main()