Skip to content

Commit 33e3ded

Browse files
committed
adding the torch guard and test changes
1 parent 1eb0058 commit 33e3ded

13 files changed

Lines changed: 155 additions & 46 deletions

README.md

Lines changed: 24 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -88,12 +88,21 @@ You can install TGB via [pip](https://pypi.org/project/py-tgb/). **Requires pyth
8888
pip install py-tgb
8989
```
9090

91-
TGB uses [PyTorch](https://pytorch.org/) and [PyG](https://pyg.org/) at runtime, but **PyTorch is
92-
not installed automatically** — install the build that matches your platform/CUDA first (see the
93-
[PyTorch install guide](https://pytorch.org/get-started/locally/)). For example, a CPU build:
91+
By default this is a **pure Python/numpy install — PyTorch and PyG are not required** and are not
92+
pulled in. This is enough for dataset loading (`LinkPropPredDataset`/`NodePropPredDataset`) and
93+
evaluation (`Evaluator`, `NegativeEdgeSampler`) with plain numpy arrays.
94+
95+
If you want to use the PyG dataset wrappers (`PyGLinkPropPredDataset`/`PyGNodePropPredDataset`) or
96+
run the example model scripts, install the `torch` extra, which pulls in `torch-geometric`:
97+
```
98+
pip install py-tgb[torch]
99+
```
100+
`torch-geometric` doesn't pin a specific PyTorch build, so install the [PyTorch build matching your
101+
platform/CUDA](https://pytorch.org/get-started/locally/) first if you need a non-default one, e.g.:
94102
```
95103
pip install torch # CPU (macOS / no CUDA)
96104
# or a CUDA build, e.g.: pip install torch --index-url https://download.pytorch.org/whl/cu121
105+
pip install py-tgb[torch]
97106
```
98107

99108
### Development Install (uv, optional)
@@ -106,13 +115,13 @@ built or how `pip install py-tgb` behaves for users.**
106115
# install uv: https://docs.astral.sh/uv/getting-started/installation/
107116
uv venv --python 3.11 && source .venv/bin/activate
108117
109-
uv pip install -e . # builds via poetry-core, installs runtime deps
118+
uv pip install -e . # builds via poetry-core, installs runtime deps (pure numpy, no torch)
110119
uv pip install -r requirements-dev.txt # pytest, mkdocs, ...
111120
112-
# PyTorch is not pulled in automatically — install the build for your platform.
121+
# Only needed for the PyG dataset wrappers / example scripts:
113122
uv pip install torch # CPU (macOS / no CUDA)
114123
# CUDA example: uv pip install torch==2.0.0 --index-url https://download.pytorch.org/whl/cu117
115-
uv pip install torch_geometric
124+
uv pip install -e ".[torch]" # installs torch-geometric
116125
```
117126

118127
### Testing
@@ -122,18 +131,20 @@ TGB ships an offline-by-default `pytest` suite that guards dataset metadata (URL
122131

123132
```
124133
pip install -r requirements-dev.txt # pytest, pytest-mock
125-
pip install -e . # runtime deps (torch-geometric, pandas, ...)
126-
pip install torch # torch is needed by most tests but is not
127-
# a declared package dependency; install the
128-
# build matching your platform/CUDA
134+
pip install -e . # runtime deps (pure numpy, no torch required)
129135
130136
pytest -m "not network" # fast, offline suite (default for CI)
131137
pytest --run-network # also check that each dataset URL is reachable
132138
```
133139

134-
Tests that import `tgb` model/utility code are skipped automatically if `torch` is not installed;
135-
the `tgb/utils/info.py` metadata and URL tests run with no heavy dependencies. See
136-
[`test/README.md`](test/README.md) for the full layout and the network-test opt-in.
140+
Nearly all tests run with no torch installed at all — `tgb/utils/utils.py`,
141+
`Evaluator`, `NegativeEdgeSampler`, and the dataset classes all guard their torch imports and fall
142+
back to plain numpy behavior. `test/test_no_torch.py` simulates torch being uninstalled and asserts
143+
this directly, including that the PyG-only wrapper modules (which do require `pip install
144+
py-tgb[torch]`) still fail with a clear `ImportError` rather than silently misbehaving. Only a
145+
handful of tests that construct actual tensors (e.g. `test_eval_torch_matches_numpy`) skip when
146+
torch isn't installed. See [`test/README.md`](test/README.md) for the full layout and the
147+
network-test opt-in.
137148

138149
### Links and Datasets
139150

pyproject.toml

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,14 +8,17 @@ packages = [{include = "tgb"}]
88

99
[tool.poetry.dependencies]
1010
python = "^3.9"
11-
torch-geometric = "^2.3.0"
11+
torch-geometric = {version = "^2.3.0", optional = true}
1212
tqdm = "^4.65.0"
1313
numpy = "^2.0.2"
1414
clint = "^0.5.1"
1515
requests = "^2.28.2"
1616
pandas = ">=2.2.3"
1717
scikit-learn = "^1.2.2"
1818

19+
[tool.poetry.extras]
20+
torch = ["torch-geometric"]
21+
1922
[tool.poetry.group.dev.dependencies]
2023
mkdocs = "^1.4.3"
2124
mkdocs-material = "^9.1.15"

test/README.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ TGB_RUN_NETWORK=1 pytest # same, via env var
2525
| `test_utils.py` | helpers in `tgb/utils/utils.py` | no |
2626
| `test_pre_process.py` | CSV loaders in `tgb/utils/pre_process.py` | no |
2727
| `test_negative_sampler.py` | `NegativeEdgeSampler` load/query behavior | no |
28+
| `test_no_torch.py` | simulates torch/torch_geometric being uninstalled; guards issue #127 | no |
2829

2930
All dataset lists are derived from the dictionaries in `info.py`, so new datasets
3031
are covered automatically. Shared fixtures and the `--run-network` gate live in

test/test_download.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,6 @@
1616

1717
import pytest
1818

19-
# The dataset classes import torch transitively (tgb.utils.utils); skip this
20-
# module cleanly if torch is not installed rather than failing collection.
21-
pytest.importorskip("torch")
22-
2319
from tgb.linkproppred.dataset import LinkPropPredDataset
2420
from tgb.nodeproppred.dataset import NodePropPredDataset
2521

test/test_evaluate.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -6,9 +6,6 @@
66
import numpy as np
77
import pytest
88

9-
# Evaluator imports torch transitively (tgb.utils.utils); skip cleanly if absent.
10-
pytest.importorskip("torch")
11-
129
from tgb.linkproppred.evaluate import Evaluator
1310

1411
# For pos score 1.0 against neg scores [0.5, 0.9, 2.0]:

test/test_negative_sampler.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,6 @@
88
import numpy as np
99
import pytest
1010

11-
# negative_sampler / utils import torch at module top; skip cleanly if absent.
12-
pytest.importorskip("torch")
13-
1411
from tgb.linkproppred.negative_sampler import NegativeEdgeSampler
1512
from tgb.utils.utils import save_pkl
1613

test/test_no_torch.py

Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
"""
2+
Regression test for https://github.com/shenyangHuang/TGB/issues/127.
3+
4+
Simulates torch/torch_geometric being uninstalled and verifies that the
5+
numpy-only import surface (dataset classes, ``Evaluator``, negative samplers,
6+
``tgb.utils.utils``) still imports and works. Also locks in the intentional
7+
boundary: the PyG-only wrapper modules (``dataset_pyg``) still require torch
8+
and should fail with a clear ``ImportError`` rather than something more
9+
confusing, since guarding them would strip their only purpose.
10+
"""
11+
import builtins
12+
import importlib
13+
import sys
14+
15+
import numpy as np
16+
import pytest
17+
18+
TGB_MODULES = [
19+
"tgb.utils.utils",
20+
"tgb.utils.pre_process",
21+
"tgb.linkproppred.dataset",
22+
"tgb.linkproppred.evaluate",
23+
"tgb.linkproppred.negative_sampler",
24+
"tgb.linkproppred.thg_negative_sampler",
25+
"tgb.linkproppred.tkg_negative_sampler",
26+
"tgb.nodeproppred.dataset",
27+
"tgb.nodeproppred.evaluate",
28+
]
29+
30+
31+
@pytest.fixture
32+
def no_torch(monkeypatch):
33+
"""Simulate torch/torch_geometric being uninstalled and force TGB
34+
modules to re-import under that condition."""
35+
saved_modules = sys.modules.copy()
36+
for name in list(sys.modules):
37+
if name == "torch" or name.startswith(("torch.", "torch_geometric")) \
38+
or name.startswith("tgb."):
39+
sys.modules.pop(name, None)
40+
41+
real_import = builtins.__import__
42+
43+
def fake_import(name, *args, **kwargs):
44+
if name == "torch" or name.startswith(("torch.", "torch_geometric")):
45+
raise ImportError(f"simulated: {name} is not installed")
46+
return real_import(name, *args, **kwargs)
47+
48+
monkeypatch.setattr(builtins, "__import__", fake_import)
49+
yield
50+
sys.modules.clear()
51+
sys.modules.update(saved_modules)
52+
53+
54+
def test_core_modules_import_without_torch(no_torch):
55+
for name in TGB_MODULES:
56+
mod = importlib.import_module(name)
57+
assert getattr(mod, "torch", None) is None
58+
59+
60+
def test_negative_sampler_works_without_torch(no_torch):
61+
negative_sampler = importlib.import_module("tgb.linkproppred.negative_sampler")
62+
sampler = negative_sampler.NegativeEdgeSampler(dataset_name="tgbl-mock")
63+
sampler.eval_set["test"] = {(0, 1, 0): [5, 6, 7]}
64+
neg = sampler.query_batch(
65+
np.array([0]), np.array([1]), np.array([0]), split_mode="test"
66+
)
67+
assert neg == [[5, 6, 7]]
68+
69+
70+
def test_evaluator_works_without_torch(no_torch):
71+
evaluate = importlib.import_module("tgb.linkproppred.evaluate")
72+
evaluator = evaluate.Evaluator(name="tgbl-wiki", k_value=10)
73+
result = evaluator.eval(
74+
{
75+
"y_pred_pos": np.array([1.0]),
76+
"y_pred_neg": np.array([[0.5, 0.9, 2.0]]),
77+
"eval_metric": ["mrr"],
78+
}
79+
)
80+
assert result["mrr"] == pytest.approx(0.5)
81+
82+
83+
def test_set_random_seed_works_without_torch(no_torch):
84+
utils = importlib.import_module("tgb.utils.utils")
85+
utils.set_random_seed(0) # should not raise even though torch is None
86+
87+
88+
def test_pyg_dataset_still_requires_torch(no_torch):
89+
with pytest.raises(ImportError):
90+
importlib.import_module("tgb.linkproppred.dataset_pyg")

test/test_pre_process.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,6 @@
77
import numpy as np
88
import pytest
99

10-
# pre_process imports torch transitively (tgb.utils.utils); skip cleanly if absent.
11-
pytest.importorskip("torch")
12-
1310
from tgb.utils.pre_process import load_edgelist_trade, load_edgelist_wiki
1411

1512

test/test_utils.py

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,6 @@
55
import pandas as pd
66
import pytest
77

8-
# tgb.utils.utils imports torch at module top; skip cleanly if torch is absent.
9-
pytest.importorskip("torch")
10-
118
from tgb.utils.utils import (
129
add_inverse_quadruples,
1310
find_nearest,

tgb/linkproppred/negative_sampler.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,14 @@
33
Load already generated negative edges from file, batch them based on the positive edge, and return the evaluation set
44
"""
55

6-
import torch
7-
from torch import Tensor
86
import numpy as np
7+
8+
try:
9+
import torch
10+
from torch import Tensor
11+
except ImportError:
12+
torch = None
13+
Tensor = None
914
from tgb.utils.utils import save_pkl, load_pkl
1015
from tgb.utils.info import PROJ_DIR
1116
import os

0 commit comments

Comments
 (0)