Sanity checks for GPU tests, 3rd party framework tensor conversion ch…

…anges (#646)

* Add sanity checks to handle cases of hidden GPU devices (`CUDA_VISIBLE_DEVICES=-1`)
Add `has_tensorflow_gpu`

* Allow `...2xp` utility methods to accept a target `Ops` object for conversions
Refactor 3rd party framework GPU tensor detection

* Handle `cupy.fromDlpack` deprecation for `cupy >= 10.0.0`

* Make `ops` arg in `...2xp` functions keyword-only

* Short-circuit `..._gpu_array` functions
Add `Ops` type to `..2xp` functions
Defer `cupy` deprecation-related changes to a different PR

* Fix type annotation

* Defer sanity check in `_custom_kernels.py` to another PR

thinc/backends/cupy_ops.py

@@ -20,7 +20,8 @@
 from . import _custom_kernels
 from ..types import DeviceTypes
 from ..util import torch2xp, tensorflow2xp, mxnet2xp
-from ..util import is_torch_array, is_tensorflow_array, is_mxnet_array
+from ..util import is_cupy_array
+from ..util import is_torch_gpu_array, is_tensorflow_gpu_array, is_mxnet_gpu_array
 
 
 @registry.ops("CupyOps")
@@ -79,15 +80,15 @@ def asarray(self, data, dtype=None):
  # We'll try to perform a zero-copy conversion if possible.
  array = None
  cast_array = False
- if isinstance(data, cupy.ndarray):
+ if is_cupy_array(data):
  array = self.xp.asarray(data, **dtype)
- elif is_torch_array(data) and data.device.type == "cuda":
+ elif is_torch_gpu_array(data):
  array = torch2xp(data)
  cast_array = True
- elif is_tensorflow_array(data) and "GPU:" in data.device:
+ elif is_tensorflow_gpu_array(data):
  array = tensorflow2xp(data)
  cast_array = True
- elif is_mxnet_array(data) and data.context.device_type != "cpu":
+ elif is_mxnet_gpu_array(data):
  array = mxnet2xp(data)
  cast_array = True
  else:

thinc/tests/backends/test_ops.py

@@ -8,7 +8,7 @@
 from packaging.version import Version
 from thinc.api import NumpyOps, CupyOps, Ops, get_ops
 from thinc.api import get_current_ops, use_ops
-from thinc.util import has_torch, torch2xp, xp2torch, torch_version
+from thinc.util import has_torch, torch2xp, xp2torch, torch_version, gpu_is_available
 from thinc.api import fix_random_seed
 from thinc.api import LSTM
 from thinc.types import Floats2d
@@ -25,7 +25,7 @@
 BLIS_OPS = NumpyOps(use_blis=True)
 CPU_OPS = [NUMPY_OPS, VANILLA_OPS]
 XP_OPS = [NUMPY_OPS]
-if CupyOps.xp is not None:
+if CupyOps.xp is not None and gpu_is_available():
  XP_OPS.append(CupyOps())
 ALL_OPS = XP_OPS + [VANILLA_OPS]
 
@@ -570,7 +570,9 @@ def test_backprop_seq2col_window_two(ops, dtype):
  ops.xp.testing.assert_allclose(seq, expected, atol=0.001, rtol=0.001)
 
 
-@pytest.mark.skipif(CupyOps.xp is None, reason="needs GPU/CuPy")
+@pytest.mark.skipif(
+ CupyOps.xp is None or not gpu_is_available(), reason="needs GPU/CuPy"
+)
 @pytest.mark.parametrize("nW", [1, 2])
 def test_large_seq2col_gpu_against_cpu(nW):
  cupy_ops = CupyOps()
@@ -592,7 +594,9 @@ def test_large_seq2col_gpu_against_cpu(nW):
  assert_allclose(cols, cols_gpu.get())
 
 
-@pytest.mark.skipif(CupyOps.xp is None, reason="needs GPU/CuPy")
+@pytest.mark.skipif(
+ CupyOps.xp is None or not gpu_is_available(), reason="needs GPU/CuPy"
+)
 @pytest.mark.parametrize("nW", [1, 2])
 def test_large_backprop_seq2col_gpu_against_cpu(nW):
  cupy_ops = CupyOps()

thinc/tests/layers/test_tensorflow_wrapper.py

@@ -2,7 +2,7 @@
 import pytest
 from thinc.api import Adam, ArgsKwargs, Linear, Model, TensorFlowWrapper
 from thinc.api import get_current_ops, keras_subclass, tensorflow2xp, xp2tensorflow
-from thinc.util import has_cupy, has_tensorflow, to_categorical
+from thinc.util import gpu_is_available, has_tensorflow, to_categorical
 
 from ..util import check_input_converters, make_tempdir
 
@@ -342,7 +342,7 @@ def test_tensorflow_wrapper_to_cpu(tf_model):
 
 
 @pytest.mark.skipif(not has_tensorflow, reason="needs TensorFlow")
-@pytest.mark.skipif(not has_cupy, reason="needs cupy")
+@pytest.mark.skipif(not gpu_is_available(), reason="needs GPU/cupy")
 def test_tensorflow_wrapper_to_gpu(model, X):
  model.to_gpu(0)
 

thinc/types.py

@@ -9,7 +9,7 @@
  import cupy
 
  get_array_module = cupy.get_array_module
-except ImportError:
+except (ImportError, AttributeError):
  get_array_module = lambda obj: numpy
 
 # Use typing_extensions for Python versions < 3.8
@@ -684,7 +684,7 @@ def __setitem__(self, key: _4_Key2d, value: Floats2d) -> None: ...
  def __setitem__(self, key: _4_Key3d, value: Floats3d) -> None: ...
  @overload
  def __setitem__(self, key: _4_Key4d, value: "Floats4d") -> None: ...
- 
+
  def __setitem__(self, key: _4_AllKeys, value: _F4_AllReturns) -> None: ...
 
  @overload
@@ -792,7 +792,7 @@ def copy(self):
  self.data.copy(),
  self.size_at_t.copy(),
  self.lengths.copy(),
- self.indices.copy()
+ self.indices.copy(),
  )
 
  def __len__(self) -> int:
@@ -923,7 +923,7 @@ def __len__(self) -> int:
 class ArgsKwargs:
  """A tuple of (args, kwargs) that can be spread into some function f:
 
-  f(*args, **kwargs)
+ f(*args, **kwargs)
  """
 
  args: Tuple[Any, ...]

thinc/util.py

@@ -48,8 +48,10 @@
  import tensorflow as tf
 
  has_tensorflow = True
+ has_tensorflow_gpu = len(tf.config.get_visible_devices("GPU")) > 0
 except ImportError: # pragma: no cover
  has_tensorflow = False
+ has_tensorflow_gpu = False
 
 
 try: # pragma: no cover
@@ -61,6 +63,10 @@
 
 from .types import ArrayXd, ArgsKwargs, Ragged, Padded, FloatsXd, IntsXd # noqa: E402
 from . import types # noqa: E402
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+ from .api import Ops
 
 
 def get_array_module(arr): # pragma: no cover
@@ -71,6 +77,9 @@ def get_array_module(arr): # pragma: no cover
 
 
 def gpu_is_available():
+ if not has_cupy:
+ return False
+
  try:
  cupy.cuda.runtime.getDeviceCount()
  return True
@@ -98,7 +107,7 @@ def is_xp_array(obj: Any) -> bool:
 
 
 def is_cupy_array(obj: Any) -> bool: # pragma: no cover
- """Check whether an object is a cupy array"""
+ """Check whether an object is a cupy array."""
  if not has_cupy:
  return False
  elif isinstance(obj, cupy.ndarray):
@@ -108,7 +117,7 @@ def is_cupy_array(obj: Any) -> bool: # pragma: no cover
 
 
 def is_numpy_array(obj: Any) -> bool:
- """Check whether an object is a numpy array"""
+ """Check whether an object is a numpy array."""
  if isinstance(obj, numpy.ndarray):
  return True
  else:
@@ -124,6 +133,10 @@ def is_torch_array(obj: Any) -> bool: # pragma: no cover
  return False
 
 
+def is_torch_gpu_array(obj: Any) -> bool: # pragma: no cover
+ return is_torch_array(obj) and obj.is_cuda
+
+
 def is_tensorflow_array(obj: Any) -> bool: # pragma: no cover
  if not has_tensorflow:
  return False
@@ -133,6 +146,10 @@ def is_tensorflow_array(obj: Any) -> bool: # pragma: no cover
  return False
 
 
+def is_tensorflow_gpu_array(obj: Any) -> bool: # pragma: no cover
+ return is_tensorflow_array(obj) and "GPU:" in obj.device
+
+
 def is_mxnet_array(obj: Any) -> bool: # pragma: no cover
  if not has_mxnet:
  return False
@@ -142,6 +159,10 @@ def is_mxnet_array(obj: Any) -> bool: # pragma: no cover
  return False
 
 
+def is_mxnet_gpu_array(obj: Any) -> bool: # pragma: no cover
+ return is_mxnet_array(obj) and obj.context.device_type != "cpu"
+
+
 def to_numpy(data): # pragma: no cover
  if isinstance(data, numpy.ndarray):
  return data
@@ -341,6 +362,7 @@ def xp2torch(
  xp_tensor: ArrayXd, requires_grad: bool = False
 ) -> "torch.Tensor": # pragma: no cover
  """Convert a numpy or cupy tensor to a PyTorch tensor."""
+ assert_pytorch_installed()
  if hasattr(xp_tensor, "toDlpack"):
  dlpack_tensor = xp_tensor.toDlpack() # type: ignore
  torch_tensor = torch.utils.dlpack.from_dlpack(dlpack_tensor)
@@ -351,12 +373,25 @@ def xp2torch(
  return torch_tensor
 
 
-def torch2xp(torch_tensor: "torch.Tensor") -> ArrayXd: # pragma: no cover
- """Convert a torch tensor to a numpy or cupy tensor."""
- if torch_tensor.is_cuda:
- return cupy.fromDlpack(torch.utils.dlpack.to_dlpack(torch_tensor))
+def torch2xp(
+ torch_tensor: "torch.Tensor", *, ops: Optional["Ops"] = None
+) -> ArrayXd: # pragma: no cover
+ """Convert a torch tensor to a numpy or cupy tensor depending on the `ops` parameter.
+ If `ops` is `None`, the type of the resultant tensor will be determined by the source tensor's device.
+ """
+ from .api import NumpyOps
+
+ assert_pytorch_installed()
+ if is_torch_gpu_array(torch_tensor):
+ if isinstance(ops, NumpyOps):
+ return torch_tensor.detach().cpu().numpy()
+ else:
+ return cupy.fromDlpack(torch.utils.dlpack.to_dlpack(torch_tensor))
  else:
- return torch_tensor.detach().numpy()
+ if isinstance(ops, NumpyOps) or ops is None:
+ return torch_tensor.detach().numpy()
+ else:
+ return cupy.asarray(torch_tensor)
 
 
 def xp2tensorflow(
@@ -382,24 +417,33 @@ def xp2tensorflow(
  return tf_tensor
 
 
-def tensorflow2xp(tf_tensor: "tf.Tensor") -> ArrayXd: # pragma: no cover
- """Convert a Tensorflow tensor to numpy or cupy tensor."""
+def tensorflow2xp(
+ tf_tensor: "tf.Tensor", *, ops: Optional["Ops"] = None
+) -> ArrayXd: # pragma: no cover
+ """Convert a Tensorflow tensor to numpy or cupy tensor depending on the `ops` parameter.
+ If `ops` is `None`, the type of the resultant tensor will be determined by the source tensor's device.
+ """
+ from .api import NumpyOps
+
  assert_tensorflow_installed()
- if tf_tensor.device is not None:
- _, device_type, device_num = tf_tensor.device.rsplit(":", 2)
- else:
- device_type = "CPU"
- if device_type == "CPU" or not has_cupy:
- return tf_tensor.numpy()
+ if is_tensorflow_gpu_array(tf_tensor):
+ if isinstance(ops, NumpyOps):
+  return tf_tensor.numpy()
+ else:
+  dlpack_tensor = tensorflow.experimental.dlpack.to_dlpack(tf_tensor)
+  return cupy.fromDlpack(dlpack_tensor)
  else:
- dlpack_tensor = tensorflow.experimental.dlpack.to_dlpack(tf_tensor)
- return cupy.fromDlpack(dlpack_tensor)
+ if isinstance(ops, NumpyOps) or ops is None:
+ return tf_tensor.numpy()
+ else:
+ return cupy.asarray(tf_tensor.numpy())
 
 
 def xp2mxnet(
  xp_tensor: ArrayXd, requires_grad: bool = False
 ) -> "mx.nd.NDArray": # pragma: no cover
  """Convert a numpy or cupy tensor to a MXNet tensor."""
+ assert_mxnet_installed()
  if hasattr(xp_tensor, "toDlpack"):
  dlpack_tensor = xp_tensor.toDlpack() # type: ignore
  mx_tensor = mx.nd.from_dlpack(dlpack_tensor)
@@ -410,12 +454,23 @@ def xp2mxnet(
  return mx_tensor
 
 
-def mxnet2xp(mx_tensor: "mx.nd.NDArray") -> ArrayXd: # pragma: no cover
+def mxnet2xp(
+ mx_tensor: "mx.nd.NDArray", *, ops: Optional["Ops"] = None
+) -> ArrayXd: # pragma: no cover
  """Convert a MXNet tensor to a numpy or cupy tensor."""
- if mx_tensor.context.device_type != "cpu":
- return cupy.fromDlpack(mx_tensor.to_dlpack_for_write())
+ from .api import NumpyOps
+
+ assert_mxnet_installed()
+ if is_mxnet_gpu_array(mx_tensor):
+ if isinstance(ops, NumpyOps):
+ return mx_tensor.detach().asnumpy()
+ else:
+ return cupy.fromDlpack(mx_tensor.to_dlpack_for_write())
  else:
- return mx_tensor.detach().asnumpy()
+ if isinstance(ops, NumpyOps) or ops is None:
+ return mx_tensor.detach().asnumpy()
+ else:
+ return cupy.asarray(mx_tensor.asnumpy())
 
 
 # This is how functools.partials seems to do it, too, to retain the return type