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