add virtual cluster

Signed-off-by: 黑驰 <[email protected]>

Author: wumuzi520

BUILD.bazel

@@ -493,6 +493,7 @@ ray_cc_library(
     name = "gcs_pub_sub_lib",
     srcs = [
         "src/ray/gcs/pubsub/gcs_pub_sub.cc",
+        "src/ray/gcs/pubsub/gcs_pub_sub.ant.cc",
     ],
     hdrs = [
         "src/ray/gcs/pubsub/gcs_pub_sub.h",
@@ -1833,6 +1834,34 @@ ray_cc_test(
     ],
 )
 
+ray_cc_test(
+    name = "gcs_virtual_cluster_manager_test",
+    size = "small",
+    srcs = [
+        "src/ray/gcs/gcs_server/test/gcs_virtual_cluster_manager_test.cc",
+    ],
+    tags = ["team:core"],
+    deps = [
+        ":gcs_server_lib",
+        ":gcs_server_test_util",
+        ":gcs_test_util_lib",
+        ":ray_mock",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
+ray_cc_test(
+    name = "virtual_cluster_manager_test",
+    size = "small",
+    srcs = ["src/ray/raylet/virtual_cluster_manager_test.cc"],
+    tags = ["team:core"],
+    deps = [
+        ":ray_mock",
+        ":raylet_lib",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 ray_cc_library(
     name = "gcs_table_storage_lib",
     srcs = glob(
@@ -1956,6 +1985,7 @@ ray_cc_library(
     name = "gcs_client_lib",
     srcs = [
         "src/ray/gcs/gcs_client/accessor.cc",
+        "src/ray/gcs/gcs_client/accessor.ant.cc",
         "src/ray/gcs/gcs_client/gcs_client.cc",
     ],
     hdrs = [

python/ray/_private/ray_constants.py

@@ -118,6 +118,8 @@ def env_set_by_user(key):
 # the local working_dir and py_modules to be uploaded, or these files might get
 # garbage collected before the job starts.
 RAY_RUNTIME_ENV_URI_PIN_EXPIRATION_S_DEFAULT = 10 * 60
+# Environment variable to specify the virtual cluster ID a Ray job belongs to.
+RAY_VIRTUAL_CLUSTER_ID_ENV_VAR = "VIRTUAL_CLUSTER_ID"
 # If set to 1, then `.gitignore` files will not be parsed and loaded into "excludes"
 # when using a local working_dir or py_modules.
 RAY_RUNTIME_ENV_IGNORE_GITIGNORE = "RAY_RUNTIME_ENV_IGNORE_GITIGNORE"

python/ray/_private/state.py

@@ -158,15 +158,15 @@ def _gen_actor_info(self, actor_table_data):
         }
         return actor_info
 
-    def node_table(self):
+    def node_table(self, virtual_cluster_id=None):
         """Fetch and parse the Gcs node info table.
 
         Returns:
             Information about the node in the cluster.
         """
         self._check_connected()
 
-        return self.global_state_accessor.get_node_table()
+        return self.global_state_accessor.get_node_table(virtual_cluster_id)
 
     def job_table(self):
         """Fetch and parse the gcs job table.
@@ -724,7 +724,7 @@ def update_worker_num_paused_threads(self, worker_id, num_paused_threads_delta):
             worker_id, num_paused_threads_delta
         )
 
-    def cluster_resources(self):
+    def cluster_resources(self, virtual_cluster_id=None):
         """Get the current total cluster resources.
 
         Note that this information can grow stale as nodes are added to or
@@ -738,23 +738,25 @@ def cluster_resources(self):
 
         # Calculate total resources.
         total_resources = defaultdict(int)
-        for node_total_resources in self.total_resources_per_node().values():
+        for node_total_resources in self.total_resources_per_node(
+            virtual_cluster_id
+        ).values():
             for resource_id, value in node_total_resources.items():
                 total_resources[resource_id] += value
 
         return dict(total_resources)
 
-    def _live_node_ids(self):
+    def _live_node_ids(self, virtual_cluster_id=None):
         """Returns a set of node IDs corresponding to nodes still alive."""
-        return set(self.total_resources_per_node().keys())
+        return set(self.total_resources_per_node(virtual_cluster_id).keys())
 
-    def available_resources_per_node(self):
+    def available_resources_per_node(self, virtual_cluster_id=None):
         """Returns a dictionary mapping node id to available resources."""
         self._check_connected()
         available_resources_by_id = {}
 
         all_available_resources = (
-            self.global_state_accessor.get_all_available_resources()
+            self.global_state_accessor.get_all_available_resources(virtual_cluster_id)
         )
         for available_resource in all_available_resources:
             message = gcs_pb2.AvailableResources.FromString(available_resource)
@@ -769,11 +771,15 @@ def available_resources_per_node(self):
         return available_resources_by_id
 
     # returns a dict that maps node_id(hex string) to a dict of {resource_id: capacity}
-    def total_resources_per_node(self) -> Dict[str, Dict[str, int]]:
+    def total_resources_per_node(
+        self, virtual_cluster_id=None
+    ) -> Dict[str, Dict[str, int]]:
         self._check_connected()
         total_resources_by_node = {}
 
-        all_total_resources = self.global_state_accessor.get_all_total_resources()
+        all_total_resources = self.global_state_accessor.get_all_total_resources(
+            virtual_cluster_id
+        )
         for node_total_resources in all_total_resources:
             message = gcs_pb2.TotalResources.FromString(node_total_resources)
             # Calculate total resources for this node.
@@ -786,7 +792,7 @@ def total_resources_per_node(self) -> Dict[str, Dict[str, int]]:
 
         return total_resources_by_node
 
-    def available_resources(self):
+    def available_resources(self, virtual_cluster_id=None):
         """Get the current available cluster resources.
 
         This is different from `cluster_resources` in that this will return
@@ -802,7 +808,9 @@ def available_resources(self):
         """
         self._check_connected()
 
-        available_resources_by_id = self.available_resources_per_node()
+        available_resources_by_id = self.available_resources_per_node(
+            virtual_cluster_id
+        )
 
         # Calculate total available resources.
         total_available_resources = defaultdict(int)
@@ -879,13 +887,19 @@ def next_job_id():
 
 @DeveloperAPI
 @client_mode_hook
-def nodes():
+def nodes(virtual_cluster_id=None):
     """Get a list of the nodes in the cluster (for debugging only).
 
     Returns:
         Information about the Ray clients in the cluster.
     """
-    return state.node_table()
+    if not virtual_cluster_id:
+        virtual_cluster_id = ray.get_runtime_context().virtual_cluster_id
+    elif type(virtual_cluster_id) is not str:
+        raise TypeError(
+            f"virtual_cluster_id must be a string, got {type(virtual_cluster_id)}"
+        )
+    return state.node_table(virtual_cluster_id)
 
 
 def workers():
@@ -999,7 +1013,7 @@ def object_transfer_timeline(filename=None):
 
 @DeveloperAPI
 @client_mode_hook
-def cluster_resources():
+def cluster_resources(virtual_cluster_id=None):
     """Get the current total cluster resources.
 
     Note that this information can grow stale as nodes are added to or removed
@@ -1009,12 +1023,18 @@ def cluster_resources():
         A dictionary mapping resource name to the total quantity of that
             resource in the cluster.
     """
-    return state.cluster_resources()
+    if not virtual_cluster_id:
+        virtual_cluster_id = ray.get_runtime_context().virtual_cluster_id
+    elif type(virtual_cluster_id) is not str:
+        raise TypeError(
+            f"virtual_cluster_id must be a string, got {type(virtual_cluster_id)}"
+        )
+    return state.cluster_resources(virtual_cluster_id)
 
 
 @DeveloperAPI
 @client_mode_hook
-def available_resources():
+def available_resources(virtual_cluster_id=None):
     """Get the current available cluster resources.
 
     This is different from `cluster_resources` in that this will return idle
@@ -1028,7 +1048,13 @@ def available_resources():
             is currently not available (i.e., quantity is 0), it will not
             be included in this dictionary.
     """
-    return state.available_resources()
+    if not virtual_cluster_id:
+        virtual_cluster_id = ray.get_runtime_context().virtual_cluster_id
+    elif type(virtual_cluster_id) is not str:
+        raise TypeError(
+            f"virtual_cluster_id must be a string, got {type(virtual_cluster_id)}"
+        )
+    return state.available_resources(virtual_cluster_id)
 
 
 @DeveloperAPI
@@ -1045,7 +1071,7 @@ def available_resources_per_node():
 
 
 @DeveloperAPI
-def total_resources_per_node():
+def total_resources_per_node(virtual_cluster_id=None):
     """Get the current total resources of each live node.
 
     Note that this information can grow stale as tasks start and finish.
@@ -1054,6 +1080,12 @@ def total_resources_per_node():
         A dictionary mapping node hex id to total resources dictionary.
     """
 
+    if not virtual_cluster_id:
+        virtual_cluster_id = ray.get_runtime_context().virtual_cluster_id
+    elif type(virtual_cluster_id) is not str:
+        raise TypeError(
+            f"virtual_cluster_id must be a string, got {type(virtual_cluster_id)}"
+        )
     return state.total_resources_per_node()
 
 

python/ray/_private/worker.py

@@ -570,6 +570,10 @@ def current_cluster_and_job(self):
         assert isinstance(self.current_job_id, ray.JobID)
         return self.node.cluster_id, self.current_job_id
 
+    @property
+    def current_virtual_cluster_id(self):
+        return os.environ.get(ray_constants.RAY_VIRTUAL_CLUSTER_ID_ENV_VAR, "")
+
     @property
     def runtime_env(self):
         """Get the runtime env in json format"""

python/ray/_raylet.pyx

@@ -144,6 +144,7 @@ from ray.includes.unique_ids cimport (
     CObjectID,
     CPlacementGroupID,
     ObjectIDIndexType,
+    CVirtualClusterID,
 )
 from ray.includes.libcoreworker cimport (
     ActorHandleSharedPtr,

python/ray/dashboard/modules/job/common.py

@@ -9,12 +9,17 @@
 
 from ray._private import ray_constants
 from ray._private.event.export_event_logger import get_export_event_logger
-from ray._private.gcs_utils import GcsAioClient
+from ray._private.gcs_utils import GcsAioClient, GcsChannel
 from ray._private.runtime_env.packaging import parse_uri
+from ray.core.generated import gcs_service_pb2_grpc
 from ray.core.generated.export_event_pb2 import ExportEvent
 from ray.core.generated.export_submission_job_event_pb2 import (
     ExportSubmissionJobEventData,
 )
+from ray.core.generated.gcs_service_pb2 import (
+    CreateJobClusterRequest,
+    RemoveVirtualClusterRequest,
+)
 from ray.util.annotations import PublicAPI
 
 # NOTE(edoakes): these constants should be considered a public API because
@@ -104,6 +109,8 @@ class JobInfo:
     #: The driver process exit code after the driver executed. Return None if driver
     #: doesn't finish executing
     driver_exit_code: Optional[int] = None
+    #: The job cluster this job belongs to.
+    job_cluster_id: Optional[str] = None
 
     def __post_init__(self):
         if isinstance(self.status, str):
@@ -378,6 +385,62 @@ async def get_job_info(job_id: str):
         }
 
 
+class VirtualClusterClient:
+    def __init__(self, gcs_address):
+        self._gcs_channel = GcsChannel(gcs_address=gcs_address, aio=True)
+        self._gcs_channel.connect()
+
+        self._gcs_virtual_cluster_info_stub = (
+            gcs_service_pb2_grpc.VirtualClusterInfoGcsServiceStub(
+                self._gcs_channel.channel()
+            )
+        )
+
+    def build_job_cluster_id(self, job_id, virtual_cluster_id):
+        """
+        Constructs a unique job cluster ID by combining
+        the virtual cluster ID and job ID.
+        Note:
+            The format needs to remain consistent with `virtual_cluster_id.h`.
+        """
+        return f"{virtual_cluster_id}##{job_id}"
+
+    async def create_job_cluster(self, job_id, virtual_cluster_id, replica_sets):
+        request = CreateJobClusterRequest(
+            job_id=job_id,
+            virtual_cluster_id=virtual_cluster_id,
+            replica_sets=replica_sets,
+        )
+
+        reply = await self._gcs_virtual_cluster_info_stub.CreateJobCluster(request)
+
+        if reply.status.code != 0:
+            logger.warning(
+                f"Failed to create job cluster for job ID '{job_id}' in "
+                f"virtual cluster '{virtual_cluster_id}', "
+                f"message: {reply.status.message}"
+            )
+            return None, reply.status.message
+
+        return reply.job_cluster_id, None
+
+    async def remove_job_cluster(self, job_cluster_id):
+        if job_cluster_id is None:
+            return
+
+        request = RemoveVirtualClusterRequest(
+            virtual_cluster_id=job_cluster_id,
+        )
+
+        reply = await self._gcs_virtual_cluster_info_stub.RemoveVirtualCluster(request)
+
+        if reply.status.code != 0:
+            logger.warning(
+                f"Failed to remove job cluster '{job_cluster_id}',"
+                f" message: {reply.status.message}"
+            )
+
+
 def uri_to_http_components(package_uri: str) -> Tuple[str, str]:
     suffix = Path(package_uri).suffix
     if suffix not in {".zip", ".whl"}:
@@ -426,6 +489,10 @@ class JobSubmitRequest:
     # to reserve for the entrypoint command, separately from any Ray tasks
     # or actors that are created by it.
     entrypoint_resources: Optional[Dict[str, float]] = None
+    # Optional virtual cluster ID for job.
+    virtual_cluster_id: Optional[str] = None
+    # Optional replica sets for job
+    replica_sets: Optional[Dict[str, int]] = None
 
     def __post_init__(self):
         if not isinstance(self.entrypoint, str):
@@ -511,6 +578,31 @@ def __post_init__(self):
                             f"got {type(v)}"
                         )
 
+        if self.virtual_cluster_id is not None and not isinstance(
+            self.virtual_cluster_id, str
+        ):
+            raise TypeError(
+                "virtual_cluster_id must be a string if provided, "
+                f"got {type(self.virtual_cluster_id)}"
+            )
+
+        if self.replica_sets is not None:
+            if not isinstance(self.replica_sets, dict):
+                raise TypeError(
+                    "replica_sets must be a dict, " f"got {type(self.replica_sets)}"
+                )
+            else:
+                for k in self.replica_sets.keys():
+                    if not isinstance(k, str):
+                        raise TypeError(
+                            "replica_sets keys must be strings, " f"got {type(k)}"
+                        )
+                for v in self.replica_sets.values():
+                    if not isinstance(v, int):
+                        raise TypeError(
+                            "replica_sets values must be integers, " f"got {type(v)}"
+                        )
+
 
 @dataclass
 class JobSubmitResponse: