chore: import upstream snapshot with attribution

python/ray/_private/accelerators/furiosa.py

@@ -0,0 +1,209 @@
+import logging
+import os
+from functools import lru_cache
+from typing import List, Optional, Tuple
+
+from ray._private.accelerators.accelerator import AcceleratorManager
+from ray._private.ray_constants import env_bool
+
+logger = logging.getLogger(__name__)
+
+# Ray uses ``FURIOSA_DEVICES`` to track which Furiosa NPUs are assigned to a
+# worker/actor process. The value uses ``npu:<index>`` notation. Ray's
+# scheduler operates at the device level, so the value Ray writes is always
+# the device-level form (e.g. ``npu:0,npu:3``). Bare integer IDs are also
+# accepted on read for convenience.
+#
+# Note that ``furiosa-llm``'s Python API does not honor ``FURIOSA_DEVICES``
+# automatically; callers must pass ``devices=os.environ["FURIOSA_DEVICES"]``
+# (or an equivalent list) explicitly to ``furiosa_llm.LLM(...)``. The
+# ``furiosa-llm`` CLI does read the value but accepts a richer
+# ``npu:X:Y`` (PE-level) form that Ray does not currently preserve through
+# worker scheduling; see ``_strip_npu_prefix`` below.
+FURIOSA_VISIBLE_DEVICES_ENV_VAR = "FURIOSA_DEVICES"
+NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR = "RAY_EXPERIMENTAL_NOSET_FURIOSA_DEVICES"
+
+_FURIOSA_DEVICE_PREFIX = "npu:"
+
+
+@lru_cache(maxsize=None)
+def _ensure_furiosa_initialized() -> bool:
+    """Run ``furiosa_smi_py.init()`` exactly once per process."""
+    from furiosa_smi_py import init
+
+    init()
+    return True
+
+
+def _get_furiosa_list_devices():
+    """Lazy import + one-shot init of ``furiosa_smi_py``.
+
+    Returns the current ``list_devices`` callable on success, or ``None`` if
+    the SDK is unavailable or initialization fails. ``list_devices`` is
+    re-imported on each call so test monkeypatches on the module attribute
+    take effect.
+    """
+    try:
+        _ensure_furiosa_initialized()
+        from furiosa_smi_py import list_devices
+    except Exception as e:
+        logger.debug("furiosa_smi_py is unavailable: %s", e)
+        return None
+    return list_devices
+
+
+def _strip_npu_prefix(token: str) -> str:
+    """Return the numeric device index from an ``npu:<id>`` token.
+
+    Accepts bare integers (``"3"``) as well as the prefixed form
+    (``"npu:3"``) so that values written by other tooling round-trip
+    cleanly.
+    """
+    token = token.strip()
+    if token.startswith(_FURIOSA_DEVICE_PREFIX):
+        token = token[len(_FURIOSA_DEVICE_PREFIX) :]
+    # ``furiosa-llm`` accepts both ``npu:X`` (whole NPU) and ``npu:X:Y``
+    # (PE-level, e.g. ``npu:0:0-3`` for fused PE 0-3 of NPU 0). Ray's
+    # scheduler currently operates at the device level only, so we keep
+    # the device index and drop any trailing PE selector. Round-tripping
+    # PE-level partitioning through worker scheduling is tracked as a
+    # follow-up enhancement.
+    return token.split(":", 1)[0]
+
+
+class FuriosaAcceleratorManager(AcceleratorManager):
+    """FuriosaAI NPU accelerators.
+
+    Resource name is ``FURIOSA``. The accelerator type is reported as
+    ``FURIOSA_<ARCH>`` where ``<ARCH>`` is the architecture identifier
+    that the Furiosa SMI SDK exposes via its ``Arch`` enum. The current
+    SDK variants are ``Rngd``, ``RngdS``, ``RngdMax`` and ``RngdPlus``,
+    which surface here as ``FURIOSA_RNGD``, ``FURIOSA_RNGDS``,
+    ``FURIOSA_RNGDMAX`` and ``FURIOSA_RNGDPLUS`` respectively.
+    Supporting any architecture the SDK reports keeps this manager
+    forward-compatible with new SKUs as Furiosa adds them to
+    ``furiosa_smi_py``.
+
+    Device visibility is tracked through the ``FURIOSA_DEVICES``
+    environment variable, formatted as ``npu:<id>`` tokens. The value
+    can be passed to the ``furiosa-llm`` CLI (e.g.,
+    ``furiosa-llm serve --devices "$FURIOSA_DEVICES" ...``). When
+    invoking the ``furiosa_llm.LLM`` Python API directly, the assigned
+    devices must be passed explicitly, e.g.
+    ``LLM(model_path, devices=os.environ["FURIOSA_DEVICES"])``;
+    ``LLM(devices=None)`` allocates all visible NPUs and would bypass
+    Ray's per-worker isolation.
+    """
+
+    @staticmethod
+    def get_resource_name() -> str:
+        return "FURIOSA"
+
+    @staticmethod
+    def get_visible_accelerator_ids_env_var() -> str:
+        return FURIOSA_VISIBLE_DEVICES_ENV_VAR
+
+    @staticmethod
+    def get_current_process_visible_accelerator_ids() -> Optional[List[str]]:
+        visible_devices = os.environ.get(
+            FuriosaAcceleratorManager.get_visible_accelerator_ids_env_var()
+        )
+        if visible_devices is None:
+            return None
+        if visible_devices == "":
+            return []
+        return [
+            _strip_npu_prefix(token)
+            for token in visible_devices.split(",")
+            if token.strip()
+        ]
+
+    @staticmethod
+    def get_current_node_num_accelerators() -> int:
+        """Detects the number of Furiosa NPU devices on the current machine."""
+        list_devices = _get_furiosa_list_devices()
+        if list_devices is None:
+            return 0
+        try:
+            return len(list_devices())
+        except Exception as e:
+            logger.debug("Could not list Furiosa NPU devices: %s", e)
+            return 0
+
+    @staticmethod
+    def get_current_node_accelerator_type() -> Optional[str]:
+        """Gets the architecture of the Furiosa NPU on the current node.
+
+        Returns a string like ``FURIOSA_RNGD``, ``FURIOSA_RNGDMAX``,
+        ``FURIOSA_RNGDS``, or ``FURIOSA_RNGDPLUS``. Ray assumes a single
+        accelerator type per node, so the architecture of the first detected
+        device is used.
+
+        The architecture is read from
+        ``device.device_info().arch()`` to mirror the upstream Furiosa SMI
+        interface. The Arch enum string is normalized into an
+        accelerator-type label: ``+`` is mapped to ``plus`` so distinct
+        SKUs do not collide (``rngd+`` becomes ``FURIOSA_RNGDPLUS``,
+        matching the PyO3 enum form ``RngdPlus``), and any remaining
+        non-alphanumeric characters are stripped.
+        """
+        list_devices = _get_furiosa_list_devices()
+        if list_devices is None:
+            return None
+        try:
+            devices = list_devices()
+            if not devices:
+                return None
+
+            arch_obj = devices[0].device_info().arch()
+            if arch_obj is None:
+                return None
+
+            # PyO3 enums typically stringify as "<EnumName.Variant>",
+            # "EnumName.Variant", or just "Variant". Take the trailing
+            # component.
+            raw = str(arch_obj).split(".")[-1].strip()
+            if not raw:
+                return None
+            # Map special suffixes to their alphabetic equivalents so that the
+            # ``Arch::ToString`` form ("rngd+") and the PyO3 enum form
+            # ("RngdPlus") produce the same Ray accelerator type label.
+            # Without this, stripping "+" would collapse "rngd+" into
+            # "rngd", colliding with the distinct ``Rngd`` SKU.
+            raw = raw.replace("+", "plus")
+            normalized = "".join(ch for ch in raw if ch.isalnum()).upper()
+            if not normalized:
+                return None
+            return f"FURIOSA_{normalized}"
+        except Exception as e:
+            logger.debug("Failed to detect Furiosa NPU type: %s", e)
+            return None
+
+    @staticmethod
+    def validate_resource_request_quantity(
+        quantity: float,
+    ) -> Tuple[bool, Optional[str]]:
+        if isinstance(quantity, float) and not quantity.is_integer():
+            return (
+                False,
+                f"{FuriosaAcceleratorManager.get_resource_name()} resource quantity"
+                " must be a whole number. Furiosa NPUs do not support"
+                " fractional resource sharing."
+                f" The specified quantity {quantity} is invalid.",
+            )
+        return (True, None)
+
+    @staticmethod
+    def set_current_process_visible_accelerator_ids(
+        visible_furiosa_devices: List[str],
+    ) -> None:
+        if env_bool(NOSET_FURIOSA_VISIBLE_DEVICES_ENV_VAR, False):
+            return
+
+        formatted = ",".join(
+            f"{_FURIOSA_DEVICE_PREFIX}{_strip_npu_prefix(str(d))}"
+            for d in visible_furiosa_devices
+        )
+        os.environ[
+            FuriosaAcceleratorManager.get_visible_accelerator_ids_env_var()
+        ] = formatted