From 2d8bb9780f5b23da65abee37fd612a1fa2f2f4e8 Mon Sep 17 00:00:00 2001
From: Mengqing Cao <cmq0113@163.com>
Date: Fri, 17 May 2024 10:21:01 +0800
Subject: [PATCH] fix segmentfault in (#2506) (#2530)

- fix segmentfault in Kunpeng (#2506)
  - avoids the repeated initialization of deepspeed causing by (#2507)
---
 wenet/dataset/processor.py | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/wenet/dataset/processor.py b/wenet/dataset/processor.py
index 5131a13ac6..4d3a809612 100644
--- a/wenet/dataset/processor.py
+++ b/wenet/dataset/processor.py
@@ -37,6 +37,21 @@
 
 logging.getLogger('langid').setLevel(logging.INFO)
 
+import os
+try:
+    cpu_info = os.popen("lscpu | grep 'Vendor ID'").read()
+    # 0x48 --> HiSilicon
+    if (cpu_info.rstrip().split(" ")[-1] == "0x48"):
+        # NOTE (MengqingCao): set number of threads in the subprocesses to 1
+        # Why? There may be some operators ultilizing multi-threads in processor,
+        # causing possibly deadlock in Kunpeng.
+        # Similar issue in PyTorch: https://github.com/pytorch/pytorch/issues/45198
+        torch.set_num_threads(1)
+except Exception as ex:
+    logging.warning('Failed to set number of thread in Kunpeng, \
+        this may cause segmentfault while dataloading, \
+        ignore this warning if you are not using Kunpeng')
+
 
 class UrlOpenError(Exception):