Make MessagePassing interface thread-safe (#9001)

Fixes #8994

Author: rusty1s

CHANGELOG.md

@@ -13,6 +13,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Made `MessagePassing` interface thread-safe ([#9001](https://github.com/pyg-team/pytorch_geometric/pull/9001))
 - Breaking Change: Added support for `EdgeIndex` in `cugraph` GNN layers ([#8938](https://github.com/pyg-team/pytorch_geometric/pull/8937))
 - Added the `dim` arg to `torch.cross` calls ([#8918](https://github.com/pyg-team/pytorch_geometric/pull/8918))
 

benchmark/inference/README.md

@@ -7,7 +7,7 @@
    ```bash
    pip install ogb
    ```
-1. Install `autoconf` required for `jemalloc` setup
+1. Install `autoconf` required for `jemalloc` setup:
    ```bash
    sudo apt-get install autoconf
    ```

torch_geometric/nn/conv/message_passing.py

@@ -185,7 +185,6 @@ def __init__(
                     fuse=self.fuse,
                 )
 
-                # Cache to potentially disable later on:
                 self.__class__._orig_propagate = self.__class__.propagate
                 self.__class__._jinja_propagate = module.propagate
 
@@ -197,22 +196,30 @@ def __init__(
 
         # Optimize `edge_updater()` via `*.jinja` templates (if implemented):
         if (self.inspector.implements('edge_update')
-                and not self.edge_updater.__module__.startswith(jinja_prefix)
-                and self.inspector.can_read_source):
-            module = module_from_template(
-                module_name=f'{jinja_prefix}_edge_updater',
-                template_path=osp.join(root_dir, 'edge_updater.jinja'),
-                tmp_dirname='message_passing',
-                # Keyword arguments:
-                modules=self.inspector._modules,
-                collect_name='edge_collect',
-                signature=self._get_edge_updater_signature(),
-                collect_param_dict=self.inspector.get_param_dict(
-                    'edge_update'),
-            )
+                and not self.edge_updater.__module__.startswith(jinja_prefix)):
+            if self.inspector.can_read_source:
+
+                module = module_from_template(
+                    module_name=f'{jinja_prefix}_edge_updater',
+                    template_path=osp.join(root_dir, 'edge_updater.jinja'),
+                    tmp_dirname='message_passing',
+                    # Keyword arguments:
+                    modules=self.inspector._modules,
+                    collect_name='edge_collect',
+                    signature=self._get_edge_updater_signature(),
+                    collect_param_dict=self.inspector.get_param_dict(
+                        'edge_update'),
+                )
 
-            self.__class__.edge_updater = module.edge_updater
-            self.__class__.edge_collect = module.edge_collect
+                self.__class__._orig_edge_updater = self.__class__.edge_updater
+                self.__class__._jinja_edge_updater = module.edge_updater
+
+                self.__class__.edge_updater = module.edge_updater
+                self.__class__.edge_collect = module.edge_collect
+            else:
+                self.__class__._orig_edge_updater = self.__class__.edge_updater
+                self.__class__._jinja_edge_updater = (
+                    self.__class__.edge_updater)
 
         # Explainability:
         self._explain: Optional[bool] = None

torch_geometric/template.py

@@ -1,13 +1,11 @@
 import importlib
-import os
 import os.path as osp
 import sys
+import tempfile
 from typing import Any
 
 from jinja2 import Environment, FileSystemLoader
 
-from torch_geometric import get_home_dir
-
 
 def module_from_template(
     module_name: str,
@@ -23,13 +21,15 @@ def module_from_template(
     template = env.get_template(osp.basename(template_path))
     module_repr = template.render(**kwargs)
 
-    instance_dir = osp.join(get_home_dir(), tmp_dirname)
-    os.makedirs(instance_dir, exist_ok=True)
-    instance_path = osp.join(instance_dir, f'{module_name}.py')
-    with open(instance_path, 'w') as f:
-        f.write(module_repr)
+    with tempfile.NamedTemporaryFile(
+            mode='w',
+            prefix=f'{module_name}_',
+            suffix='.py',
+            delete=False,
+    ) as tmp:
+        tmp.write(module_repr)
 
-    spec = importlib.util.spec_from_file_location(module_name, instance_path)
+    spec = importlib.util.spec_from_file_location(module_name, tmp.name)
     assert spec is not None
     module = importlib.util.module_from_spec(spec)
     sys.modules[module_name] = module