Merge pull request #3 from jantonguirao/janton_imperative_mode

dali2: operator placement handling (string, Device, torch.device) and tests

Author: mzient

dali/python/nvidia/dali/experimental/dali2/_device.py

@@ -14,35 +14,43 @@
 
 import nvidia.dali.backend as _backend
 from threading import local
-
+from typing import Union, Optional
 
 class Device:
     _thread_local = local()
 
     def __init__(self, name: str, device_id: int = None):
+        device_type, name_device_id = Device.split_device_type_and_id(name)
+        if name_device_id is not None and device_id is not None:
+            raise ValueError(f"Invalid device name: {name}\n"
+                             f"Ordinal ':{name_device_id}' should not appear "
+                             "in device name when device_id is provided")
         if device_id is None:
-            type_and_id = name.split(":")
-            if len(type_and_id) < 1 or len(type_and_id) > 2:
-                raise ValueError(f"Invalid device name: {name}")
-            device_type = type_and_id[0]
-            if len(type_and_id) == 2:
-                device_id = int(type_and_id[1])
-        else:
-            if ":" in name:
-                raise ValueError(
-                    f"Invalid device name: {name}\n"
-                    f"':' should not appear in device name when device_id is provided"
-                )
-            device_type = name
+            device_id = name_device_id
+
+        if device_type == "cuda":
+            device_type = "gpu"
 
         Device.validate_device_type(device_type)
         if device_id is not None:
             Device.validate_device_id(device_id, device_type)
         else:
             device_id = Device.default_device_id(device_type)
+
         self.device_type = device_type
         self.device_id = device_id
 
+    @staticmethod
+    def split_device_type_and_id(name: str) -> tuple[str, int]:
+        type_and_id = name.split(":")
+        if len(type_and_id) < 1 or len(type_and_id) > 2:
+            raise ValueError(f"Invalid device name: {name}")
+        device_type = type_and_id[0]
+        device_id = None
+        if len(type_and_id) == 2:
+            device_id = int(type_and_id[1])
+        return device_type, device_id
+
     @staticmethod
     def default_device_id(device_type: str) -> int:
         if device_type == "cpu":
@@ -135,3 +143,43 @@ def __exit__(self, exc_type, exc_value, traceback):
 
 Device._thread_local.devices = None
 Device._thread_local.previous_device_ids = None
+
+
+def device(obj: Union[Device, str, "torch.device"], id: Optional[int] = None) -> Device:
+    """
+    Returns a Device object from various input types.
+
+    - If `obj` is already a `Device`, returns it. In this case, `id` must be `None`.
+    - If `obj` is a `str`, parses it as a device name (e.g., `"gpu"`, `"cpu:0"`, `"cuda:1"`). In this case, `id` can be specified.
+      Note: If the string already contains a device id and `id` is also provided, a `ValueError` is raised.
+    - If `obj` is a `torch.device`, converts it to a `Device`. In this case, `id` must be `None`.
+    - If `obj` is None, returns it.
+    - If `obj` is not a `Device`, `str`, or `torch.device` or None, raises a `TypeError`.
+    """
+
+    # None
+    if obj is None:
+        return obj
+
+    # Device instance
+    if isinstance(obj, Device):
+        if id is not None:
+            raise ValueError("Cannot specify id when passing a Device instance")
+        return obj
+
+    if isinstance(obj, str):
+        return Device(obj, id)
+
+    # torch.device detected by duck-typing
+    is_torch_device = (
+        obj.__class__.__module__ == "torch" and
+        obj.__class__.__name__ == "device" and
+        hasattr(obj, "type") and
+        hasattr(obj, "index"))
+    if is_torch_device:
+        dev_type = "gpu" if obj.type == "cuda" else obj.type
+        if id is not None:
+            raise ValueError("Cannot specify id when passing a torch.device")
+        return Device(dev_type, obj.index)
+
+    raise TypeError(f"Cannot convert {type(obj)} to Device")

dali/python/nvidia/dali/experimental/dali2/_tensor.py

@@ -67,7 +67,7 @@ def __init__(
         self,
         data: Optional[Any] = None,
         dtype: Optional[Any] = None,
-        device: Optional[Device] = None,
+        device: Optional[Union[Device, str, "torch.device"]] = None,
         layout: Optional[str] = None,
         batch: Optional[Any] = None,
         index_in_batch: Optional[int] = None,
@@ -92,6 +92,9 @@ def __init__(
 
         copied = False
 
+        from ._device import device as _to_device
+        device = _to_device(device)
+
         from . import _fn
 
         if dtype is not None:

dali/python/nvidia/dali/experimental/dali2/ops.py

@@ -38,17 +38,10 @@ def __init__(
         self._num_inputs = num_inputs
         self._call_arg_names = None if call_arg_names is None else tuple(call_arg_names)
         self._api_type = None
-        if isinstance(device, str):
-            self._device = _device.Device(
-                name=device,
-                device_id=kwargs.get("device_id", _device.Device.default_device_id(device)),
-            )
-        else:
-            if not isinstance(device, _device.Device):
-                raise TypeError(
-                    f"`device` must be a Device instance or a string, got {type(device)}"
-                )
-            self._device = device
+
+        from ._device import device as _to_device
+        self._device = _to_device(device)
+
         self._input_meta = []
         self._arg_meta = {}
         self._num_outputs = None