feat(RHOAIENG-26482): add gcs fault tolerance

Author: kryanbeane

src/codeflare_sdk/ray/rayjobs/__init__.py

@@ -1,2 +1,3 @@
 from .rayjob import RayJob, RayJobClusterConfig
 from .status import RayJobDeploymentStatus, CodeflareRayJobStatus, RayJobInfo
+from .config import RayJobClusterConfig

src/codeflare_sdk/ray/rayjobs/config.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 """
-The config sub-module contains the definition of the RayJobClusterConfigV2 dataclass,
+The config sub-module contains the definition of the RayJobClusterConfig dataclass,
 which is used to specify resource requirements and other details when creating a
 Cluster object.
 """
@@ -139,6 +139,16 @@ class RayJobClusterConfig:
             A list of V1Volume objects to add to the Cluster
         volume_mounts:
             A list of V1VolumeMount objects to add to the Cluster
+        enable_gcs_ft:
+            A boolean indicating whether to enable GCS fault tolerance.
+        enable_usage_stats:
+            A boolean indicating whether to capture and send Ray usage stats externally.
+        redis_address:
+            The address of the Redis server to use for GCS fault tolerance, required when enable_gcs_ft is True.
+        redis_password_secret:
+            Kubernetes secret reference containing Redis password. ex: {"name": "secret-name", "key": "password-key"}
+        external_storage_namespace:
+            The storage namespace to use for GCS fault tolerance. By default, KubeRay sets it to the UID of RayCluster.
     """
 
     head_cpu_requests: Union[int, str] = 2
@@ -165,8 +175,39 @@ class RayJobClusterConfig:
     annotations: Dict[str, str] = field(default_factory=dict)
     volumes: list[V1Volume] = field(default_factory=list)
     volume_mounts: list[V1VolumeMount] = field(default_factory=list)
+    enable_gcs_ft: bool = False
+    enable_usage_stats: bool = False
+    redis_address: Optional[str] = None
+    redis_password_secret: Optional[Dict[str, str]] = None
+    external_storage_namespace: Optional[str] = None
 
     def __post_init__(self):
+        if self.enable_usage_stats:
+            self.envs["RAY_USAGE_STATS_ENABLED"] = "1"
+        else:
+            self.envs["RAY_USAGE_STATS_ENABLED"] = "0"
+
+        if self.enable_gcs_ft:
+            if not self.redis_address:
+                raise ValueError(
+                    "redis_address must be provided when enable_gcs_ft is True"
+                )
+
+            if self.redis_password_secret and not isinstance(
+                self.redis_password_secret, dict
+            ):
+                raise ValueError(
+                    "redis_password_secret must be a dictionary with 'name' and 'key' fields"
+                )
+
+            if self.redis_password_secret and (
+                "name" not in self.redis_password_secret
+                or "key" not in self.redis_password_secret
+            ):
+                raise ValueError(
+                    "redis_password_secret must contain both 'name' and 'key' fields"
+                )
+
         self._validate_types()
         self._memory_to_string()
         self._validate_gpu_config(self.head_accelerators)
@@ -251,6 +292,11 @@ def build_ray_cluster_spec(self, cluster_name: str) -> Dict[str, Any]:
             "workerGroupSpecs": [self._build_worker_group_spec(cluster_name)],
         }
 
+        # Add GCS fault tolerance if enabled
+        if self.enable_gcs_ft:
+            gcs_ft_options = self._build_gcs_ft_options()
+            ray_cluster_spec["gcsFaultToleranceOptions"] = gcs_ft_options
+
         return ray_cluster_spec
 
     def _build_head_group_spec(self) -> Dict[str, Any]:
@@ -453,3 +499,25 @@ def _generate_volumes(self) -> list:
     def _build_env_vars(self) -> list:
         """Build environment variables list."""
         return [V1EnvVar(name=key, value=value) for key, value in self.envs.items()]
+
+    def _build_gcs_ft_options(self) -> Dict[str, Any]:
+        """Build GCS fault tolerance options."""
+        gcs_ft_options = {"redisAddress": self.redis_address}
+
+        if (
+            hasattr(self, "external_storage_namespace")
+            and self.external_storage_namespace
+        ):
+            gcs_ft_options["externalStorageNamespace"] = self.external_storage_namespace
+
+        if hasattr(self, "redis_password_secret") and self.redis_password_secret:
+            gcs_ft_options["redisPassword"] = {
+                "valueFrom": {
+                    "secretKeyRef": {
+                        "name": self.redis_password_secret["name"],
+                        "key": self.redis_password_secret["key"],
+                    }
+                }
+            }
+
+        return gcs_ft_options

src/codeflare_sdk/ray/rayjobs/rayjob.py

@@ -140,7 +140,6 @@ def __init__(
             self.cluster_name = cluster_name
             logger.info(f"Using existing cluster: {self.cluster_name}")
 
-        # Initialize the KubeRay job API client
         self._api = RayjobApi()
 
         logger.info(f"Initialized RayJob: {self.name} in namespace: {self.namespace}")

src/codeflare_sdk/ray/rayjobs/test_rayjob.py

@@ -971,3 +971,24 @@ def test_rayjob_user_override_shutdown_behavior(mocker):
     )
 
     assert rayjob_override_priority.shutdown_after_job_finishes is True
+
+
+def test_build_ray_cluster_spec_with_gcs_ft(mocker):
+    """Test build_ray_cluster_spec with GCS fault tolerance enabled."""
+    from codeflare_sdk.ray.rayjobs.config import RayJobClusterConfig
+
+    # Create a test cluster config with GCS FT enabled
+    cluster_config = RayJobClusterConfig(
+        enable_gcs_ft=True,
+        redis_address="redis://redis-service:6379",
+        external_storage_namespace="storage-ns",
+    )
+
+    # Build the spec using the method on the cluster config
+    spec = cluster_config.build_ray_cluster_spec("test-cluster")
+
+    # Verify GCS fault tolerance options
+    assert "gcsFaultToleranceOptions" in spec
+    gcs_ft = spec["gcsFaultToleranceOptions"]
+    assert gcs_ft["redisAddress"] == "redis://redis-service:6379"
+    assert gcs_ft["externalStorageNamespace"] == "storage-ns"