From a6b7895612ba94e7fbe34a6195b87b82a2d157b8 Mon Sep 17 00:00:00 2001
From: ZL Deng <dawnmsg@gmail.com>
Date: Mon, 7 Mar 2022 13:08:20 +0100
Subject: [PATCH] solved the SLURM compatability issue

---
 README.md                  |  8 ++++++--
 ribodetector/detect_cpu.py | 20 ++++++++++++--------
 setup.py                   |  5 +++--
 3 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index f077e0b..3ba8b04 100644
--- a/README.md
+++ b/README.md
@@ -17,10 +17,12 @@ conda create -n ribodetector python=3.8
 conda activate ribodetector
 ```
 
-#### 2. Install `pytorch`
+#### 2. Install `pytorch` in the ribodetector env if GPU is available
 
 To install `pytorch` compatible with your CUDA version, please fellow this instruction:
-https://pytorch.org/get-started/locally/. Our code was tested with `pytorch v1.7` and `v1.7.1`.
+https://pytorch.org/get-started/locally/. Our code was tested with `pytorch v1.7`, `v1.7.1`, `v1.10.2`.
+
+Note: you can skip this step if you don't use GPU
 
 ### Installation
 
@@ -106,6 +108,8 @@ ribodetector_cpu -t 20 \
 ```
 The above command line excutes ribodetector for paired-end reads with length 100 using 20 CPU cores.
 
+When using SLURM job submission system, you need to specify `--cpus-per-task` to the number you CPU cores you need and set `--threads-per-core` to 1.
+
 #### Full help
 
 ```shell
diff --git a/ribodetector/detect_cpu.py b/ribodetector/detect_cpu.py
index cd87db0..fd8d24f 100644
--- a/ribodetector/detect_cpu.py
+++ b/ribodetector/detect_cpu.py
@@ -5,7 +5,7 @@
 Created Date: January 1st 2020
 Author: ZL Deng <dawnmsg(at)gmail.com>
 ---------------------------------------
-Last Modified: 6th December 2020 11:02:19 pm
+Last Modified: 7th March 2022 12:02:19 pm
 '''
 
 import os
@@ -70,19 +70,23 @@ def load_model(self):
         self.model_file = os.path.join(
             cd, self.config['state_file'][model_file_ext]).replace('.pth', '.onnx')
 
-        # self.model = onnxruntime.InferenceSession(self.model_file)
+        self.logger.info('Using high {} model file: {}{}{}{} on CPU'.format(model_file_ext.upper(),
+                                                                            colors.BOLD,
+                                                                            colors.OKCYAN,
+                                                                            self.model_file,
+                                                                            colors.ENDC))
+        
         so = onnxruntime.SessionOptions()
-        # so.intra_op_num_threads = 2
+        so.intra_op_num_threads = 1
+        so.inter_op_num_threads = 1
+        # so.so.enable_mem_pattern = False
+        # so.enable_cpu_mem_arena = False
         # so.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL
         so.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
 
         self.model = onnxruntime.InferenceSession(self.model_file, so)
 
-        self.logger.info('Using high {} model file: {}{}{}{} on CPU'.format(model_file_ext.upper(),
-                                                                            colors.BOLD,
-                                                                            colors.OKCYAN,
-                                                                            self.model_file,
-                                                                            colors.ENDC))
+        
 
     def run(self):
         """
diff --git a/setup.py b/setup.py
index a9a9697..48e891c 100644
--- a/setup.py
+++ b/setup.py
@@ -11,13 +11,14 @@
     "tqdm",
     "numpy",
     "biopython",
-    "onnxruntime >= 1.10.0",
-    "torch >= 1.7.1",
+    "onnxruntime >= 1.10.0"
+    # "torch >= 1.7.1",
 ]
 
 setup(
     name="ribodetector",
     version="0.2.4",
+    python_requires=">=3.8",
     author="Z-L Deng",
     author_email="dawnmsg@gmail.com",
     description="Accurate and rapid RiboRNA sequences Detector based on deep learning.",