Repository: pyg-team/pytorch_geometric

Issue:
Possible overwriting scenario with Jinja
### 🐛 Describe the bug

I am getting the following error, not always but from time to time:
```
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/conv/cg_conv.py", line 57, in __init__
    super().__init__(aggr=aggr, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/conv/message_passing.py", line 193, in __init__
    self.__class__._jinja_propagate = module.propagate
AttributeError: module 'torch_geometric.nn.conv.cg_conv_CGConv_propagate' has no attribute 'propagate'
```
I am using pyg in a parallel setting with mpi. I think there is possibility of overwriting when pyg uses jina template here:
https://github.com/pyg-team/pytorch_geometric/blob/9b660ac6ca882604d1ae521912d20ded1d180ecf/torch_geometric/nn/conv/message_passing.py#L170

I put some following debug message around line 186:
```
print ("module:", module)
print ("dir(module):", dir(module))
```
Here is what I got from one process:
```
module: <module 'torch_geometric.nn.conv.pna_conv_PNAConv_propagate' from '/root/.cache/pyg/message_passing/torch_geometric.nn.conv.pna_conv_PNAConv_propagate.py'>
dir(module): ['__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__']
```
And this is the output from the other process:
```
module: <module 'torch_geometric.nn.conv.pna_conv_PNAConv_propagate' from '/root/.cache/pyg/message_passing/torch_geometric.nn.conv.pna_conv_PNAConv_propagate.py'>
dir(module): ['Adj', 'Any', 'Callable', 'CollectArgs', 'DataLoader', 'DegreeScalerAggregation', 'Dict', 'Linear', 'List', 'MessagePassing', 'ModuleList', 'NamedTuple', 'OptTensor', 'Optional', 'PNAConv', 'Sequential', 'Size', 'SparseTensor', 'Tensor', 'Union', '__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__spec__', 'activation_resolver', 'collect', 'degree', 'is_compiling', 'is_sparse', 'is_torch_sparse_tensor', 'propagate', 'ptr2index', 'reset', 'torch', 'torch_geometric', 'typing']
```
It looks to me this can happen when two processes in the same node generate the same template file. One process read the python script in the middle, while the other process overwrites it.

This is just my thought. Anyhow, I am getting such error when using with MPI. Any help will be appreciated.

### Versions

```
PyTorch version: 2.0.1+cpu
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.8.10 (default, Nov 22 2023, 10:22:35)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-6.6.12-linuxkit-x86_64-with-glibc2.29
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.4
[pip3] torch==2.0.1+cpu
[pip3] torch-cluster==1.6.3+pt20cpu
[pip3] torch_geometric==2.5.0
[pip3] torch-scatter==2.1.2+pt20cpu
[pip3] torch-sparse==0.6.18+pt20cpu
[pip3] torch-spline-conv==1.2.2+pt20cpu
[pip3] torchaudio==2.0.2+cpu
[pip3] torchvision==0.15.2+cpu
[conda] Could not collect

```


Discussion:
Thanks for the issue. I will take a look :)

Resolving patch:
diff --git a/torch_geometric/nn/conv/message_passing.py b/torch_geometric/nn/conv/message_passing.py
--- a/torch_geometric/nn/conv/message_passing.py
+++ b/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
diff --git a/torch_geometric/template.py b/torch_geometric/template.py
--- a/torch_geometric/template.py
+++ b/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


Test patch: