Skip to content

Commit c6f012c

Browse files
authored
prevent AOTI library collisions between delegates (pytorch#21287)
Summary: The CUDA backend writes each compiled AOTInductor library to a temporary file, named from its content key plus the process id, before loading it. Two identical CUDA partitions can share the same content key, so both delegates computed the same path; loading the second library overwrote the first while it was still in use, and symbol lookup could then crash. Append a per-process atomic counter to the temporary file name so every delegate in a process gets a distinct path. The file is still removed when the delegate is destroyed, as before; only the name changes. Differential Revision: D113322459
1 parent fbd14e6 commit c6f012c

1 file changed

Lines changed: 11 additions & 3 deletions

File tree

backends/cuda/runtime/cuda_backend.cpp

Lines changed: 11 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -330,10 +330,18 @@ class ET_EXPERIMENTAL CudaBackend final
330330
so_blob_key.c_str(),
331331
static_cast<uint32_t>(aoti_dso_buffer.error()));
332332

333-
// Generate dynamic temporary file path
333+
// Generate a unique temporary file path. so_blob_key already selected the
334+
// blob above and is deliberately kept out of the filename: it can be an
335+
// untrusted, variable-length payload key, so embedding it risks path
336+
// traversal and cross-key collisions. Uniqueness comes from the pid
337+
// (across processes) and an atomic counter (within a process, since two
338+
// identical CUDA partitions would otherwise clobber each other's .so).
339+
static std::atomic<uint64_t> so_file_counter{0};
334340
filesystem::path temp_dir = filesystem::temp_directory_path();
335-
filesystem::path so_path =
336-
temp_dir / (so_blob_key + to_string(get_process_id()) + ".so");
341+
filesystem::path so_path = temp_dir /
342+
("executorch_cuda_" + to_string(get_process_id()) + "_" +
343+
to_string(so_file_counter.fetch_add(1, std::memory_order_relaxed)) +
344+
".so");
337345

338346
// Create a temporary file
339347
ofstream outfile(so_path, ios::binary);

0 commit comments

Comments
 (0)