From ae6c139d3c7eafa3015ec779628321185f76d381 Mon Sep 17 00:00:00 2001
From: HR Wu <5631010+heiruwu@users.noreply.github.com>
Date: Thu, 11 Apr 2024 04:24:17 +0800
Subject: [PATCH] feat(ray): support custom accelerator type (#134)

Because

- We need to support GPU types other than the official accelerator type
supported by
[ray](https://docs.ray.io/en/latest/ray-core/accelerator-types.html)

This commit

- support custom accelerator type env
---
 instill/helpers/const.py      | 1 +
 instill/helpers/ray_config.py | 9 +++++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/instill/helpers/const.py b/instill/helpers/const.py
index 67ad01d..3193ce3 100644
--- a/instill/helpers/const.py
+++ b/instill/helpers/const.py
@@ -114,6 +114,7 @@ class VisualQuestionAnsweringInput:
 ENV_MEMORY = "RAY_MEMORY"
 ENV_TOTAL_VRAM = "RAY_TOTAL_VRAM"
 ENV_RAY_ACCELERATOR_TYPE = "RAY_ACCELERATOR_TYPE"
+ENV_RAY_CUSTOM_RESOURCE = "RAY_CUSTOM_RESOURCE"
 ENV_NUM_OF_GPUS = "RAY_NUM_OF_GPUS"
 ENV_NUM_OF_CPUS = "RAY_NUM_OF_CPUS"
 ENV_NUM_OF_MIN_REPLICAS = "RAY_NUM_OF_MIN_REPLICAS"
diff --git a/instill/helpers/ray_config.py b/instill/helpers/ray_config.py
index a69cb30..9b3b589 100644
--- a/instill/helpers/ray_config.py
+++ b/instill/helpers/ray_config.py
@@ -15,6 +15,7 @@
     ENV_NUM_OF_MAX_REPLICAS,
     ENV_NUM_OF_MIN_REPLICAS,
     ENV_RAY_ACCELERATOR_TYPE,
+    ENV_RAY_CUSTOM_RESOURCE,
     ENV_TOTAL_VRAM,
     RAM_MINIMUM_RESERVE,
     RAM_UPSCALE_FACTOR,
@@ -48,6 +49,10 @@ def __init__(self, deployable: Deployment) -> None:
         if accelerator_type is not None and accelerator_type != "":
             self._update_accelerator_type(accelerator_type)
 
+        custom_resource = os.getenv(ENV_RAY_CUSTOM_RESOURCE)
+        if custom_resource is not None and custom_resource != "":
+            self._update_custom_resource(custom_resource)
+
         num_of_min_replicas = os.getenv(ENV_NUM_OF_MIN_REPLICAS)
         if num_of_min_replicas is not None and num_of_min_replicas != "":
             self._update_min_replicas(int(num_of_min_replicas))
@@ -132,10 +137,10 @@ def _update_accelerator_type(self, accelerator_type: str):
 
         return self
 
-    def _update_num_custom_resource(self, resource_name: str, num: float):
+    def _update_custom_resource(self, resource_name: str):
         if self._deployment.ray_actor_options is not None:
             self._deployment.ray_actor_options.update(
-                {"resources": {resource_name: num}}
+                {"resources": {resource_name: 0.001}}
             )
 
         return self