explicitly set *_NUM_THREADS in Returnn run (#424)

Author: JackTemaki

returnn/extract_prior.py

@@ -163,7 +163,10 @@ def create_files(self):
 
     def run(self):
         cmd = self._get_run_cmd()
-        sp.check_call(cmd)
+        env = os.environ.copy()
+        env["OMP_NUM_THREADS"] = str(self.rqmt["cpu"])
+        env["MKL_NUM_THREADS"] = str(self.rqmt["cpu"])
+        sp.check_call(cmd, env=env)
 
         merged_scores = np.loadtxt(self.out_prior_txt_file.get_path(), delimiter=" ")
 

returnn/forward.py

@@ -124,7 +124,10 @@ def run(self):
             ]
 
             try:
-                sp.check_call(call, cwd=d)
+                env = os.environ.copy()
+                env["OMP_NUM_THREADS"] = str(self.rqmt["cpu"])
+                env["MKL_NUM_THREADS"] = str(self.rqmt["cpu"])
+                sp.check_call(call, cwd=d, env=env)
             except Exception as e:
                 print("Run crashed - copy temporary work folder as 'crash_dir'")
                 shutil.copytree(d, "crash_dir")

returnn/search.py

@@ -127,7 +127,10 @@ def run(self):
             os.path.join(self.returnn_root.get_path(), "rnn.py"),
             self.out_returnn_config_file.get_path(),
         ]
-        sp.check_call(call)
+        env = os.environ.copy()
+        env["OMP_NUM_THREADS"] = str(self.rqmt["cpu"])
+        env["MKL_NUM_THREADS"] = str(self.rqmt["cpu"])
+        sp.check_call(call, env=env)
 
     @classmethod
     def create_returnn_config(

returnn/training.py

@@ -353,7 +353,10 @@ def run(self):
                         print("Cannot read:", exc)
             sys.stdout.flush()
 
-        sp.check_call(self._get_run_cmd())
+        env = os.environ.copy()
+        env["OMP_NUM_THREADS"] = str(self.rqmt["cpu"])
+        env["MKL_NUM_THREADS"] = str(self.rqmt["cpu"])
+        sp.check_call(self._get_run_cmd(), env=env)
 
         lrf = self.returnn_config.get("learning_rate_file", "learning_rates")
         self._relink(lrf, self.out_learning_rates.get_path())