Apply the cuFile error work around to data_sink as well (#15335)

Issue #14140

Follow-up on #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: #15335

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); }
 

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");

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

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.
  */