minor changes on command

Author: yli1 user

matrix/__main__.py

@@ -0,0 +1,4 @@
+from .cli import main
+
+if __name__ == "__main__":
+    main()

matrix/cli.py

@@ -80,6 +80,7 @@ def start_cluster(
         slurm: tp.Dict[str, tp.Union[str, int]] | None = None,
         local: tp.Dict[str, tp.Union[str, int]] | None = None,
         enable_grafana: bool = False,
+        force_new_head: bool = False,
     ):
         """
         Starts the Ray cluster with additional keyword arguments. Only do this for new clusters.
@@ -97,11 +98,18 @@ def start_cluster(
             slurm (dict, optional): resources for slurm cluster.
             local (dict, optional): resources for local cluster.
             enable_grafana (bool, optional): If True, enable prometheus and grafana dashboard.
-        
+            force_new_head (bool): force to remove head.json if haven't run 'matrix stop_cluster'.
+            
         Returns:
             None
         """
-        self.cluster.start(add_workers, slurm, local, enable_grafana=enable_grafana)
+        self.cluster.start(
+            add_workers,
+            slurm,
+            local,
+            enable_grafana=enable_grafana,
+            force_new_head=force_new_head,
+        )
 
     def stop_cluster(self):
         """
@@ -370,7 +378,3 @@ def check_health(
 
 def main():
     fire.Fire(Cli)
-
-
-if __name__ == "__main__":
-    main()

matrix/cluster/ray_cluster.py

@@ -149,7 +149,11 @@ def start_grafana(self, force: bool):
                 num_gpus=0,
                 max_restarts=3,  # Allow 3 automatic retries
                 max_task_retries=-1,
-            ).remote(cluster_info)
+            ).remote(
+                cluster_info.temp_dir,
+                cluster_info.prometheus_port,
+                cluster_info.grafana_port,
+            )
             ray.get(actor.start.remote())
             return "Successfully started Grafana dashboard"
 
@@ -159,6 +163,7 @@ def start(
         slurm: tp.Dict[str, tp.Union[str, int]] | None,
         local: tp.Dict[str, tp.Union[str, int]] | None,
         enable_grafana: bool = False,
+        force_new_head: bool = False,
     ):
         """
         Starts a Ray cluster on Slurm.
@@ -168,19 +173,24 @@ def start(
 
         Args:
             add_workers (int): The number of worker nodes to start.
-            requirements (dict): Slurm resource requirements for worker nodes.
-                                 e.g., {'qos': '...', 'partition': '...', 'gpus-per-node': 8}.
-            head_requirements (dict): optional to specify head requirements when launching head.
-            executor (str): Slurm executor to use (default: "slurm").
+            slurm (dict, optional): resources requirements for slurm cluster.
+                                    e.g., {'qos': '...', 'partition': '...', 'gpus-per-node': 8}.
+            local (dict, optional): resources requirements for local cluster.
             enable_grafana (bool): Whether to start Prometheus and Grafana
                                           for monitoring (default: True).
+            force_new_head (bool): force to remove head.json if haven't run 'matrix stop_cluster'.
         """
         common_params = {"account", "partition", "qos", "exclusive"}
         start_wait_time_seconds = 60
         worker_wait_timeout_seconds = 60
         requirements = slurm or local or {}
         executor = "slurm" if slurm else "local"
 
+        if force_new_head:
+            # remove existing head.json
+            if self._cluster_json.exists():
+                self._cluster_json.unlink()
+
         if self._cluster_json.exists():
             print(f"Adding workers to existing cluster:\n{self.cluster_info()}")
             # todo: check the cluser is alive

matrix/cluster/ray_dashboard_job.py

@@ -15,8 +15,6 @@
 
 import ray
 
-from matrix.common.cluster_info import ClusterInfo
-
 # Configure logging
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
@@ -38,14 +36,14 @@ class RayDashboardJob:
     properly when the actor is killed.
     """
 
-    def __init__(self, cluster_info: ClusterInfo):
+    def __init__(self, temp_dir: str, prometheus_port: int, grafana_port: int):
         """
         Initialize the RayDashboardJob actor.
         """
         self.head_env = os.environ.copy()
-        self.temp_dir = cluster_info.temp_dir
-        self.prometheus_port = cluster_info.prometheus_port
-        self.grafana_port = cluster_info.grafana_port
+        self.temp_dir = temp_dir
+        self.prometheus_port = prometheus_port
+        self.grafana_port = grafana_port
         self.processes: List[subprocess.Popen[str]] = []
         self.monitor_thread: threading.Thread | None = None
         self.should_run = True