fix: Secure torch model inits with global locks (#120)

Signed-off-by: Christoph Auer <[email protected]>

Author: cau-git

docling_ibm_models/code_formula_model/code_formula_predictor.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MIT
 #
 import logging
+import threading
 from typing import List, Optional, Union
 
 import numpy as np
@@ -17,6 +18,9 @@
 
 _log = logging.getLogger(__name__)
 
+# Global lock for model initialization to prevent threading issues
+_model_init_lock = threading.Lock()
+
 
 class StopOnString(StoppingCriteria):
     def __init__(self, tokenizer, stop_string):
@@ -80,13 +84,17 @@ def __init__(
         if device == "cpu":
             torch.set_num_threads(self._num_threads)
 
-        self._tokenizer = AutoTokenizer.from_pretrained(
-            artifacts_path, use_fast=True, padding_side="left"
-        )
-        self._model = SamOPTForCausalLM.from_pretrained(artifacts_path).to(self._device)
-        self._model.eval()
+        # Use lock to prevent threading issues during model initialization
+        with _model_init_lock:
+            self._tokenizer = AutoTokenizer.from_pretrained(
+                artifacts_path, use_fast=True, padding_side="left"
+            )
+            self._model = SamOPTForCausalLM.from_pretrained(artifacts_path).to(
+                self._device
+            )
+            self._model.eval()
 
-        self._image_processor = SamOptImageProcessor.from_pretrained(artifacts_path)
+            self._image_processor = SamOptImageProcessor.from_pretrained(artifacts_path)
 
         _log.debug("CodeFormulaModel settings: {}".format(self.info()))
 

docling_ibm_models/document_figure_classifier_model/document_figure_classifier_predictor.py

@@ -3,6 +3,7 @@
 # SPDX-License-Identifier: MIT
 #
 import logging
+import threading
 from typing import List, Tuple, Union
 
 import numpy as np
@@ -13,6 +14,9 @@
 
 _log = logging.getLogger(__name__)
 
+# Global lock for model initialization to prevent threading issues
+_model_init_lock = threading.Lock()
+
 
 class DocumentFigureClassifierPredictor:
     r"""
@@ -85,22 +89,23 @@ def __init__(
         if device == "cpu":
             torch.set_num_threads(self._num_threads)
 
-        model = AutoModelForImageClassification.from_pretrained(artifacts_path)
-        self._model = model.to(device)
-        self._model.eval()
-
-        self._image_processor = transforms.Compose(
-            [
-                transforms.Resize((224, 224)),
-                transforms.ToTensor(),
-                transforms.Normalize(
-                    mean=[0.485, 0.456, 0.406],
-                    std=[0.47853944, 0.4732864, 0.47434163],
-                ),
-            ]
-        )
-
-        config = AutoConfig.from_pretrained(artifacts_path)
+        with _model_init_lock:
+            model = AutoModelForImageClassification.from_pretrained(artifacts_path)
+            self._model = model.to(device)
+            self._model.eval()
+
+            self._image_processor = transforms.Compose(
+                [
+                    transforms.Resize((224, 224)),
+                    transforms.ToTensor(),
+                    transforms.Normalize(
+                        mean=[0.485, 0.456, 0.406],
+                        std=[0.47853944, 0.4732864, 0.47434163],
+                    ),
+                ]
+            )
+
+            config = AutoConfig.from_pretrained(artifacts_path)
 
         self._classes = list(config.id2label.values())
         self._classes.sort()

docling_ibm_models/layoutmodel/layout_predictor.py

@@ -4,6 +4,7 @@
 #
 import logging
 import os
+import threading
 from collections.abc import Iterable
 from typing import Set, Union
 
@@ -15,6 +16,9 @@
 
 _log = logging.getLogger(__name__)
 
+# Global lock for model initialization to prevent threading issues
+_model_init_lock = threading.Lock()
+
 
 class LayoutPredictor:
     """
@@ -87,10 +91,13 @@ def __init__(
         processor_config = os.path.join(artifact_path, "preprocessor_config.json")
         model_config = os.path.join(artifact_path, "config.json")
         self._image_processor = RTDetrImageProcessor.from_json_file(processor_config)
-        self._model = RTDetrForObjectDetection.from_pretrained(
-            artifact_path, config=model_config
-        ).to(self._device)
-        self._model.eval()
+
+        # Use lock to prevent threading issues during model initialization
+        with _model_init_lock:
+            self._model = RTDetrForObjectDetection.from_pretrained(
+                artifact_path, config=model_config
+            ).to(self._device)
+            self._model.eval()
 
         _log.debug("LayoutPredictor settings: {}".format(self.info()))
 

docling_ibm_models/tableformer/data_management/tf_predictor.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+import threading
 from itertools import groupby
 from pathlib import Path
 
@@ -35,6 +36,9 @@
 
 logger = s.get_custom_logger(__name__, LOG_LEVEL)
 
+# Global lock for model initialization to prevent threading issues
+_model_init_lock = threading.Lock()
+
 
 class bcolors:
     HEADER = "\033[95m"
@@ -175,34 +179,39 @@ def _load_model(self):
         """
 
         self._model_type = self._config["model"]["type"]
-        model = TableModel04_rs(self._config, self._init_data, self._device)
-
-        if model is None:
-            err_msg = "Not able to initiate a model for {}".format(self._model_type)
-            self._log().error(err_msg)
-            raise ValueError(err_msg)
-
-        self._remove_padding = False
-        if self._model_type == "TableModel02":
-            self._remove_padding = True
-
-        # Load model from safetensors
-        save_dir = self._config["model"]["save_dir"]
-        models_fn = glob.glob(f"{save_dir}/tableformer_*.safetensors")
-        if not models_fn:
-            err_msg = "Not able to find a model file for {}".format(self._model_type)
-            self._log().error(err_msg)
-            raise ValueError(err_msg)
-        model_fn = models_fn[
-            0
-        ]  # Take the first tableformer safetensors file inside the save_dir
-        missing, unexpected = load_model(model, model_fn, device=self._device)
-        if missing or unexpected:
-            err_msg = "Not able to load the model weights for {}".format(
-                self._model_type
-            )
-            self._log().error(err_msg)
-            raise ValueError(err_msg)
+
+        # Use lock to prevent threading issues during model initialization
+        with _model_init_lock:
+            model = TableModel04_rs(self._config, self._init_data, self._device)
+
+            if model is None:
+                err_msg = "Not able to initiate a model for {}".format(self._model_type)
+                self._log().error(err_msg)
+                raise ValueError(err_msg)
+
+            self._remove_padding = False
+            if self._model_type == "TableModel02":
+                self._remove_padding = True
+
+            # Load model from safetensors
+            save_dir = self._config["model"]["save_dir"]
+            models_fn = glob.glob(f"{save_dir}/tableformer_*.safetensors")
+            if not models_fn:
+                err_msg = "Not able to find a model file for {}".format(
+                    self._model_type
+                )
+                self._log().error(err_msg)
+                raise ValueError(err_msg)
+            model_fn = models_fn[
+                0
+            ]  # Take the first tableformer safetensors file inside the save_dir
+            missing, unexpected = load_model(model, model_fn, device=self._device)
+            if missing or unexpected:
+                err_msg = "Not able to load the model weights for {}".format(
+                    self._model_type
+                )
+                self._log().error(err_msg)
+                raise ValueError(err_msg)
 
         return model