Fix intermittent segfault on interpreter shutdown (#1513)

* Remove static reference to the `cupy` python module in `CupyUtil` which was causing a segfault on interpreter shutdown.
 
## By Submitting this PR I confirm:
- I am familiar with the [Contributing Guidelines](https://github.com/nv-morpheus/Morpheus/blob/main/docs/source/developer_guide/contributing.md).
- When the PR is ready for review, new or existing tests cover these changes.
- When the PR is ready for review, the documentation is up to date with these changes.

Authors:
  - David Gardner (https://github.com/dagardner-nv)

Approvers:
  - Devin Robison (https://github.com/drobison00)

URL: #1513

morpheus/_lib/include/morpheus/utilities/cupy_util.hpp

@@ -43,8 +43,6 @@ struct CupyUtil
     using tensor_map_t    = std::map<std::string, TensorObject>;
     using py_tensor_map_t = std::map<std::string, pybind11::object>;
 
-    static pybind11::object cp_module;  // handle to cupy module
-
     /**
      * @brief Import and return the cupy module. Requires GIL to have already been aqcuired.
      *

morpheus/_lib/src/utilities/cupy_util.cpp

@@ -23,8 +23,7 @@
 #include "morpheus/utilities/tensor_util.hpp"
 
 #include <cuda_runtime.h>
-#include <glog/logging.h>  // for COMPACT_GOOGLE_LOG_FATAL, DCHECK, LogMessageFatal
-#include <pybind11/cast.h>
+#include <glog/logging.h>         // for COMPACT_GOOGLE_LOG_FATAL, DCHECK, LogMessageFatal
 #include <pybind11/functional.h>  // IWYU pragma: keep
 #include <pybind11/gil.h>         // IWYU pragma: keep
 #include <pybind11/pybind11.h>
@@ -33,7 +32,6 @@
 #include <rmm/cuda_stream_view.hpp>  // for cuda_stream_per_thread
 #include <rmm/device_buffer.hpp>     // for device_buffer
 
-#include <array>    // for array
 #include <cstdint>  // for uintptr_t
 #include <memory>   // for make_shared
 #include <optional>
@@ -46,20 +44,10 @@ namespace morpheus {
 
 namespace py = pybind11;
 
-pybind11::object CupyUtil::cp_module = pybind11::none();
-
 pybind11::module_ CupyUtil::get_cp()
 {
     DCHECK(PyGILState_Check() != 0);
-
-    if (cp_module.is_none())
-    {
-        cp_module = pybind11::module_::import("cupy");
-    }
-
-    auto m = pybind11::cast<pybind11::module_>(cp_module);
-
-    return m;
+    return pybind11::cast<pybind11::module_>(pybind11::module_::import("cupy"));
 }
 
 bool CupyUtil::is_cupy_array(pybind11::object test_obj)

tests/test_tensor_memory.py

@@ -220,3 +220,21 @@ def test_tensorindex_bug(config: Config, tensor_cls: type, shape: typing.Tuple[i
     tensor_a = mem.get_tensor('a')
     assert tensor_a.shape == shape
     assert tensor_a.nbytes == shape[0] * shape[1] * 4
+
+
+def test_tensor_update(config: Config):
+    tensor_data = {
+        "input_ids": cp.array([1, 2, 3]), "input_mask": cp.array([1, 1, 1]), "segment_ids": cp.array([0, 0, 1])
+    }
+    tensor_memory = TensorMemory(count=3, tensors=tensor_data)
+
+    # Update tensors with new data
+    new_tensors = {
+        "input_ids": cp.array([4, 5, 6]), "input_mask": cp.array([1, 0, 1]), "segment_ids": cp.array([1, 1, 0])
+    }
+
+    tensor_memory.set_tensors(new_tensors)
+
+    for (key, cp_arr) in new_tensors.items():
+        tensor = tensor_memory.get_tensor(key)
+        cp.allclose(tensor, cp_arr)