Make hooks seedable and added instance-level properties (#396)

* Made requires and produces to instance level variables

* WIP

* Added id to each hook

* Updated with add attribute to hook

* Updated BaseDGHook

* Applied ruff format

* Applied add_attribute_to_batch to add produce attribute to the batch

* typo

* Rename add attribute to batch function and added comments to SeedableHook

Author: ntgbaoo

examples/linkproppred/tgn.py

@@ -205,7 +205,7 @@ def eval(
 )
 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)

test/unit/test_hooks/test_batch_analytics_hook.py

@@ -50,6 +50,29 @@ def test_hook_dependancies():
         'num_repeated_node_events',
     }
 
+    hook_with_id = BatchAnalyticsHook(id='foo')
+    assert hook_with_id.requires == {
+        'edge_src',
+        'edge_dst',
+        'edge_time',
+        'node_x_time',
+        'node_x_nids',
+    }
+    assert hook_with_id.produces == {
+        'num_edge_events_foo',
+        'num_node_events_foo',
+        'num_unique_timestamps_foo',
+        'num_unique_nodes_foo',
+        'avg_degree_foo',
+        'num_repeated_edge_events_foo',
+        'num_repeated_node_events_foo',
+    }
+
+
+def test_hook_repre():
+    hook_with_id = BatchAnalyticsHook(id='foo')
+    assert 'foo' in hook_with_id.__repr__()
+
 
 def test_hook_reset_state():
     assert BatchAnalyticsHook.has_state is False
@@ -60,8 +83,6 @@ def test_basic_analytics_num_events_and_timestamps(dg):
     batch = dg.materialize()
     processed_batch = hook(dg, batch)
 
-    # assert batch.node_x_nids is not None
-
     # edge and node events
     assert processed_batch.num_edge_events == 3
     assert processed_batch.num_node_events == 3
@@ -121,3 +142,24 @@ def test_basic_analytics_empty_edges_and_nodes(dg):
 
     # _count_repeated_node_events: node_x_nids.numel() == 0 -> 0
     assert processed_batch.num_repeated_node_events == 0
+
+
+def test_hook_with_id(dg):
+    hook = BatchAnalyticsHook(id='foo')
+    batch = dg.materialize()
+
+    batch = dg.materialize()
+    processed_batch = hook(dg, batch)
+
+    expected_produce = [
+        'num_edge_events_foo',
+        'num_node_events_foo',
+        'num_unique_timestamps_foo',
+        'num_unique_nodes_foo',
+        'avg_degree_foo',
+        'num_repeated_edge_events_foo',
+        'num_repeated_node_events_foo',
+    ]
+
+    for produce in expected_produce:
+        assert hasattr(processed_batch, produce)

test/unit/test_hooks/test_deduplication_hook.py

@@ -16,8 +16,18 @@ def dg():
 
 
 def test_hook_dependancies():
-    assert DeduplicationHook.requires == {'edge_src', 'edge_dst'}
-    assert DeduplicationHook.produces == {'unique_nids', 'global_to_local'}
+    hook = DeduplicationHook()
+    assert hook.requires == {'edge_src', 'edge_dst'}
+    assert hook.produces == {'unique_nids', 'global_to_local'}
+
+    hook_with_id = DeduplicationHook(id='foo')
+    assert hook_with_id.requires == {'edge_src', 'edge_dst'}
+    assert hook_with_id.produces == {'unique_nids_foo', 'global_to_local_foo'}
+
+
+def test_hook_repre():
+    hook_with_id = DeduplicationHook(id='foo')
+    assert 'foo' in hook_with_id.__repr__()
 
 
 def test_hook_reset_state():
@@ -39,8 +49,23 @@ def test_dedup(dg):
     )
 
 
+def test_dedup_with_id(dg):
+    hook = DeduplicationHook(id='foo')
+    batch = dg.materialize()
+    processed_batch = hook(dg, batch)
+    torch.testing.assert_close(
+        processed_batch.unique_nids_foo, torch.IntTensor([1, 2, 4, 8])
+    )
+    torch.testing.assert_close(
+        processed_batch.global_to_local_foo(batch.edge_src), torch.IntTensor([1, 1, 0])
+    )
+    torch.testing.assert_close(
+        processed_batch.global_to_local_foo(batch.edge_dst), torch.IntTensor([1, 2, 3])
+    )
+
+
 def test_dedup_with_negatives(dg):
-    hook = DeduplicationHook()
+    hook = DeduplicationHook(seed_nodes_keys=['neg'])
     batch = dg.materialize()
     batch.neg = torch.IntTensor([1, 5, 10])  # add some mock negatives
 
@@ -60,7 +85,7 @@ def test_dedup_with_negatives(dg):
 
 
 def test_dedup_with_nbrs(dg):
-    hook = DeduplicationHook()
+    hook = DeduplicationHook(seed_nodes_keys=['nbr_nids'])
     batch = dg.materialize()
     batch.nbr_nids = [  # add some mock neighbours
         torch.IntTensor([1, 5]),  # First hop
@@ -107,7 +132,7 @@ def node_only_graph():
 
 def test_dedup_node_only_batch(node_only_graph):
     hm = HookManager(keys=['unit'])
-    hm.register('unit', DeduplicationHook())
+    hm.register('unit', DeduplicationHook(seed_nodes_keys=['node_x_nids']))
     loader = DGDataLoader(node_only_graph, batch_size=3, hook_manager=hm)
     with hm.activate('unit'):
         batch_iter = iter(loader)

test/unit/test_hooks/test_device_transfer_hook.py

@@ -15,8 +15,9 @@ def dg():
 
 
 def test_hook_dependancies():
-    assert DeviceTransferHook.requires == set()
-    assert DeviceTransferHook.produces == set()
+    hook = DeviceTransferHook('cpu')
+    assert hook.requires == set()
+    assert hook.produces == set()
 
 
 def test_hook_reset_state():

test/unit/test_hooks/test_hook_manager.py

@@ -11,29 +11,47 @@
 
 
 class MockHook(StatelessHook):
-    produces = {'foo'}
+    _cls_produces = {'foo'}
+
+    def __init__(self, id: str = None):
+        super().__init__()
+        self._id = id
+        self.__post_init__()
 
     def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
         batch.edge_time *= 2
         return batch
 
 
 class MockHookRequires(StatelessHook):
-    requires = {'foo'}
+    _cls_requires = {'foo'}
+
+    def __init__(self, id: str = None):
+        super().__init__()
+        self._id = id
+        self.__post_init__()
 
     def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
         return batch
 
 
-class DeduplicationMockHook(StatelessHook):
+class MockHookRequiresWoof(StatelessHook):
+    _cls_requires = {'foo_woof'}
+
+    def __init__(self, id: str = None):
+        super().__init__()
+        self._id = id
+        self.__post_init__()
+
     def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
         return batch
 
 
 class MockHookWithState(StatefulHook):
-    has_state: bool = True
-
-    def __init__(self) -> None:
+    def __init__(self, id: str = None) -> None:
+        super().__init__()
+        self._id = id
+        self.has_state = True
         self.x = 0
 
     def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
@@ -165,8 +183,8 @@ def test_resolve_hooks_by_key():
 def test_resolve_hooks_no_solution_no_dag():
     h1 = MockHook()
     h2 = MockHook()
-    h1.requires, h1.produces = {'x'}, {'y'}
-    h2.requires, h2.produces = {'y'}, {'x'}
+    h1._requires, h1._produces = {'x'}, {'y'}
+    h2._requires, h2._produces = {'y'}, {'x'}
 
     # Cycle-like missing dependency
     hm = HookManager(keys=['train'])
@@ -210,8 +228,8 @@ def test_topo_sort_cached(dg, monkeypatch):
     hm = HookManager(keys=['train'])
 
     h1, h2 = MockHook(), MockHook()
-    h1.requires, h1.produces = set(), {'x'}
-    h2.requires, h2.produces = {'x'}, {'y'}
+    h1._requires, h1._produces = set(), {'x'}
+    h2._requires, h2._produces = {'x'}, {'y'}
 
     hm.register('train', h1)
     hm.register('train', h2)
@@ -237,7 +255,7 @@ def test_topo_sort_cached_invalidated(dg, monkeypatch):
     hm = HookManager(keys=['train'])
 
     h1 = MockHook()
-    h1.requires, h1.produces = set(), {'x'}
+    h1._requires, h1._produces = set(), {'x'}
 
     hm.register('train', h1)
     call_count = {'n': 0}
@@ -266,8 +284,8 @@ def fake_topo_sort(hooks_list):
 def test_topo_sort_no_solution_no_dag(dg):
     h1 = MockHook()
     h2 = MockHook()
-    h1.requires, h1.produces = {'x'}, {'y'}
-    h2.requires, h2.produces = {'y'}, {'x'}
+    h1._requires, h1._produces = {'x'}, {'y'}
+    h2._requires, h2._produces = {'y'}, {'x'}
 
     # Cycle-like missing dependency
     hm = HookManager(keys=['train'])
@@ -362,8 +380,8 @@ def test_activate_ctx():
 
 def test_topo_sort_neg_before_nbr():
     mock_neg_hook, mock_nbr_hook = MockHook(), MockHook()
-    mock_neg_hook.requires, mock_neg_hook.produces = set(), {'neg'}
-    mock_nbr_hook.requires, mock_nbr_hook.produces = set(), {'nbr_nids'}
+    mock_neg_hook._requires, mock_neg_hook._produces = set(), {'neg'}
+    mock_nbr_hook._requires, mock_nbr_hook._produces = set(), {'nbr_nids'}
 
     # Register neg first in foo, nbr first in bar
     hm = HookManager(keys=['foo', 'bar'])
@@ -381,25 +399,14 @@ def test_topo_sort_neg_before_nbr():
     assert bar_hooks.index(mock_neg_hook) < bar_hooks.index(mock_nbr_hook)
 
 
-def test_force_last_dedup_hook():
-    # @TODO: Dedup hook is temporarily forced to run at the end.
-    # This test potentially needs to be updated once instance-level require is introduced to DGHook.
-    h1 = MockHook()
-    h2 = MockHookRequires()
-    h3 = DeduplicationMockHook()
+def test_resolve_hooks_with_id_by_key():
+    h1 = MockHook(id='woof')  # MockHook produces with _woof suffix
+    h2 = MockHookRequiresWoof()
 
-    hm = HookManager(keys=['train', 'val'])
-    hm.register('train', h3)
+    hm = HookManager(keys=['train'])
     hm.register('train', h2)
     hm.register('train', h1)
-    hm.register('val', h3)
-    hm.register('val', h2)
-    hm.register('val', h1)
 
-    hm.resolve_hooks()
-    assert len(hm._key_to_hooks['train']) == 3
-    assert len(hm._key_to_hooks['val']) == 3
+    hm.resolve_hooks('train')
+    assert len(hm._key_to_hooks['train']) == 2
     assert hm._key_to_hooks['train'].index(h1) < hm._key_to_hooks['train'].index(h2)
-    assert hm._key_to_hooks['val'].index(h1) < hm._key_to_hooks['val'].index(h2)
-    assert hm._key_to_hooks['train'].index(h2) < hm._key_to_hooks['train'].index(h3)
-    assert hm._key_to_hooks['val'].index(h2) < hm._key_to_hooks['val'].index(h3)

test/unit/test_hooks/test_negative_edge_sampler_hook.py

@@ -14,8 +14,18 @@ def data():
 
 
 def test_hook_dependancies():
-    assert NegativeEdgeSamplerHook.requires == {'edge_src', 'edge_dst', 'edge_time'}
-    assert NegativeEdgeSamplerHook.produces == {'neg', 'neg_time'}
+    hook = NegativeEdgeSamplerHook(low=0, high=10)
+    assert hook.requires == {'edge_src', 'edge_dst', 'edge_time'}
+    assert hook.produces == {'neg', 'neg_time'}
+
+    hook_with_id = NegativeEdgeSamplerHook(low=0, high=10, id='foo')
+    assert hook_with_id.requires == {'edge_src', 'edge_dst', 'edge_time'}
+    assert hook_with_id.produces == {'neg_foo', 'neg_time_foo'}
+
+
+def test_hook_repre():
+    hook_with_id = NegativeEdgeSamplerHook(low=0, high=10, id='foo')
+    assert 'foo' in hook_with_id.__repr__()
 
 
 def test_hook_reset_state():
@@ -42,6 +52,17 @@ def test_negative_edge_sampler(data):
     assert batch.neg_time.shape == batch.neg.shape
 
 
+def test_negative_edge_sampler_with_id(data):
+    dg = DGraph(data)
+    hook = NegativeEdgeSamplerHook(low=0, high=10, id='foo')
+    batch = hook(dg, dg.materialize())
+    assert isinstance(batch, DGBatch)
+    assert torch.is_tensor(batch.neg_foo)
+    assert torch.is_tensor(batch.neg_time_foo)
+    assert batch.neg_foo.shape == batch.edge_dst.shape
+    assert batch.neg_time_foo.shape == batch.neg_foo.shape
+
+
 @pytest.fixture
 def node_only_data():
     edge_index = torch.IntTensor([[1, 2], [2, 3], [3, 4]])