From 0a9e5626b01c8d79bfeb5c9d6bd612bb7ed4e514 Mon Sep 17 00:00:00 2001 From: Ngo Bao Date: Sat, 23 May 2026 20:33:41 -0500 Subject: [PATCH 1/3] Fixed bug in examples that failed integration tests --- .gitignore | 3 +++ examples/graphproppred/gcn.py | 2 +- examples/graphproppred/tgcn.py | 2 +- examples/linkproppred/ctan.py | 2 +- examples/linkproppred/roland.py | 2 +- examples/linkproppred/tgb_seq/edgebank.py | 9 ++++++++- test/integration/test_base3.py | 2 +- test/integration/test_edgebank.py | 2 +- tgm/hooks/hook_manager.py | 2 +- 9 files changed, 18 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index d6cb73df..70e365ce 100644 --- a/.gitignore +++ b/.gitignore @@ -166,3 +166,6 @@ submitit* # Benchmarks .benchmarks + +# TGB-seq data cache +data/ diff --git a/examples/graphproppred/gcn.py b/examples/graphproppred/gcn.py index c3f2d357..9a60b481 100644 --- a/examples/graphproppred/gcn.py +++ b/examples/graphproppred/gcn.py @@ -317,7 +317,7 @@ def eval( log_metrics_dict(train_results, epoch=epoch) log_metrics_dict(val_results, epoch=epoch) - val_score = val_results['BinaryAUROC'] + val_score = val_results['ValidationBinaryAUROC'] if val_score > best_val: best_val = val_score test_results = eval(test_loader, test_labels, encoder, decoder, test_metrics) diff --git a/examples/graphproppred/tgcn.py b/examples/graphproppred/tgcn.py index 86d31926..79fd3dbf 100644 --- a/examples/graphproppred/tgcn.py +++ b/examples/graphproppred/tgcn.py @@ -298,7 +298,7 @@ def eval( log_metrics_dict(train_results, epoch=epoch) log_metrics_dict(val_results, epoch=epoch) - val_score = val_results['BinaryAUROC'] + val_score = val_results['ValidationBinaryAUROC'] if val_score > best_val: best_val = val_score test_results, h_0 = eval( diff --git a/examples/linkproppred/ctan.py b/examples/linkproppred/ctan.py index db2d60f3..66f77182 100644 --- a/examples/linkproppred/ctan.py +++ b/examples/linkproppred/ctan.py @@ -234,7 +234,7 @@ def compute_delta_t_stats(train_dg: DGraph) -> Tuple[float, float]: ) train_key, val_key, test_key = hm.keys hm.register_shared(nbr_hook) -hm.register_shared(DeduplicationHook()) +hm.register_shared(DeduplicationHook(seed_nodes_keys=['neg', 'nbr_nids'])) train_loader = DGDataLoader(train_dg, args.bsize, hook_manager=hm) val_loader = DGDataLoader(val_dg, args.bsize, hook_manager=hm) diff --git a/examples/linkproppred/roland.py b/examples/linkproppred/roland.py index a84cb5f6..7e8f187e 100644 --- a/examples/linkproppred/roland.py +++ b/examples/linkproppred/roland.py @@ -204,7 +204,7 @@ def eval( except StopIteration: pass - z = z.detach() + z[0], z[1] = z[0].detach(), z[1].detach() return float(np.mean(perf_list)), z diff --git a/examples/linkproppred/tgb_seq/edgebank.py b/examples/linkproppred/tgb_seq/edgebank.py index c84c3d63..0062e262 100644 --- a/examples/linkproppred/tgb_seq/edgebank.py +++ b/examples/linkproppred/tgb_seq/edgebank.py @@ -1,4 +1,5 @@ import argparse +from typing import Set import numpy as np import torch @@ -64,7 +65,13 @@ def eval( class TGBSEQ_NegativeEdgeSamplerHook(StatelessHook): - produces = {'neg', 'neg_time'} + @property + def produces(self) -> Set[str]: + return {'neg', 'neg_time'} + + @property + def requires(self) -> Set[str]: + return {'edge_src', 'edge_dst', 'edge_time'} def __init__( self, dataset_name: str, split_mode: str, dgraph: DGraph, root: str = './data' diff --git a/test/integration/test_base3.py b/test/integration/test_base3.py index e0e1d02d..15b2be1d 100644 --- a/test/integration/test_base3.py +++ b/test/integration/test_base3.py @@ -8,7 +8,7 @@ '--partition=main', '--cpus-per-task=2', '--mem=4G', - '--time=0:05:00', + '--time=0:10:00', ] ) def test_base3_linkprop_inf_EB_memory(slurm_job_runner, dataset): diff --git a/test/integration/test_edgebank.py b/test/integration/test_edgebank.py index b45c82df..06b8861d 100644 --- a/test/integration/test_edgebank.py +++ b/test/integration/test_edgebank.py @@ -44,7 +44,7 @@ def test_edgebank_linkprop_pred_fixed_memory(slurm_job_runner, dataset): '--partition=main', '--cpus-per-task=2', '--mem=8G', - '--time=0:15:00', + '--time=0:25:00', ] ) def test_edgebank_tgb_seq_unlimited_memory(slurm_job_runner, dataset): diff --git a/tgm/hooks/hook_manager.py b/tgm/hooks/hook_manager.py index fd94a33b..9c5fc43f 100644 --- a/tgm/hooks/hook_manager.py +++ b/tgm/hooks/hook_manager.py @@ -210,7 +210,7 @@ def activate(self, key: str) -> Iterator[None]: def _ensure_valid_hook(self, hook: Any) -> None: if not isinstance(hook, DGHook): raise BadHookProtocolError( - f'Cannot register hook {type(hook).__name__}: must implement __call__(dg: DGraph, batch: DGBatch) -> DGBatch and reset_state()' + f'Cannot register hook {type(hook).__name__}: must implement __call__(dg: DGraph, batch: DGBatch) -> DGBatch, reset_state(), requires and produces properties.' ) def _ensure_no_active_key(self) -> None: From 0d2bb34fd8d7a5df7370a5fdc860bbdd65736bea Mon Sep 17 00:00:00 2001 From: Ngo Bao Date: Sun, 24 May 2026 13:38:02 -0500 Subject: [PATCH 2/3] Increased time budget for base3 and edgebank integration tests --- test/integration/test_base3.py | 2 +- test/integration/test_edgebank.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/test/integration/test_base3.py b/test/integration/test_base3.py index 15b2be1d..6e3f0b7a 100644 --- a/test/integration/test_base3.py +++ b/test/integration/test_base3.py @@ -8,7 +8,7 @@ '--partition=main', '--cpus-per-task=2', '--mem=4G', - '--time=0:10:00', + '--time=0:20:00', ] ) def test_base3_linkprop_inf_EB_memory(slurm_job_runner, dataset): diff --git a/test/integration/test_edgebank.py b/test/integration/test_edgebank.py index 06b8861d..39dff96a 100644 --- a/test/integration/test_edgebank.py +++ b/test/integration/test_edgebank.py @@ -116,7 +116,7 @@ def test_edgebank_linkprop_pred_unlimited_memory_tkgl(slurm_job_runner, dataset) '--partition=main', '--cpus-per-task=2', '--mem=8G', - '--time=2:00:00', + '--time=3:00:00', ] ) def test_edgebank_linkprop_pred_fixed_memory_tkgl(slurm_job_runner, dataset): From c08f117d892f039ea6f56840088c268bc5b67cb2 Mon Sep 17 00:00:00 2001 From: Ngo Bao Date: Sun, 24 May 2026 19:32:23 -0500 Subject: [PATCH 3/3] Increase time budget for tkgl edgebank --- test/integration/test_edgebank.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/integration/test_edgebank.py b/test/integration/test_edgebank.py index 39dff96a..fd9f22fd 100644 --- a/test/integration/test_edgebank.py +++ b/test/integration/test_edgebank.py @@ -98,7 +98,7 @@ def test_edgebank_linkprop_pred_fixed_memory_thgl(slurm_job_runner, dataset): '--partition=main', '--cpus-per-task=2', '--mem=8G', - '--time=1:15:00', + '--time=3:00:00', ] ) def test_edgebank_linkprop_pred_unlimited_memory_tkgl(slurm_job_runner, dataset): @@ -116,7 +116,7 @@ def test_edgebank_linkprop_pred_unlimited_memory_tkgl(slurm_job_runner, dataset) '--partition=main', '--cpus-per-task=2', '--mem=8G', - '--time=3:00:00', + '--time=4:00:00', ] ) def test_edgebank_linkprop_pred_fixed_memory_tkgl(slurm_job_runner, dataset):