From f11416d45fb71dad9ded765fc4f5a67bef4153da Mon Sep 17 00:00:00 2001
From: Heiru Wu <heiru.wu@instill.tech>
Date: Tue, 5 Mar 2024 15:45:01 +0800
Subject: [PATCH] feat(ray): expose cpu/gpu resource allocation config

---
 instill/helpers/build.py      | 10 +++++++++-
 instill/helpers/ray_config.py | 32 +++++++++++++++++++++++---------
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/instill/helpers/build.py b/instill/helpers/build.py
index 9fbee9a..85cd517 100644
--- a/instill/helpers/build.py
+++ b/instill/helpers/build.py
@@ -7,15 +7,19 @@
 
 import instill
 from instill.helpers.const import DEFAULT_DEPENDENCIES
+from instill.utils.logger import Logger
 
 if __name__ == "__main__":
+    Logger.i("[Instill Builder] Setup docker...")
     client = docker.from_env()
     shutil.copyfile(
         __file__.replace("build.py", "Dockerfile"), os.getcwd() + "/Dockerfile"
     )
 
     try:
+        Logger.i("[Instill Builder] Loading config file...")
         with open("instill.yaml", "r", encoding="utf8") as f:
+            Logger.i("[Instill Builder] Parsing config file...")
             config = yaml.safe_load(f)
 
         build = config["build"]
@@ -37,6 +41,7 @@
             packages_str += p + " "
         packages_str += f"instill-sdk=={instill_version}"
 
+        Logger.i("[Instill Builder] Building model image...")
         img, _ = client.images.build(
             path="./",
             rm=True,
@@ -50,8 +55,11 @@
             quiet=False,
         )
         img.tag(f"{registry}/{repo}", tag)
+        Logger.i("[Instill Builder] Pushing model image...")
         client.images.push(f"{registry}/{repo}", tag=tag)
     except Exception as e:
-        print(e)
+        Logger.e("[Instill Builder] Build failed")
+        Logger.e(e)
     finally:
         os.remove("Dockerfile")
+        Logger.i("[Instill Builder] Build successful")
diff --git a/instill/helpers/ray_config.py b/instill/helpers/ray_config.py
index f28753b..b8348ea 100644
--- a/instill/helpers/ray_config.py
+++ b/instill/helpers/ray_config.py
@@ -32,24 +32,34 @@ def __init__(
         self.use_gpu = use_gpu
         # params
         if use_gpu:
-            self._update_num_cpus(0.25)
-            self._update_num_gpus(0.2)
+            self.update_num_cpus(0.25)
+            self.update_num_gpus(0.2)
         else:
-            self._update_num_cpus(0.25)
+            self.update_num_cpus(0.25)
 
-    def _update_num_cpus(self, num_cpus: float):
+    def update_num_cpus(self, num_cpus: float):
         if self._deployment.ray_actor_options is not None:
             self._deployment.ray_actor_options.update({"num_cpus": num_cpus})
 
-    def _update_memory(self, memory: float):
+        return self
+
+    def update_memory(self, memory: float):
         if self._deployment.ray_actor_options is not None:
             self._deployment.ray_actor_options.update({"memory": memory})
 
-    def _update_num_gpus(self, num_gpus: float):
+        return self
+
+    def update_num_gpus(self, num_gpus: float):
         if self._deployment.ray_actor_options is not None:
             self._deployment.ray_actor_options.update({"num_gpus": num_gpus})
 
+        return self
+
     def _determine_vram_usage(self, model_path: str, total_vram: str):
+        warn(
+            "determine vram usage base on file size will soon be removed",
+            PendingDeprecationWarning,
+        )
         if total_vram == "":
             return 0.25
         if os.path.isfile(model_path):
@@ -75,6 +85,10 @@ def _determine_vram_usage(self, model_path: str, total_vram: str):
         raise ModelPathException
 
     def _determine_ram_usage(self, model_path: str):
+        warn(
+            "determine ram usage base on file size will soon be removed",
+            PendingDeprecationWarning,
+        )
         if os.path.isfile(model_path):
             return max(
                 RAM_MINIMUM_RESERVE * (1024 * 1024 * 1024),
@@ -124,13 +138,13 @@ def deploy(self, model_folder_path: str, ray_addr: str, total_vram: str):
 
         if self.use_gpu:
             if model_name in MODEL_VRAM_OVERRIDE_LIST:
-                self._update_num_gpus(MODEL_VRAM_OVERRIDE_LIST[model_name])
+                self.update_num_gpus(MODEL_VRAM_OVERRIDE_LIST[model_name])
             else:
-                self._update_num_gpus(
+                self.update_num_gpus(
                     self._determine_vram_usage(model_folder_path, total_vram)
                 )
         else:
-            self._update_memory(self._determine_ram_usage(model_folder_path))
+            self.update_memory(self._determine_ram_usage(model_folder_path))
 
         if model_name in MODEL_VRAM_OVERRIDE_LIST:
             self.update_min_replicas(1)