Wait device task finish when python exit

mindspore/ccsrc/pipeline/jit/pipeline.cc

@@ -1788,6 +1788,10 @@ void MemoryRecycle() {
 
 void ClearResAtexit() {
   MS_LOG(INFO) << "Pipeline clear all resource";
+  // When the python process exits, the kernels on the device may not have finished executing.
+  device::KernelRuntimeManager::Instance().WaitTaskFinishOnDevice();
+  device::DeviceContextManager::GetInstance().WaitTaskFinishOnDevice();
+
   RecordExitStatus();
 #if ((defined ENABLE_CPU) && (!defined _WIN32))
   if (ps::PSContext::instance()->is_ps_mode() && ps::PSContext::instance()->is_worker()) {

mindspore/ccsrc/runtime/device/kernel_runtime_manager.cc

@@ -140,5 +140,20 @@ void KernelRuntimeManager::ReleaseKernelRuntime(const std::string &device_name,
   runtime->ReleaseDeviceRes();
   runtime_map_.erase(runtime_iter);
 }
+
+void KernelRuntimeManager::WaitTaskFinishOnDevice() const {
+  for (const auto &iter : runtime_map_) {
+    auto kernel_runtime = iter.second;
+    try {
+      if (kernel_runtime != nullptr && !kernel_runtime->SyncStream()) {
+        MS_LOG(ERROR) << "SyncStream failed";
+        return;
+      }
+    } catch (const std::exception &ex) {
+      MS_LOG(ERROR) << "SyncStream failed, exception:" << ex.what();
+      return;
+    }
+  }
+}
 }  // namespace device
 }  // namespace mindspore

mindspore/ccsrc/runtime/device/kernel_runtime_manager.h

@@ -40,6 +40,7 @@ class KernelRuntimeManager {
   void ReleaseKernelRuntime(const std::string &device_name, uint32_t device_id);
   void ClearRuntimeResource();
   void ClearGraphResource(uint32_t graph_id);
+  void WaitTaskFinishOnDevice() const;
 
  private:
   KernelRuntimeManager() = default;

mindspore/ccsrc/runtime/hardware/device_context_manager.cc

@@ -66,5 +66,20 @@ void DeviceContextManager::UpdateDeviceContextKey(const DeviceContextKey &old_ke
   handle.key() = new_key_str;
   (void)device_contexts_.insert(std::move(handle));
 }
+
+void DeviceContextManager::WaitTaskFinishOnDevice() const {
+  for (const auto &item : device_contexts_) {
+    auto device_context = item.second;
+    try {
+      if (device_context != nullptr && !device_context->SyncStream()) {
+        MS_LOG(ERROR) << "SyncStream failed";
+        return;
+      }
+    } catch (const std::exception &ex) {
+      MS_LOG(ERROR) << "SyncStream failed, exception:" << ex.what();
+      return;
+    }
+  }
+}
 }  // namespace device
 }  // namespace mindspore

mindspore/ccsrc/runtime/hardware/device_context_manager.h

@@ -39,6 +39,7 @@ class DeviceContextManager {
   DeviceContext *GetOrCreateDeviceContext(const DeviceContextKey &device_context_key);
   void UpdateDeviceContextKey(const DeviceContextKey &old_key, const DeviceContextKey &new_key);
   void ClearDeviceContexts();
+  void WaitTaskFinishOnDevice() const;
 
  private:
   DeviceContextManager() = default;