fix the closing pool only for tf < 2.10.0

kermitt2 · Sep 1, 2023 · f8df013 · f8df013
1 parent 8a03944
commit f8df013
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 5 deletions.
diff --git a/delft/sequenceLabelling/tagger.py b/delft/sequenceLabelling/tagger.py
@@ -2,6 +2,7 @@
 
 import numpy as np
 import tensorflow as tf
+from packaging import version
 
 from delft.sequenceLabelling.data_generator import DataGeneratorTransformers
 from delft.sequenceLabelling.preprocess import Preprocessor
@@ -23,16 +24,16 @@ def __init__(self,
         self.model_config = model_config
         self.embeddings = embeddings
 
-
     def tag(self, texts, output_format, features=None, multi_gpu=False):
         if multi_gpu:
             strategy = tf.distribute.MirroredStrategy()
             print('Running with multi-gpu. Number of devices: {}'.format(strategy.num_replicas_in_sync))
 
             # This trick avoid an exception being through when the --multi-gpu approach is used on a single GPU system.
             # It might be removed with TF 2.10 https://github.com/tensorflow/tensorflow/issues/50487
-            import atexit
-            atexit.register(strategy._extended._collective_ops._pool.close) # type: ignore
+            if version.parse(tf.__version__) < version.parse('2.10.0'):
+                import atexit
+                atexit.register(strategy._extended._collective_ops._pool.close) # type: ignore
 
             with strategy.scope():
                 return self.tag_(texts, output_format, features)

diff --git a/delft/sequenceLabelling/wrapper.py b/delft/sequenceLabelling/wrapper.py
@@ -1,5 +1,7 @@
 import os
 
+from packaging import version
+
 # ask tensorflow to be quiet and not print hundred lines of logs
 from delft.utilities.Transformer import TRANSFORMER_CONFIG_FILE_NAME, DEFAULT_TRANSFORMER_TOKENIZER_DIR
 from delft.utilities.misc import print_parameters
@@ -473,8 +475,9 @@ def tag(self, texts, output_format, features=None, batch_size=None, multi_gpu=Fa
 
             # This trick avoid an exception being through when the --multi-gpu approach is used on a single GPU system.
             # It might be removed with TF 2.10 https://github.com/tensorflow/tensorflow/issues/50487
-            import atexit
-            atexit.register(strategy._extended._collective_ops._pool.close) # type: ignore
+            if version.parse(tf.__version__) < version.parse('2.10.0'):
+                import atexit
+                atexit.register(strategy._extended._collective_ops._pool.close) # type: ignore
 
             with strategy.scope():
                 return self.tag_(texts, output_format, features, batch_size)