fix: GPU selection on both CUDA and MPS

Signed-off-by: Luca Foppiano <luca@foppiano.org>

Author: lfoppiano

delft/sequenceLabelling/data_loader.py

@@ -22,6 +22,17 @@
 from delft.utilities.Utilities import len_until_first_pad, truncate_batch_values
 
 
+def _worker_init_fn(worker_id):
+    # Each fork-mode worker inherits the parent's LMDB env handle, which is unsafe.
+    # Re-open per worker so each gets an independent reader-locktable slot.
+    info = torch.utils.data.get_worker_info()
+    if info is None:
+        return
+    embeddings = getattr(info.dataset, "embeddings", None)
+    if embeddings is not None and hasattr(embeddings, "reopen_lmdb"):
+        embeddings.reopen_lmdb()
+
+
 def collate_fn(batch):
     """
     Custom collate function to handle variable-length sequences.
@@ -431,6 +442,7 @@ def create_dataloader(
         num_workers=num_workers,
         pin_memory=pin_memory,
         collate_fn=collate_fn,
+        worker_init_fn=_worker_init_fn if num_workers > 0 else None,
     )
 
 

delft/sequenceLabelling/tagger.py

@@ -5,6 +5,7 @@
 
 from delft.sequenceLabelling.data_loader import create_dataloader
 from delft.utilities.Tokenizer import tokenizeAndFilter
+from delft.utilities.Utilities import pick_device
 
 
 class Tagger(object):
@@ -13,7 +14,7 @@ def __init__(self, model, model_config, embeddings=None, preprocessor=None, devi
         self.preprocessor = preprocessor
         self.model_config = model_config
         self.embeddings = embeddings
-        self.device = device if device else torch.device("cpu")
+        self.device = pick_device(device)
 
     def tag(self, texts, output_format, features=None):
         if output_format == "json":

delft/sequenceLabelling/trainer.py

@@ -19,6 +19,7 @@
 from delft.sequenceLabelling.config import ModelConfig, TrainingConfig
 from delft.sequenceLabelling.evaluation import classification_report
 from delft.sequenceLabelling.preprocess import Preprocessor
+from delft.utilities.Utilities import pick_device
 
 # Default file names
 DEFAULT_WEIGHT_FILE_NAME = "model_weights.pt"
@@ -150,7 +151,7 @@ def __init__(
 
         # Set device
         if device is None:
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            self.device = pick_device()
         else:
             self.device = torch.device(device)
 

delft/sequenceLabelling/wrapper.py

@@ -23,6 +23,7 @@
 from delft.sequenceLabelling.data_loader import create_dataloader
 from delft.sequenceLabelling.evaluation import classification_report
 from delft.sequenceLabelling.models import get_model
+from delft.utilities.Utilities import pick_device
 from delft.sequenceLabelling.preprocess import Preprocessor, prepare_preprocessor
 from delft.sequenceLabelling.trainer import (
     CONFIG_FILE_NAME,
@@ -97,10 +98,7 @@ def __init__(
             self.nb_workers = nb_workers
 
         # Set device
-        if device is None:
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        else:
-            self.device = torch.device(device)
+        self.device = pick_device(device)
 
         word_emb_size = 0
         self.embeddings = None

delft/textClassification/data_loader.py

@@ -10,6 +10,17 @@
 from delft.textClassification.preprocess import to_indices_single, to_vector_single
 
 
+def _worker_init_fn(worker_id):
+    # Each fork-mode worker inherits the parent's LMDB env handle, which is unsafe.
+    # Re-open per worker so each gets an independent reader-locktable slot.
+    info = torch.utils.data.get_worker_info()
+    if info is None:
+        return
+    embeddings = getattr(info.dataset, "embeddings", None)
+    if embeddings is not None and hasattr(embeddings, "reopen_lmdb"):
+        embeddings.reopen_lmdb()
+
+
 class TextClassificationDataset(Dataset):
     """
     Dataset for text classification.
@@ -132,6 +143,7 @@ def create_dataloader(
         batch_size=batch_size,
         shuffle=shuffle,
         num_workers=num_workers,
+        worker_init_fn=_worker_init_fn if num_workers > 0 else None,
     )
 
     return loader

delft/textClassification/wrapper.py

@@ -13,6 +13,7 @@
 from delft.textClassification.trainer import Trainer
 from delft.utilities.Embeddings import Embeddings, load_resource_registry
 from delft.utilities.misc import print_parameters, to_wandb_table
+from delft.utilities.Utilities import pick_device
 
 # File names for saving/loading
 PREPROCESSOR_FILE = "preprocessor.json"
@@ -87,10 +88,7 @@ def __init__(
         self.report_to_wandb = report_to_wandb
         self.wandb = None
 
-        if device is None:
-            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        else:
-            self.device = torch.device(device)
+        self.device = pick_device(device)
 
         self.registry = load_resource_registry(os.path.join(DELFT_PROJECT_DIR, "resources-registry.json"))
 

delft/utilities/Utilities.py

@@ -1,12 +1,14 @@
 # some convenient methods for all models
 import numpy as np
 import regex as re
+import torch
 
 # seed is fixed for reproducibility
 from numpy.random import seed
 
 seed(7)
 import argparse
+import os
 import os.path
 import shutil
 from urllib.parse import urlparse
@@ -16,6 +18,46 @@
 from tqdm import tqdm
 
 
+def best_device() -> torch.device:
+    """
+    Pick the best available compute device: CUDA → MPS (Apple Silicon) → CPU.
+
+    When MPS is selected, transparently enables CPU fallback for ops that don't
+    have MPS kernels yet, so unsupported ops degrade gracefully instead of
+    raising NotImplementedError mid-training.
+    """
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if torch.backends.mps.is_available():
+        os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def pick_device(device=None) -> torch.device:
+    """
+    Resolve a device (auto-pick when device=None, else honor the caller's choice)
+    and print a one-line summary so the user can see which compute is in use.
+    """
+    if device is None:
+        d = best_device()
+    elif isinstance(device, torch.device):
+        d = device
+    else:
+        d = torch.device(device)
+
+    if d.type == "cuda":
+        n = torch.cuda.device_count()
+        name = torch.cuda.get_device_name(d)
+        plural = "" if n == 1 else "s"
+        print(f"Running on {d} ({name}); {n} CUDA device{plural} available")
+    elif d.type == "mps":
+        print(f"Running on {d} (Apple Silicon GPU via Metal Performance Shaders)")
+    else:
+        print(f"Running on {d}")
+    return d
+
+
 def truncate_batch_values(batch_values: list, max_sequence_length: int) -> list:
     return [row[:max_sequence_length] for row in batch_values]