From 0a8807eb2f8f87cbfdc49538b73ff498526adf66 Mon Sep 17 00:00:00 2001
From: Vukasin Milovanovic <vmilovanovic@nvidia.com>
Date: Mon, 1 Apr 2024 14:31:16 -0700
Subject: [PATCH] Apply the cuFile error work around to data_sink as well
 (#15335)

Issue #14140

Follow-up on https://github.com/rapidsai/cudf/pull/15293

Moving the `cudaFree(0)` call to a function called both by file `datasource` and `data_sink`.

Authors:
  - Vukasin Milovanovic (https://github.com/vuule)

Approvers:
  - David Wendt (https://github.com/davidwendt)
  - Yunsong Wang (https://github.com/PointKernel)
  - Nghia Truong (https://github.com/ttnghia)

URL: https://github.com/rapidsai/cudf/pull/15335
---
 cpp/src/io/utilities/data_sink.cpp         | 1 +
 cpp/src/io/utilities/datasource.cpp        | 6 +-----
 cpp/src/io/utilities/file_io_utilities.cpp | 8 ++++++++
 cpp/src/io/utilities/file_io_utilities.hpp | 3 +++
 4 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/cpp/src/io/utilities/data_sink.cpp b/cpp/src/io/utilities/data_sink.cpp
index 5557648ebbe..66905c5256f 100644
--- a/cpp/src/io/utilities/data_sink.cpp
+++ b/cpp/src/io/utilities/data_sink.cpp
@@ -36,6 +36,7 @@ class file_sink : public data_sink {
  public:
   explicit file_sink(std::string const& filepath)
   {
+    detail::force_init_cuda_context();
     _output_stream.open(filepath, std::ios::out | std::ios::binary | std::ios::trunc);
     if (!_output_stream.is_open()) { detail::throw_on_file_open_failure(filepath, true); }
 
diff --git a/cpp/src/io/utilities/datasource.cpp b/cpp/src/io/utilities/datasource.cpp
index 54e7c6bf1d6..d8dbd3614c8 100644
--- a/cpp/src/io/utilities/datasource.cpp
+++ b/cpp/src/io/utilities/datasource.cpp
@@ -43,12 +43,8 @@ class file_source : public datasource {
  public:
   explicit file_source(char const* filepath) : _file(filepath, O_RDONLY)
   {
+    detail::force_init_cuda_context();
     if (detail::cufile_integration::is_kvikio_enabled()) {
-      // Workaround for https://github.com/rapidsai/cudf/issues/14140, where cuFileDriverOpen errors
-      // out if no CUDA calls have been made before it. This is a no-op if the CUDA context is
-      // already initialized
-      cudaFree(0);
-
       _kvikio_file = kvikio::FileHandle(filepath);
       CUDF_LOG_INFO("Reading a file using kvikIO, with compatibility mode {}.",
                     _kvikio_file.is_compat_mode_on() ? "on" : "off");
diff --git a/cpp/src/io/utilities/file_io_utilities.cpp b/cpp/src/io/utilities/file_io_utilities.cpp
index 01090a43a0e..39031526fc8 100644
--- a/cpp/src/io/utilities/file_io_utilities.cpp
+++ b/cpp/src/io/utilities/file_io_utilities.cpp
@@ -34,6 +34,14 @@ namespace cudf {
 namespace io {
 namespace detail {
 
+void force_init_cuda_context()
+{
+  // Workaround for https://github.com/rapidsai/cudf/issues/14140, where cuFileDriverOpen errors
+  // out if no CUDA calls have been made before it. This is a no-op if the CUDA context is already
+  // initialized.
+  cudaFree(0);
+}
+
 [[noreturn]] void throw_on_file_open_failure(std::string const& filepath, bool is_create)
 {
   // save errno because it may be overwritten by subsequent calls
diff --git a/cpp/src/io/utilities/file_io_utilities.hpp b/cpp/src/io/utilities/file_io_utilities.hpp
index 0d5a5b218da..74a2ae53961 100644
--- a/cpp/src/io/utilities/file_io_utilities.hpp
+++ b/cpp/src/io/utilities/file_io_utilities.hpp
@@ -37,6 +37,9 @@ namespace detail {
 
 [[noreturn]] void throw_on_file_open_failure(std::string const& filepath, bool is_create);
 
+// Call before any cuFile API calls to ensure the CUDA context is initialized.
+void force_init_cuda_context();
+
 /**
  * @brief Class that provides RAII for file handling.
  */