Rayにおけるアクセラレータ(GPU)のサポート

[Accelerator Support — Ray 2\.35\.0](https://docs.ray.io/en/latest/ray-core/scheduling/accelerators.html)の翻訳です。

:::note warn
本書は著者が手動で翻訳したものであり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

アクセラレータ(GPU)は多くの機械学習アプリケーションで重要なものです。Ray Coreはネイティブで事前定義済みの[リソース](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#core-resources)タイプとして多くのアクセラレータをサポートしており、タスクやアクターでそれらのアクセラレータの[リソース要件](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#resource-requirements)を指定することができます。

Ray Coreがネイティブでサポートしているアクセラレータは以下の通りです:

| アクセラレータ | Rayのリソース名 | サポートレベル |
|:--|:--|:--|
| Nvidia GPU  | GPU  | 完全にテストされておりRayチームがサポート  |
| AMD GPU  |  GPU | 実験段階でありコミュニティがサポート  |
| Intel GPU  |  GPU | 実験段階でありコミュニティがサポート  |
| [AWS Neuron Core](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/model-architecture-fit.html)|  neuron_cores | 実験段階でありコミュニティがサポート  |
| Google TPU  |  TPU | 実験段階でありコミュニティがサポート  |
| Intel Gaudi  |  HPU |  実験段階でありコミュニティがサポート |
| Huawei Ascend  | NPU  | 実験段階でありコミュニティがサポート  |

# アクセラレータを持つRayノードの起動

デフォルトでは、Rayは、ノードにおけるアクセラレータリソースの数量を、Rayによって自動検知された物理的なアクセラレータの数量に設定します。必要な場合には、これを[上書き](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#specify-node-resources)することができます。

:::note info
**注意**
Rayのリソースは[論理的](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#logical-resources)なので、マシンに搭載されているアクセラレータの真の数より大きい数量のアクセラレータのリソース(`num_gpus`など)を指定することを妨げたりはしません。この場合、Rayはアクセラレータを必要とするタスクやアクターのスケジューリングの際に、あなたが指定した個数のアクセラレータがあるかのように動作します。タスクやアクターが実際には存在しないアクセラレータを使用しようとした際にのみトラブルが発生します。
:::

**Nvidia GPU**

Rayで利用できるNvidia GPUの数を制限するために、Rayノードを起動する前に環境変数`CUDA_VISIBLE_DEVICES`を設定することができます。例えば、`CUDA_VISIBLE_DEVICES=1,3 ray start --head --num-gpus=2`によって、Rayはデバイス1と3のみを参照できます。

**AMD GPU**

Rayで利用できるAMD GPUの数を制限するために、Rayノードを起動する前に環境変数`ROCR_VISIBLE_DEVICES`を設定することができます。例えば、`ROCR_VISIBLE_DEVICES=1,3 ray start --head --num-gpus=2`によって、Rayはデバイス1と3のみを参照できます。

**Intel GPU**

Rayで利用できるIntel GPUの数を制限するために、Rayノードを起動する前に環境変数`ONEAPI_DEVICE_SELECTOR`を設定することができます。例えば、`ONEAPI_DEVICE_SELECTOR=1,3 ray start --head --num-gpus=2`によって、Rayはデバイス1と3のみを参照できます。

**AWS Neuron Core**

Rayで利用できるAWS Neuro Coreの数を制限するために、Rayノードを起動する前に環境変数`NEURON_RT_VISIBLE_CORES`を設定することができます。例えば、`NEURON_RT_VISIBLE_CORES=1,3 ray start --head --resources='{"neuron_cores": 2}'`によって、Rayはデバイス1と3のみを参照できます。

オーケストレーションの基盤としてEKSとRay on Neuronの他のサンプルに関しては、[AWSのドキュメント](https://awslabs.github.io/data-on-eks/docs/category/inference-on-eks)をご覧ください。

**Google TPU**

Rayで利用できるGoogle TPUの数を制限するために、Rayノードを起動する前に環境変数`TPU_VISIBLE_CHIPS`を設定することができます。例えば、`TPU_VISIBLE_CHIPS=1,3 ray start --head --resources='{"TPU": 2}'`によって、Rayはデバイス1と3のみを参照できます。

**Intel Gaudi**

Rayで利用できるIntel Gaudi HPUの数を制限するために、Rayノードを起動する前に環境変数`HABANA_VISIBLE_MODULES`を設定することができます。例えば、`HABANA_VISIBLE_MODULES=1,3 ray start --head --resources='{"HPU": 2}'`によって、Rayはデバイス1と3のみを参照できます。

**Huawei Ascend**

Rayで利用できるHuawei Ascend NPUの数を制限するために、Rayノードを起動する前に環境変数`ASCEND_RT_VISIBLE_DEVICES`を設定することができます。例えば、`ASCEND_RT_VISIBLE_DEVICES=1,3 ray start --head --resources='{"NPU": 2}'`によって、Rayはデバイス1と3のみを参照できます。

# タスクやアクターでのアクセラレータの使用

タスクやアクターでアクセラレータが必要な場合、対応する[リソース要件](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#resource-requirements)を指定することができます(例: `@ray.remote(num_gpus=1)`)。すると、Rayは十分なアクセラレータリソースを持つノードにタスクやアクターをスケジュールし、タスクやアクターのコードを実行する前に、対応する環境変数(`CUDA_VISIBLE_DEVICES`など)を設定することで、タスクやアクターにアクセラレータを割り当てます。

**Nvidia GPU**

```py
import os
import ray

ray.init(num_gpus=2)

@ray.remote(num_gpus=1)
class GPUActor:
    def ping(self):
        print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
        print("CUDA_VISIBLE_DEVICES: {}".format(os.environ["CUDA_VISIBLE_DEVICES"]))

@ray.remote(num_gpus=1)
def gpu_task():
    print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
    print("CUDA_VISIBLE_DEVICES: {}".format(os.environ["CUDA_VISIBLE_DEVICES"]))

gpu_actor = GPUActor.remote()
ray.get(gpu_actor.ping.remote())
# The actor uses the first GPU so the task uses the second one.
ray.get(gpu_task.remote())
```
```
(GPUActor pid=52420) GPU IDs: [0]
(GPUActor pid=52420) CUDA_VISIBLE_DEVICES: 0
(gpu_task pid=51830) GPU IDs: [1]
(gpu_task pid=51830) CUDA_VISIBLE_DEVICES: 1
```

**AMD GPU**

```py
import os
import ray

ray.init(num_gpus=2)

@ray.remote(num_gpus=1)
class GPUActor:
    def ping(self):
        print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
        print("ROCR_VISIBLE_DEVICES: {}".format(os.environ["ROCR_VISIBLE_DEVICES"]))

@ray.remote(num_gpus=1)
def gpu_task():
    print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
    print("ROCR_VISIBLE_DEVICES: {}".format(os.environ["ROCR_VISIBLE_DEVICES"]))

gpu_actor = GPUActor.remote()
ray.get(gpu_actor.ping.remote())
# The actor uses the first GPU so the task uses the second one.
ray.get(gpu_task.remote())
```
```
(GPUActor pid=52420) GPU IDs: [0]
(GPUActor pid=52420) ROCR_VISIBLE_DEVICES: 0
(gpu_task pid=51830) GPU IDs: [1]
(gpu_task pid=51830) ROCR_VISIBLE_DEVICES: 1
```

**Intel GPU**

```py
import os
import ray

ray.init(num_gpus=2)

@ray.remote(num_gpus=1)
class GPUActor:
    def ping(self):
        print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
        print("ONEAPI_DEVICE_SELECTOR: {}".format(os.environ["ONEAPI_DEVICE_SELECTOR"]))

@ray.remote(num_gpus=1)
def gpu_task():
    print("GPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))
    print("ONEAPI_DEVICE_SELECTOR: {}".format(os.environ["ONEAPI_DEVICE_SELECTOR"]))

gpu_actor = GPUActor.remote()
ray.get(gpu_actor.ping.remote())
# The actor uses the first GPU so the task uses the second one.
ray.get(gpu_task.remote())
```
```
(GPUActor pid=52420) GPU IDs: [0]
(GPUActor pid=52420) ONEAPI_DEVICE_SELECTOR: 0
(gpu_task pid=51830) GPU IDs: [1]
(gpu_task pid=51830) ONEAPI_DEVICE_SELECTOR: 1
```

**AWS Neuron Core**

```py
import os
import ray

ray.init(resources={"neuron_cores": 2})

@ray.remote(resources={"neuron_cores": 1})
class NeuronCoreActor:
    def ping(self):
        print("Neuron Core IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["neuron_cores"]))
        print("NEURON_RT_VISIBLE_CORES: {}".format(os.environ["NEURON_RT_VISIBLE_CORES"]))

@ray.remote(resources={"neuron_cores": 1})
def neuron_core_task():
    print("Neuron Core IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["neuron_cores"]))
    print("NEURON_RT_VISIBLE_CORES: {}".format(os.environ["NEURON_RT_VISIBLE_CORES"]))

neuron_core_actor = NeuronCoreActor.remote()
ray.get(neuron_core_actor.ping.remote())
# The actor uses the first Neuron Core so the task uses the second one.
ray.get(neuron_core_task.remote())
```
```
(NeuronCoreActor pid=52420) Neuron Core IDs: [0]
(NeuronCoreActor pid=52420) NEURON_RT_VISIBLE_CORES: 0
(neuron_core_task pid=51830) Neuron Core IDs: [1]
(neuron_core_task pid=51830) NEURON_RT_VISIBLE_CORES: 1
```

**Google TPU**

```py
import os
import ray

ray.init(resources={"TPU": 2})

@ray.remote(resources={"TPU": 1})
class TPUActor:
    def ping(self):
        print("TPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["TPU"]))
        print("TPU_VISIBLE_CHIPS: {}".format(os.environ["TPU_VISIBLE_CHIPS"]))

@ray.remote(resources={"TPU": 1})
def tpu_task():
    print("TPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["TPU"]))
    print("TPU_VISIBLE_CHIPS: {}".format(os.environ["TPU_VISIBLE_CHIPS"]))

tpu_actor = TPUActor.remote()
ray.get(tpu_actor.ping.remote())
# The actor uses the first TPU so the task uses the second one.
ray.get(tpu_task.remote())
```
```
(TPUActor pid=52420) TPU IDs: [0]
(TPUActor pid=52420) TPU_VISIBLE_CHIPS: 0
(tpu_task pid=51830) TPU IDs: [1]
(tpu_task pid=51830) TPU_VISIBLE_CHIPS: 1
```

**Intel Gaudi**

```py
import os
import ray

ray.init(resources={"HPU": 2})

@ray.remote(resources={"HPU": 1})
class HPUActor:
    def ping(self):
        print("HPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["HPU"]))
        print("HABANA_VISIBLE_MODULES: {}".format(os.environ["HABANA_VISIBLE_MODULES"]))

@ray.remote(resources={"HPU": 1})
def hpu_task():
    print("HPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["HPU"]))
    print("HABANA_VISIBLE_MODULES: {}".format(os.environ["HABANA_VISIBLE_MODULES"]))

hpu_actor = HPUActor.remote()
ray.get(hpu_actor.ping.remote())
# The actor uses the first HPU so the task uses the second one.
ray.get(hpu_task.remote())
```
```
(HPUActor pid=52420) HPU IDs: [0]
(HPUActor pid=52420) HABANA_VISIBLE_MODULES: 0
(hpu_task pid=51830) HPU IDs: [1]
(hpu_task pid=51830) HABANA_VISIBLE_MODULES: 1
```

**Huawei Ascend**

```py
import os
import ray

ray.init(resources={"NPU": 2})

@ray.remote(resources={"NPU": 1})
class NPUActor:
    def ping(self):
        print("NPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["NPU"]))
        print("ASCEND_RT_VISIBLE_DEVICES: {}".format(os.environ["ASCEND_RT_VISIBLE_DEVICES"]))

@ray.remote(resources={"NPU": 1})
def npu_task():
    print("NPU IDs: {}".format(ray.get_runtime_context().get_accelerator_ids()["NPU"]))
    print("ASCEND_RT_VISIBLE_DEVICES: {}".format(os.environ["ASCEND_RT_VISIBLE_DEVICES"]))

npu_actor = NPUActor.remote()
ray.get(npu_actor.ping.remote())
# The actor uses the first NPU so the task uses the second one.
ray.get(npu_task.remote())
```
```
(NPUActor pid=52420) NPU IDs: [0]
(NPUActor pid=52420) ASCEND_RT_VISIBLE_DEVICES: 0
(npu_task pid=51830) NPU IDs: [1]
(npu_task pid=51830) ASCEND_RT_VISIBLE_DEVICES: 1
```

タスクやアクター内では、[ray\.get\_runtime\_context\(\)\.get\_accelerator\_ids\(\)](https://docs.ray.io/en/latest/ray-core/api/doc/ray.runtime_context.RuntimeContext.get_accelerator_ids.html#ray.runtime_context.RuntimeContext.get_accelerator_ids)によって、タスクやアクターで利用可能なアクセラレータIDのリストが返却されます。通常は、Rayは多くのMLフレームワークがアクセラレータの割り当ての目的のために考慮する、対応する環境変数を自動で設定するので、`get_accelerator_ids()`を呼び出す必要はありません。

:::note info
**注意:** 上で定義したリモート関数やアクターでは、実際にどのアクセラレータも使用しません。Rayは少なくとも1つのアクセラレータを持つノードにそれをスケジュールし、実行される期間一つのアクセラレータを予約しますが、実際にアクセラレータを使うかどうかは関数に依存します。これは、通常TensorFlowのような外部ライブラリを通じて行われます。いかに、実際にアクセラレータを使う例を示します。このサンプルを動作させるには、GPUバージョンのTensorFlowをインストールする必要があります。

```py
@ray.remote(num_gpus=1)
def gpu_task():
    import tensorflow as tf

    # Create a TensorFlow session. TensorFlow restricts itself to use the
    # GPUs specified by the CUDA_VISIBLE_DEVICES environment variable.
    tf.Session()
```
:::

:::note info
**注意:** 割り当てられたアクセラレータを無視したり、マシンのすべてのアクセラレータを使うようにすること自体は可能です。Rayはこれを防がず、同時に同じアクセラレータを使う大量のタスクやアクターが存在する状況に陥ることがあります。しかし、Rayはユーザーに上書きされないものと仮定して、多くのディープラーニングフレームワークによって使われるアクセラレータを制限するための環境変数(`CUDA_VISIBLE_DEVICES`など)を自動で設定します。
:::

# 部分的アクセラレータ

Rayでは、[部分的リソース要件](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#fractional-resource-requirements)をサポートしているので、同じアクセラレータを複数のタスクやアクターで共有することができます。

**Nvidia GPU**

```py
ray.init(num_cpus=4, num_gpus=1)

@ray.remote(num_gpus=0.25)
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same GPU.
ray.get([f.remote() for _ in range(4)])
```

**AMD GPU**

```py
ray.init(num_cpus=4, num_gpus=1)

@ray.remote(num_gpus=0.25)
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same GPU.
ray.get([f.remote() for _ in range(4)])
```

**Intel GPU**

```py
ray.init(num_cpus=4, num_gpus=1)

@ray.remote(num_gpus=0.25)
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same GPU.
ray.get([f.remote() for _ in range(4)])
```

**AWS Neuron Core**

AWS Neuron Coreでは、部分的リソースをサポートしていません。

**Google TPU**

Google TPUでは、部分的リソースをサポートしていません。

**Intel Gaudi**

Intel Gaudiでは、部分的リソースをサポートしていません。

**Huawei Ascend**

```py
ray.init(num_cpus=4, resources={"NPU": 1})

@ray.remote(resources={"NPU": 0.25})
def f():
    import time

    time.sleep(1)

# The four tasks created here can execute concurrently
# and share the same NPU.
ray.get([f.remote() for _ in range(4)])
```

:::note info
**注意:** 個々のタスクがアクセラレータの取り分以上を使わないようにするかどうかはユーザーの責任です。PytorchやTensorFlowではメモリー使用量を制限するように設定することができます。
:::

Rayが部分的リソース要件を持つタスクやアクターにノードのアクセラレータを割り当てる際、断片化を回避するために次のアクセラレータに移動する前に、一つのアクセラレータをパッキングします。

```py
ray.init(num_gpus=3)

@ray.remote(num_gpus=0.5)
class FractionalGPUActor:
    def ping(self):
        print("GPU id: {}".format(ray.get_runtime_context().get_accelerator_ids()["GPU"]))

fractional_gpu_actors = [FractionalGPUActor.remote() for _ in range(3)]
# Ray tries to pack GPUs if possible.
[ray.get(fractional_gpu_actors[i].ping.remote()) for i in range(3)]
```
```
(FractionalGPUActor pid=57417) GPU id: [0]
(FractionalGPUActor pid=57416) GPU id: [0]
(FractionalGPUActor pid=57418) GPU id: [1]
```

# GPUリソースを解放しないワーカー

現時点では、GPUを使用するタスク(TensorFlowを通じてなど)をワーカーが実行する際、タスクはGPUのメモリーを配分しますが、タスクが実行を終えた際に解放しない場合があります。これは、次のタイミングで、あるタスクが同じGPUを使おうとした際に問題につながる場合があります。この問題に対処するために、タスクの処理が終了した際にGPUリソースが解放されるように、RayはデフォルトでGPUタスク間でのワーカープロセスの再利用を無効化します。これは、GPUタスクのスケジュールにオーバーヘッドを引き起こするので、[ray\.remote](https://docs.ray.io/en/latest/ray-core/api/doc/ray.remote.html#ray.remote)デコレーターで、`max_calls=0`を設定することで、ワーカーの再利用の再度有効化することができます。

```py
# By default, ray does not reuse workers for GPU tasks to prevent
# GPU resource leakage.
@ray.remote(num_gpus=1)
def leak_gpus():
    import tensorflow as tf

    # This task allocates memory on the GPU and then never release it.
    tf.Session()
```

# アクセラレータのタイプ

Rayでは、リソース固有のアクセラレータのタイプをサポートします。`accelerator_type`オプションを用いることで、特定のタイプのアクセラレータを持つノードでのタスクやアクターの実行を強制することができます。内部では、このアクセラレータタイプオプションは、`"accelerator_type:<type>": 0.001`の[カスタムリソース要件](https://docs.ray.io/en/latest/ray-core/scheduling/resources.html#custom-resources)として実装されています。これによって、利用できる特定のアクセラレータタイプを持つノードへのタスクやアクターの配置を強制することができます。また、これによってマルチノードタイプのオートスケーラーは、当該タイプのリソースに対する需要を認識することができ、場合によっては当該アクセラレータを提供する新たなノードを起動することができます。

```py
from ray.util.accelerators import NVIDIA_TESLA_V100

@ray.remote(num_gpus=1, accelerator_type=NVIDIA_TESLA_V100)
def train(data):
    return "This function was run on a node with a Tesla V100 GPU"

ray.get(train.remote(1))
```

利用できるアクセラレータタイプに関しては、`ray.util.accelerators`をご覧ください。

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)