From 44c8646441ac6caf98a5664aab37f5f851221b02 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Mon, 2 Dec 2024 12:03:07 +0000 Subject: [PATCH] Update docstrings --- ray_provider/hooks.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/ray_provider/hooks.py b/ray_provider/hooks.py index bde98bb..7c6ff06 100644 --- a/ray_provider/hooks.py +++ b/ray_provider/hooks.py @@ -416,7 +416,7 @@ def _create_or_update_cluster( def _setup_gpu_driver(self, gpu_device_plugin_yaml: str) -> None: """ - Set up the GPU device plugin if GPU is enabled. Defaults to NVIDIA's plugin + Set up the GPU device plugin if GPU is enabled. :param gpu_device_plugin_yaml: Path or URL to the GPU device plugin YAML. """ @@ -425,7 +425,7 @@ def _setup_gpu_driver(self, gpu_device_plugin_yaml: str) -> None: gpu_driver_name = gpu_driver["metadata"]["name"] if not self.get_daemon_set(gpu_driver_name): - self.log.info("Creating DaemonSet for NVIDIA device plugin...") + self.log.info("Creating DaemonSet for GPU driver...") self.create_daemon_set(gpu_driver_name, gpu_driver) def _setup_load_balancer(self, name: str, namespace: str, context: Context) -> None: @@ -461,7 +461,7 @@ def setup_ray_cluster( :param context: The Airflow task context. :param ray_cluster_yaml: Path to the YAML file defining the Ray cluster. :param kuberay_version: Version of KubeRay to install. - :param gpu_device_plugin_yaml: Path or URL to the GPU device plugin YAML. Defaults to NVIDIA's plugin + :param gpu_device_plugin_yaml: Path or URL to the GPU device plugin YAML. :param update_if_exists: Whether to update the cluster if it already exists. :raises AirflowException: If there's an error setting up the Ray cluster. """ @@ -536,18 +536,21 @@ def delete_ray_cluster(self, ray_cluster_yaml: str, gpu_device_plugin_yaml: str) Execute the operator to delete the Ray cluster. :param ray_cluster_yaml: Path to the YAML file defining the Ray cluster. - :param gpu_device_plugin_yaml: Path or URL to the GPU device plugin YAML. Defaults to NVIDIA's plugin + :param gpu_device_plugin_yaml: Path or URL to the GPU device plugin YAML. :raises AirflowException: If there's an error deleting the Ray cluster. """ try: self._validate_yaml_file(ray_cluster_yaml) - """Delete the NVIDIA GPU device plugin DaemonSet if it exists.""" - gpu_driver = self.load_yaml_content(gpu_device_plugin_yaml) - gpu_driver_name = gpu_driver["metadata"]["name"] + """Delete the GPU device plugin DaemonSet if it exists.""" + if gpu_device_plugin_yaml: + gpu_driver = self.load_yaml_content(gpu_device_plugin_yaml) + gpu_driver_name = gpu_driver["metadata"]["name"] + else: + return if self.get_daemon_set(gpu_driver_name): - self.log.info("Deleting DaemonSet for NVIDIA device plugin...") + self.log.info("Deleting DaemonSet for the GPU device plugin...") self.delete_daemon_set(gpu_driver_name) self.log.info("::group:: Delete Ray Cluster")