diff --git a/gigl/distributed/base_sampler.py b/gigl/distributed/base_sampler.py index 986ba5d58..5b284b843 100644 --- a/gigl/distributed/base_sampler.py +++ b/gigl/distributed/base_sampler.py @@ -1,10 +1,14 @@ +import asyncio from collections import defaultdict from dataclasses import dataclass +from threading import Thread from typing import Optional, Union import torch +import uvloop from graphlearn_torch.channel import SampleMessage from graphlearn_torch.distributed import DistNeighborSampler as GLTDistNeighborSampler +from graphlearn_torch.distributed.dist_neighbor_sampler import ensure_device from graphlearn_torch.distributed.dist_feature import DistFeature from graphlearn_torch.distributed.event_loop import wrap_torch_future from graphlearn_torch.sampler import ( @@ -91,6 +95,18 @@ class BaseDistNeighborSampler(GLTDistNeighborSampler): sampling strategy (e.g., k-hop neighbor sampling, PPR-based sampling). """ + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._loop.close() + self._loop = uvloop.new_event_loop() + self._runner_t = Thread(target=self._run_loop) + self._runner_t.daemon = True + self._loop.call_soon_threadsafe(ensure_device, self.device) + + def _run_loop(self): + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + def _prepare_sample_loop_inputs( self, inputs: NodeSamplerInput, diff --git a/pyproject.toml b/pyproject.toml index 34cbe1129..50c4dcae1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dependencies = [ # When updating to a new tensorflow version, please ensure that the version changes align with # `transform` deps listed below. Package versions for `apache-beam`, `tensorflow-transform`, etc "tensorflow~=2.16", + "uvloop", ]