Introduce BICLEANER_AI_THREADS to control number of threads

`--processes` option is now deprecated in favor of the environment
variable. The TensorFlow threads is now set prior to initialization to
avoid errors.

6 files changed, 58 insertions, 12 deletions

diff --git a/README.md b/README.md
index 880e66d..be46422 100644
--- a/README.md
+++ b/README.md
@@ -98,6 +98,8 @@ TF_CPP_MIN_LOG_LEVEL=0 bicleaner-ai-classify
 ```
 **WARNING**: If you are experiencing slow downs because Bicleaner AI is not running in the GPU, you should check those logs to see if TensorFlow is loading all the libraries correctly.
 
+
+
 ## Cleaning
 
 `bicleaner-ai-classify` aims at detecting noisy sentence pairs in a parallel corpus. It
@@ -360,6 +362,15 @@ Also note that, if a tokenizer command is used, the word frequencies need to be
 
 If no tokenization is available for your languages, you can disable these noise option that use tokenization and use fuzzy mathing noise: `--womit_ratio 0 --freq_ratio 0 --fuzzy_ratio 6`.
 
+## Setting the number of threads
+To set the maximum number of threads/processes to be used during training or classifying, `--processes` option is no longer available.
+You will need to set `BICLEANER_AI_THREADS` environment variable to the desired value.
+For example:
+```
+BICLEANER_AI_THREADS=12 bicleaner-ai-classify ...
+```
+If the variable is not set, the program will use all the available CPU cores.
+
 ## Speed
 A comparison of the speed in number of sentences per second between different types of models and hardware:
 
diff --git a/bicleaner_ai/bicleaner_ai_classifier.py b/bicleaner_ai/bicleaner_ai_classifier.py
index 0169ade..e038d25 100755
--- a/bicleaner_ai/bicleaner_ai_classifier.py
+++ b/bicleaner_ai/bicleaner_ai_classifier.py
@@ -3,6 +3,12 @@ import os
 # Suppress Tenssorflow logging messages unless log level is explictly set
 if 'TF_CPP_MIN_LOG_LEVEL' not in os.environ:
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+# Set Tensorflow max threads before initialization
+if 'BICLEANER_AI_THREADS' in os.environ:
+    threads = int(os.environ["BICLEANER_AI_THREADS"])
+    import tensorflow as tf
+    tf.config.threading.set_intra_op_parallelism_threads(threads)
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
 import sys
 import logging
 import traceback
@@ -32,11 +38,17 @@ def initialization():
     # Set up logging
     logging_setup(args)
     logging_level = logging.getLogger().level
-    import tensorflow as tf
 
-    # Set number of processes to be used by TensorFlow
-    tf.config.threading.set_intra_op_parallelism_threads(args.processes)
-    tf.config.threading.set_inter_op_parallelism_threads(args.processes)
+    # Warn about args.processes deprecation
+    if args.processes is not None:
+        logging.warging("--processes option is not available anymore, please use BICLEANER_AI_THREADS environment variable instead.")
+
+    # Set the number of processes from the environment variable
+    # or instead use all cores
+    if "BICLEANER_AI_THREADS" in os.environ and os.environ["BICLEANER_AI_THREADS"]:
+        args.processes = int(os.environ["BICLEANER_AI_THREADS"])
+    else:
+        args.processes = max(1, cpu_count()-1)
 
     # Load metadata YAML
     args = load_metadata(args, parser)
diff --git a/bicleaner_ai/bicleaner_ai_train.py b/bicleaner_ai/bicleaner_ai_train.py
index 459811d..2eb6972 100755
--- a/bicleaner_ai/bicleaner_ai_train.py
+++ b/bicleaner_ai/bicleaner_ai_train.py
@@ -3,6 +3,12 @@ import os
 # Suppress Tenssorflow logging messages unless log level is explictly set
 if 'TF_CPP_MIN_LOG_LEVEL' not in os.environ:
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+# Set Tensorflow max threads before initialization
+if 'BICLEANER_AI_THREADS' in os.environ:
+    threads = int(os.environ["BICLEANER_AI_THREADS"])
+    import tensorflow as tf
+    tf.config.threading.set_intra_op_parallelism_threads(threads)
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
 from tempfile import TemporaryFile, NamedTemporaryFile, gettempdir
 from multiprocessing import cpu_count
 from timeit import default_timer
@@ -50,7 +56,7 @@ def initialization():
     #groupO.add_argument('-f', '--source_word_freqs', type=argparse.FileType('r'), default=None, required=False, help="L language gzipped list of word frequencies")
     groupO.add_argument('-F', '--target_word_freqs', type=argparse.FileType('r'), default=None, required=False, help="R language gzipped list of word frequencies (needed for frequence based noise)")
     groupO.add_argument('--block_size', type=check_positive, default=10000, help="Sentence pairs per block when apliying multiprocessing in the noise function")
-    groupO.add_argument('-p', '--processes', type=check_positive, default=max(1, cpu_count()-1), help="Number of process to use")
+    groupO.add_argument('-p', '--processes', default=None, help="Option no longer available, please set BICLEANER_AI_THREADS environment variable")
     groupO.add_argument('-g', '--gpu', type=check_positive_or_zero, help="Which GPU use, starting from 0. Will set the CUDA_VISIBLE_DEVICES.")
     groupO.add_argument('--mixed_precision', action='store_true', default=False, help="Use mixed precision float16 for training")
     groupO.add_argument('--save_train', type=str, default=None, help="Save the generated training dataset into a file. If the file already exists the training dataset will be loaded from there.")
@@ -110,12 +116,17 @@ def initialization():
 
     if args.gpu is not None:
         os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
-    elif "CUDA_VISIBLE_DEVICES" not in os.environ or os.environ["CUDA_VISIBLE_DEVICES"] == "":
-        import psutil
-        cpus = psutil.cpu_count(logical=False)
-        # Set number of threads for CPU training
-        tf.config.threading.set_intra_op_parallelism_threads(min(cpus, args.processes))
-        tf.config.threading.set_inter_op_parallelism_threads(min(2, args.processes))
+
+    # Warn about args.processes deprecation
+    if args.processes is not None:
+        logging.warging("--processes option is not available anymore, please use BICLEANER_AI_THREADS environment variable instead.")
+
+    # Set the number of processes from the environment variable
+    # or instead use all cores
+    if "BICLEANER_AI_THREADS" in os.environ and os.environ["BICLEANER_AI_THREADS"]:
+        args.processes = int(os.environ["BICLEANER_AI_THREADS"])
+    else:
+        args.processes = max(1, cpu_count()-1)
 
     if args.mixed_precision:
         from tensorflow.keras import mixed_precision
diff --git a/bicleaner_ai/classify.py b/bicleaner_ai/classify.py
index fbb2a51..0061702 100644
--- a/bicleaner_ai/classify.py
+++ b/bicleaner_ai/classify.py
@@ -42,7 +42,7 @@ def argument_parser():
     groupO.add_argument("--scol", default=3 if not header else "src_text", type=check_positive if not header else str, help ="Source sentence column (starting in 1). The name of the field is expected instead of the position if --header is set")
     groupO.add_argument("--tcol", default=4 if not header else "trg_text", type=check_positive if not header else str, help ="Target sentence column (starting in 1). The name of the field is expected instead of the position if --header is set")
     groupO.add_argument('-b', '--block_size', type=int, default=1000, help="Sentence pairs per block")
-    groupO.add_argument('-p', '--processes', type=int, default=0, help="Number of processes to use")
+    groupO.add_argument('-p', '--processes', default=None, help="Option no longer available, please set BICLEANER_AI_THREADS environment variable")
     groupO.add_argument('--batch_size', type=int, default=32, help="Sentence pairs per block")
 
     groupO.add_argument('--tmp_dir', default=gettempdir(), help="Temporary directory where creating the temporary files of this program")
diff --git a/scripts/bicleaner-ai-classify b/scripts/bicleaner-ai-classify
index 843b86d..8f5a678 100755
--- a/scripts/bicleaner-ai-classify
+++ b/scripts/bicleaner-ai-classify
@@ -4,6 +4,12 @@ import os
 # Suppress Tenssorflow logging messages unless log level is explictly set
 if 'TF_CPP_MIN_LOG_LEVEL' not in os.environ:
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+# Set Tensorflow max threads before initialization
+if 'BICLEANER_AI_THREADS' in os.environ:
+    threads = int(os.environ["BICLEANER_AI_THREADS"])
+    import tensorflow as tf
+    tf.config.threading.set_intra_op_parallelism_threads(threads)
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
 import sys
 import traceback
 import logging
diff --git a/scripts/bicleaner-ai-train b/scripts/bicleaner-ai-train
index 408d9de..896aada 100755
--- a/scripts/bicleaner-ai-train
+++ b/scripts/bicleaner-ai-train
@@ -3,6 +3,12 @@ import os
 # Suppress Tenssorflow logging messages unless log level is explictly set
 if 'TF_CPP_MIN_LOG_LEVEL' not in os.environ:
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+# Set Tensorflow max threads before initialization
+if 'BICLEANER_AI_THREADS' in os.environ:
+    threads = int(os.environ["BICLEANER_AI_THREADS"])
+    import tensorflow as tf
+    tf.config.threading.set_intra_op_parallelism_threads(threads)
+    tf.config.threading.set_inter_op_parallelism_threads(threads)
 import sys
 import bicleaner_ai.bicleaner_ai_train as train