Adds option to set object-store size when starting Ray cluster, uses a process group for managing ray child processes (#1274)

* Updates stop() to kill the ray process group created by init_cluster().

Signed-off-by: rlratzel <[email protected]>

* Creates a process group when running the 'ray start' command and kills the process group in RayClient.stop() to ensure all child processes are terminated.

Signed-off-by: rlratzel <[email protected]>

* Adds support for creating a process group for the ray subprocesses and killing the group to better ensure all processes are killed.

Signed-off-by: rlratzel <[email protected]>

* Wip ab/final metrics (#996)

* Adding prometheus and grafana to the nemo curator metrics path

Signed-off-by: Abhinav Garg <[email protected]>

* Implement safe extraction for tar files in file_utils.py

- Added functions `_is_safe_path` and `tar_safe_extract` to ensure safe extraction of tar files, preventing path traversal attacks.
- Included necessary imports and updated the file structure by removing outdated files from the ray-curator module.

Signed-off-by: [Your Name] [[email protected]]
Signed-off-by: Abhinav Garg <[email protected]>

* Refactor references from ray_curator to nemo_curator across multiple files

- Updated file paths and comments in api-design.md, __init__.py, client.py, and start_prometheus_grafana.py to reflect the new nemo_curator namespace.
- Changed package name in package_info.py from ray_curator to nemo_curator.

Signed-off-by: [Your Name] [[email protected]]
Signed-off-by: Abhinav Garg <[email protected]>

* Rename function `get_ray_client` to `start_prometheus_grafana` in start_prometheus_grafana.py for clarity and consistency with the new metrics path. Update the function call in the main execution block accordingly.

Signed-off-by: Abhinav Garg <[email protected]>

* Adding prometheus and grafana to the nemo curator metrics path

Signed-off-by: Abhinav Garg <[email protected]>

* Adding README for metrics

Signed-off-by: Abhinav Garg <[email protected]>

---------

Signed-off-by: Abhinav Garg <[email protected]>
Signed-off-by: [Your Name] [[email protected]]
Co-authored-by: Sarah Yurick <[email protected]>
Signed-off-by: rlratzel <[email protected]>

* undoes merge mistakes.

Signed-off-by: rlratzel <[email protected]>

* Fixes typo.

Signed-off-by: rlratzel <[email protected]>

* Adds OSError to exception handlers to handle process groups that have already terminated.

Signed-off-by: rlratzel <[email protected]>

---------

Signed-off-by: rlratzel <[email protected]>
Signed-off-by: Abhinav Garg <[email protected]>
Signed-off-by: [Your Name] [[email protected]]
Co-authored-by: Sarah Yurick <[email protected]>

Author: rlratzel

nemo_curator/core/client.py

@@ -14,6 +14,7 @@
 
 import atexit
 import os
+import signal
 import socket
 import subprocess
 from dataclasses import dataclass, field
@@ -56,7 +57,9 @@ class RayClient:
         ray_dashboard_host: The host of the Ray dashboard.
         num_gpus: The number of GPUs to use.
         num_cpus: The number of CPUs to use.
+        object_store_memory: The amount of memory to use for the object store.
         enable_object_spilling: Whether to enable object spilling.
+        ray_stdouterr_capture_file: The file to capture stdout/stderr to.
 
     Note:
         To start monitoring services (Prometheus and Grafana), use the standalone
@@ -72,18 +75,19 @@ class RayClient:
     ray_dashboard_host: str = DEFAULT_RAY_DASHBOARD_HOST
     num_gpus: int | None = None
     num_cpus: int | None = None
+    object_store_memory: int | None = None
     enable_object_spilling: bool = False
     ray_stdouterr_capture_file: str | None = None
 
-    ray_process: subprocess.Popen | None = field(init=False)
+    ray_process: subprocess.Popen | None = field(init=False, default=None)
 
     def __post_init__(self) -> None:
         if self.ray_stdouterr_capture_file and os.path.exists(self.ray_stdouterr_capture_file):
             msg = f"Capture file {self.ray_stdouterr_capture_file} already exists."
             raise FileExistsError(msg)
 
     def start(self) -> None:
-        """Start the Ray cluster, optionally capturing stdout/stderr to a file."""
+        """Start the Ray cluster if not already started, optionally capturing stdout/stderr to a file."""
         if self.include_dashboard:
             # Add Ray metrics service discovery to existing Prometheus configuration
             if is_prometheus_running() and is_grafana_running():
@@ -101,6 +105,17 @@ def start(self) -> None:
                 )
                 logger.warning(msg)
 
+        # Use the RAY_ADDRESS environment variable to determine if Ray is already running.
+        # If a Ray cluster is not running:
+        #   RAY_ADDRESS will be set below when the Ray cluster is started and self.ray_process
+        #   will be assigned the cluster process
+        # If a Ray cluster is already running:
+        #   RAY_ADDRESS will have been set prior to calling start(), presumably by a user starting
+        #   it externally, which means a cluster was already running and self.ray_process will be None.
+        #
+        # Note that the stop() method will stop the cluster only if it was started here and
+        # self.ray_process was assigned, otherwise it leaves it running with the assumption it
+        # was started externally and should not be stopped.
         if os.environ.get("RAY_ADDRESS"):
             logger.info("Ray is already running. Skipping the setup.")
         else:
@@ -119,15 +134,16 @@ def start(self) -> None:
             ip_address = socket.gethostbyname(socket.gethostname())
 
             self.ray_process = init_cluster(
-                self.ray_port,
-                self.ray_temp_dir,
-                self.ray_dashboard_port,
-                self.ray_metrics_port,
-                self.ray_client_server_port,
-                self.ray_dashboard_host,
-                self.num_gpus,
-                self.num_cpus,
-                self.enable_object_spilling,
+                ray_port=self.ray_port,
+                ray_temp_dir=self.ray_temp_dir,
+                ray_dashboard_port=self.ray_dashboard_port,
+                ray_metrics_port=self.ray_metrics_port,
+                ray_client_server_port=self.ray_client_server_port,
+                ray_dashboard_host=self.ray_dashboard_host,
+                num_gpus=self.num_gpus,
+                num_cpus=self.num_cpus,
+                object_store_memory=self.object_store_memory,
+                enable_object_spilling=self.enable_object_spilling,
                 block=True,
                 ip_address=ip_address,
                 stdouterr_capture_file=self.ray_stdouterr_capture_file,
@@ -140,8 +156,21 @@ def start(self) -> None:
 
     def stop(self) -> None:
         if self.ray_process:
-            self.ray_process.kill()
-            self.ray_process.wait()
+            # Kill the entire process group to ensure child processes are terminated
+            try:
+                os.killpg(os.getpgid(self.ray_process.pid), signal.SIGTERM)
+                self.ray_process.wait(timeout=5)
+            except subprocess.TimeoutExpired:
+                # Force kill if graceful termination doesn't work
+                try:
+                    os.killpg(os.getpgid(self.ray_process.pid), signal.SIGKILL)
+                    self.ray_process.wait()
+                except (ProcessLookupError, OSError):
+                    # Process group not found or process group already terminated
+                    pass
+            except (ProcessLookupError, OSError):
+                # Process group not found or process group already terminated
+                pass
             # Reset the environment variable for RAY_ADDRESS
             os.environ.pop("RAY_ADDRESS", None)
             # Currently there is no good way of stopping a particular Ray cluster. https://github.com/ray-project/ray/issues/54989

nemo_curator/core/utils.py

@@ -77,6 +77,7 @@ def init_cluster(  # noqa: PLR0913
     ray_dashboard_host: str,
     num_gpus: int | None = None,
     num_cpus: int | None = None,
+    object_store_memory: int | None = None,
     enable_object_spilling: bool = False,
     block: bool = True,
     ip_address: str | None = None,
@@ -99,6 +100,8 @@ def init_cluster(  # noqa: PLR0913
     ray_command.extend(["--dashboard-port", str(ray_dashboard_port)])
     ray_command.extend(["--ray-client-server-port", str(ray_client_server_port)])
     ray_command.extend(["--temp-dir", ray_temp_dir])
+    if object_store_memory is not None:
+        ray_command.extend(["--object-store-memory", str(object_store_memory)])
     ray_command.extend(["--disable-usage-stats"])
     if enable_object_spilling:
         ray_command.extend(
@@ -124,8 +127,10 @@ def init_cluster(  # noqa: PLR0913
 
     if stdouterr_capture_file:
         with open(stdouterr_capture_file, "w") as f:
-            proc = subprocess.Popen(ray_command, shell=False, stdout=f, stderr=subprocess.STDOUT)  # noqa: S603
+            proc = subprocess.Popen(  # noqa: S603
+                ray_command, shell=False, stdout=f, stderr=subprocess.STDOUT, start_new_session=True
+            )
     else:
-        proc = subprocess.Popen(ray_command, shell=False)  # noqa: S603
+        proc = subprocess.Popen(ray_command, shell=False, start_new_session=True)  # noqa: S603
     logger.info(f"Ray start command: {' '.join(ray_command)}")
     return proc

tests/core/test_get_ray_client.py

@@ -51,11 +51,11 @@ def _assert_ray_stdouterr_output(stdouterr_capture_file: str) -> None:
             with open(stdouterr_capture_file) as f:
                 if "Ray runtime started." in f.read():
                     break
-        if elapsed >= timeout:
-            msg = f"Expected output not found in {stdouterr_capture_file} after {timeout} seconds"
-            raise AssertionError(msg)
         time.sleep(1)
         elapsed += 1
+    if elapsed >= timeout:
+        msg = f"Expected output not found in {stdouterr_capture_file} after {timeout} seconds"
+        raise AssertionError(msg)
 
 
 @pytest.fixture(scope="module")