From 5d056dcb63184a58e972aeef820d9b7d377b750b Mon Sep 17 00:00:00 2001 From: apaz Date: Fri, 12 Jun 2026 23:56:19 -0500 Subject: [PATCH] kernelthing fixes --- .gitignore | 2 ++ csrc/binding.cpp | 8 +++++--- csrc/landlock.cpp | 13 +++++++++---- csrc/manager.cpp | 35 +++++++++++++++++++++++++++-------- csrc/manager.h | 13 ++++++++++--- csrc/supervisor.cpp | 7 ++++++- python/pygpubench/__init__.py | 7 ++++++- 7 files changed, 65 insertions(+), 20 deletions(-) create mode 100644 .gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c9a5a88 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +**/__pycache__/** +build/** diff --git a/csrc/binding.cpp b/csrc/binding.cpp index 48dd074..25c5be3 100644 --- a/csrc/binding.cpp +++ b/csrc/binding.cpp @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -18,10 +19,10 @@ namespace nb = nanobind; void do_bench(int result_fd, int input_fd, int supervisor_sock_fd, const std::string& kernel_qualname, const nb::object& test_generator, const nb::dict& test_kwargs, std::uintptr_t stream, bool discard, bool nvtx, bool landlock, bool mseal, - bool allow_root) { + bool allow_root, const std::vector& writable_paths) { std::vector signature_bytes(32); auto config = read_benchmark_parameters(input_fd, signature_bytes.data()); - auto mgr = make_benchmark_manager(result_fd, signature_bytes, config.Seed, discard, nvtx, landlock, mseal, allow_root, supervisor_sock_fd); + auto mgr = make_benchmark_manager(result_fd, signature_bytes, config.Seed, discard, nvtx, landlock, mseal, allow_root, supervisor_sock_fd, writable_paths); cleanse(signature_bytes.data(), 32); { @@ -63,7 +64,8 @@ NB_MODULE(_pygpubench, m) { nb::arg("nvtx") = false, nb::arg("landlock") = true, nb::arg("mseal") = true, - nb::arg("allow_root") = false + nb::arg("allow_root") = false, + nb::arg("writable_paths") = std::vector{"/tmp"} ); m.def("run_supervisor", [](int sock_fd) { diff --git a/csrc/landlock.cpp b/csrc/landlock.cpp index cec8ba6..82ca644 100644 --- a/csrc/landlock.cpp +++ b/csrc/landlock.cpp @@ -18,6 +18,8 @@ #include #include #include +#include +#include #include #include #include @@ -83,7 +85,7 @@ static void allow_path(LandlockFd& ruleset, const char *path, uint64_t access) { landlock_add_rule(ruleset, LANDLOCK_RULE_PATH_BENEATH, &attr, 0); } -void install_landlock() { +void install_landlock(const std::vector& writable_paths) { const std::uint64_t RO = LANDLOCK_ACCESS_FS_READ_FILE | LANDLOCK_ACCESS_FS_READ_DIR; @@ -111,9 +113,12 @@ void install_landlock() { // Read-only: entire filesystem allow_path(ruleset_fd, "/", RO); - // Read-write: /tmp and /dev only - allow_path(ruleset_fd, "/tmp", RW); - allow_path(ruleset_fd, "/dev", RW); // needed for /dev/null etc, used e.g., by triton + // Read-write: /dev is always needed for /dev/null etc, used e.g., by triton. + allow_path(ruleset_fd, "/dev", RW); + // Additional writable paths are caller-configurable (defaults to /tmp). + for (const std::string& path : writable_paths) { + allow_path(ruleset_fd, path.c_str(), RW); + } // required for landlock if (prctl(PR_SET_NO_NEW_PRIVS, 1, 0, 0, 0) < 0) { diff --git a/csrc/manager.cpp b/csrc/manager.cpp index 10ddb63..21a7d00 100644 --- a/csrc/manager.cpp +++ b/csrc/manager.cpp @@ -5,8 +5,10 @@ #include "manager.h" #include "utils.h" #include "check.h" +#include #include #include +#include #include #include #include @@ -27,7 +29,7 @@ static constexpr std::size_t ArenaSize = 2 * 1024 * 1024; static constexpr std::size_t BenchmarkManagerArenaSize = 128 * 1024 * 1024; extern void clear_cache(void* dummy_memory, int size, bool discard, cudaStream_t stream); -extern void install_landlock(); +extern void install_landlock(const std::vector& writable_paths); extern bool mseal_supported(); extern void seal_mappings(); extern bool supports_seccomp_notify(); @@ -139,7 +141,8 @@ void BenchmarkManagerDeleter::operator()(BenchmarkManager* p) const noexcept { BenchmarkManagerPtr make_benchmark_manager( int result_fd, const std::vector& signature, std::uint64_t seed, - bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket) + bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket, + const std::vector& writable_paths) { const std::size_t page_size = static_cast(getpagesize()); const std::size_t alloc_size = (BenchmarkManagerArenaSize + page_size - 1) & ~(page_size - 1); @@ -155,7 +158,7 @@ BenchmarkManagerPtr make_benchmark_manager( raw = new (mem) BenchmarkManager( static_cast(mem), alloc_size, result_fd, signature, seed, - discard, nvtx, landlock, mseal, allow_root, supervisor_socket); + discard, nvtx, landlock, mseal, allow_root, supervisor_socket, writable_paths); } catch (...) { // If construction throws, release the mmap'd region before propagating. if (munmap(mem, alloc_size) != 0) { @@ -170,7 +173,8 @@ BenchmarkManagerPtr make_benchmark_manager( BenchmarkManager::BenchmarkManager(std::byte* arena, std::size_t arena_size, int result_fd, const std::vector& signature, std::uint64_t seed, bool discard, - bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket) + bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket, + const std::vector& writable_paths) : mArena(arena), mResource(arena + sizeof(BenchmarkManager), arena_size - sizeof(BenchmarkManager), @@ -202,6 +206,7 @@ BenchmarkManager::BenchmarkManager(std::byte* arena, std::size_t arena_size, mNVTXEnabled = nvtx; mLandlock = landlock; + mWritablePaths = writable_paths; mSeal = mseal; mAllowRoot = allow_root; mDiscardCache = discard; @@ -225,6 +230,8 @@ BenchmarkManager::~BenchmarkManager() { for (auto& exp: mExpectedOutputs) cudaFree(exp.Value); } +static nb::tuple ensure_contiguous_tuple(const nb::tuple& tup); + std::pair, std::vector> BenchmarkManager::setup_benchmark(const nb::callable& generate_test_case, const nb::dict& kwargs, int repeats) { std::mt19937_64 rng(mSeed); std::uniform_int_distribution dist(0, std::numeric_limits::max()); @@ -244,14 +251,26 @@ std::pair, std::vector> BenchmarkManager::setu call_kwargs["seed"] = dist(rng); auto gen = nb::cast(generate_test_case(**call_kwargs)); - kernel_args[i] = nb::cast(gen[0]); - expected[i] = nb::cast(gen[1]); + kernel_args[i] = ensure_contiguous_tuple(nb::cast(gen[0])); + expected[i] = ensure_contiguous_tuple(nb::cast(gen[1])); } return std::make_pair(std::move(kernel_args), std::move(expected)); } bool can_convert_to_tensor(nb::handle obj) { - return nb::isinstance(obj); + return nb::isinstance(obj); +} + +static nb::tuple ensure_contiguous_tuple(const nb::tuple& tup) { + nb::list new_tup; + for (auto item : tup) { + if (nb::isinstance(item)) { + new_tup.append(nb::cast(item).attr("contiguous")()); + } else { + new_tup.append(item); + } + } + return nb::tuple(new_tup); } auto BenchmarkManager::make_shadow_args(const nb::tuple& args, cudaStream_t stream, @@ -332,7 +351,7 @@ void BenchmarkManager::install_protections() { // restrict access to file system if (mLandlock) - install_landlock(); + install_landlock(mWritablePaths); if (mSeal) { if (!mseal_supported()) { diff --git a/csrc/manager.h b/csrc/manager.h index 5a05de9..b183a13 100644 --- a/csrc/manager.h +++ b/csrc/manager.h @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -22,6 +23,9 @@ namespace nb = nanobind; using nb_cuda_array = nb::ndarray; +/// Broader type that accepts any CUDA ndarray (including noncontiguous), +/// used only for detection before forcing contiguity. +using nb_any_cuda_array = nb::ndarray; struct BenchmarkParameters { std::uint64_t Seed; @@ -43,7 +47,8 @@ using BenchmarkManagerPtr = std::unique_ptr& signature, std::uint64_t seed, - bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket); + bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket, + const std::vector& writable_paths); class BenchmarkManager { @@ -53,14 +58,15 @@ class BenchmarkManager { void send_report(); void clean_up(); private: - friend BenchmarkManagerPtr make_benchmark_manager(int result_fd, const std::vector& signature, std::uint64_t seed, bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket); + friend BenchmarkManagerPtr make_benchmark_manager(int result_fd, const std::vector& signature, std::uint64_t seed, bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket, const std::vector& writable_paths); friend BenchmarkManagerDeleter; /// `arena` is the mmap region that owns all memory for this object and its vectors. /// The BenchmarkManager must have been placement-newed into the front of that region; /// the rest is used as a monotonic PMR arena for internal vectors. BenchmarkManager(std::byte* arena, std::size_t arena_size, int result_fd, const std::vector& signature, std::uint64_t seed, - bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket); + bool discard, bool nvtx, bool landlock, bool mseal, bool allow_root, int supervisor_socket, + const std::vector& writable_paths); ~BenchmarkManager(); struct Expected { @@ -109,6 +115,7 @@ class BenchmarkManager { bool mNVTXEnabled = false; bool mDiscardCache = true; bool mLandlock = true; + std::vector mWritablePaths; bool mSeal = true; bool mAllowRoot = false; int mSupervisorSock = -1; diff --git a/csrc/supervisor.cpp b/csrc/supervisor.cpp index 385546e..d68d374 100644 --- a/csrc/supervisor.cpp +++ b/csrc/supervisor.cpp @@ -114,7 +114,12 @@ static bool handle_notification(int unotify_fd, const Config& cfg) { if (ioctl(unotify_fd, SECCOMP_IOCTL_NOTIF_RECV, &req) < 0) { if (errno == EINTR) return true; - if (errno == ENODEV) return false; + // ENODEV: all filter users gone. ENOENT: the notifying task died + // before we could RECV its notification (the same "target died" + // condition the NOTIF_SEND path below treats as benign). Both mean + // no live syscall is left to police -- the tracee has exited -- so + // stop the loop quietly instead of perror'ing teardown noise. + if (errno == ENODEV || errno == ENOENT) return false; perror("supervisor: SECCOMP_IOCTL_NOTIF_RECV"); return false; } diff --git a/python/pygpubench/__init__.py b/python/pygpubench/__init__.py index 46abda1..4c823f1 100644 --- a/python/pygpubench/__init__.py +++ b/python/pygpubench/__init__.py @@ -32,7 +32,8 @@ def _do_bench_impl(out_fd: "multiprocessing.connection.Connection", in_fd: "multiprocessing.connection.Connection", supervisor_sock: "socket.socket", qualname: str, test_generator: TestGeneratorInterface, test_args: dict, stream: int = None, discard: bool = True, - nvtx: bool = False, tb_conn: "multiprocessing.connection.Connection" = None, landlock=True, mseal=True, allow_root=False): + nvtx: bool = False, tb_conn: "multiprocessing.connection.Connection" = None, landlock=True, mseal=True, allow_root=False, + writable_paths=("/tmp",)): """ Benchmarks the kernel referred to by `qualname` against the test case returned by `test_generator`. :param out_fd: Writable file descriptor to which benchmark results are written. @@ -46,6 +47,7 @@ def _do_bench_impl(out_fd: "multiprocessing.connection.Connection", in_fd: "mult :param landlock: Whether to enable landlock. Enabled by default, prevents write access to the file system outside /tmp. :param mseal: Whether to enable memory sealing. Enabled by default, prevents making executable mappings writable. :param allow_root: Whether to allow the benchmark to run as root (opt-in via ``allow_root=True``). When run as root, the benchmark process's memory can be read through /proc/self/mem despite being protected. + :param writable_paths: Filesystem paths (and everything beneath them) the benchmark is allowed to write to when landlock is enabled. Defaults to ``("/tmp",)``. ``/dev`` is always writable (needed by e.g. triton); the rest of the filesystem stays read-only. """ if stream is None: import torch @@ -66,6 +68,7 @@ def _do_bench_impl(out_fd: "multiprocessing.connection.Connection", in_fd: "mult landlock, mseal, allow_root, + list(writable_paths), ) except BaseException: if tb_conn is not None: @@ -157,6 +160,7 @@ def do_bench_isolated( landlock = True, mseal = True, allow_root = False, + writable_paths = ("/tmp",), ) -> BenchmarkResult: """ Runs kernel benchmark (`do_bench_impl`) in a subprocess for proper isolation. @@ -204,6 +208,7 @@ def do_bench_isolated( landlock, mseal, allow_root, + writable_paths, ), )