[BugFix][Dataloading] Disable AsyncTransferer copying to anywhere but the GPU (#2323)

* Disable copying for anywhere but the GPU

* Remove unused import and remove references to transferring from the GPU from the docs

* Skip gpu test in cpu mode
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

docs/source/api/python/dgl.dataloading.rst

@@ -53,7 +53,7 @@ Async Copying to/from GPUs
 --------------------------
 .. currentmodule:: dgl.dataloading
 
-Data can be copied from the CPU to the GPU, or from the GPU to the CPU,
+Data can be copied from the CPU to the GPU
 while the GPU is being used for
 computation, using the :class:`AsyncTransferer`.
 For the transfer to be fully asynchronous, the context the

python/dgl/dataloading/async_transferer.py

-""" API for transferring data to/from the GPU over second stream.A """
-
+"""API for transferring data to the GPU over second stream."""
 
 from .. import backend as F
 from .. import ndarray
@@ -37,7 +36,7 @@ class Transfer(object):
 
 
 class AsyncTransferer(object):
-    """ Class for initiating asynchronous copies to/from the GPU on a second
+    """ Class for initiating asynchronous copies to the GPU on a second
     GPU stream.
 
     To initiate a transfer to a GPU:
@@ -73,8 +72,7 @@ class AsyncTransferer(object):
         to be asynchronous, the context the AsyncTranserer is created with must
         be a GPU context, and the input tensor must be in pinned memory.
 
-        Currently, transfers from the GPU to the CPU, and CPU to CPU, will
-        be synchronous.
+        Currently, only transfers to the GPU are supported.
 
         Parameters
         ----------
@@ -94,6 +92,9 @@ class AsyncTransferer(object):
         else:
             ctx = utils.to_dgl_context(device)
 
+        if ctx.device_type != ndarray.DGLContext.STR2MASK["gpu"]:
+            raise ValueError("'device' must be a GPU device.")
+
         tensor = F.zerocopy_to_dgl_ndarray(tensor)
 
         transfer_id = _CAPI_DGLAsyncTransfererStartTransfer(self._handle, tensor, ctx)

tests/compute/test_async_transferer.py

@@ -4,6 +4,8 @@ import backend as F
 
 from dgl.dataloading import AsyncTransferer
 
+@unittest.skipIf(F._default_context_str == 'cpu',
+                 reason="CPU transfer not allowed")
 def test_async_transferer_to_other():
     cpu_ones = F.ones([100,75,25], dtype=F.int32, ctx=F.cpu())
     tran = AsyncTransferer(F.ctx())
@@ -16,12 +18,16 @@ def test_async_transferer_to_other():
 def test_async_transferer_from_other():
     other_ones = F.ones([100,75,25], dtype=F.int32, ctx=F.ctx())
     tran = AsyncTransferer(F.ctx())
-    t = tran.async_copy(other_ones, F.cpu())
-    cpu_ones = t.wait()
-
-    assert F.context(cpu_ones) == F.cpu()
-    assert F.array_equal(F.copy_to(other_ones, ctx=F.cpu()), cpu_ones)
-
+    
+    try:
+        t = tran.async_copy(other_ones, F.cpu())
+    except ValueError:
+        # correctly threw an error
+        pass
+    else:
+        # should have thrown an error
+        assert False
+        
 if __name__ == '__main__':
     test_async_transferer_to_other()
     test_async_transferer_from_other()