🐛 prevent creation of separate shared clusters via Kubernetes Lease API (#290)

danielgafni · web-flow · commit 5a5719b3791c · 2026-01-25T21:17:54.000+02:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - `RayCluster`'s head pod logs are now displayed on startup timeout or failure
+- Implemented Kubernetes Lease-based leader election for coordinating shared cluster creation. This replaces the previous flaky approach and guarantees that only one step creates the shared cluster when multiple parallel steps start simultaneously (within the same Dagster run).
 
 ### Fixes
 
diff --git a/src/dagster_ray/kuberay/client/base.py b/src/dagster_ray/kuberay/client/base.py
@@ -38,8 +38,8 @@ def __init__(
         self.version = version
         self.kind = kind
         self.plural = plural
-        self.api_client = api_client
-        self._api = client.CustomObjectsApi(api_client=api_client)
+        self.api_client = api_client or client.ApiClient()
+        self._api = client.CustomObjectsApi(api_client=self.api_client)
         self._core_v1_api = client.CoreV1Api(api_client=api_client)
 
     def wait_for_service_endpoints(
diff --git a/src/dagster_ray/kuberay/leader_election.py b/src/dagster_ray/kuberay/leader_election.py
@@ -0,0 +1,165 @@
+"""Leader election via Kubernetes Lease for coordinating shared cluster creation.
+
+This module provides a Kubernetes Lease-based leader election mechanism to ensure
+only one step creates a shared RayCluster when multiple parallel steps start simultaneously.
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from collections.abc import Callable
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from kubernetes.client import ApiClient
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class LeaderElectionResult:
+    """Result of a leader election attempt."""
+
+    is_leader: bool
+    lease_name: str
+    holder_identity: str
+
+
+class LeaderElection:
+    """Kubernetes Lease-based leader election for coordinating shared cluster creation.
+
+    Uses the `coordination.k8s.io/v1` `Lease` API to implement distributed leader election.
+    """
+
+    def __init__(
+        self,
+        holder_identity: str,
+        api_client: ApiClient,
+        lease_duration_seconds: int = 60,
+        retry_period_seconds: float = 1.0,
+        acquire_timeout_seconds: float = 60.0,
+    ):
+        """Initialize leader election.
+
+        Args:
+            holder_identity: Unique identifier for this process
+            api_client: Kubernetes API client
+            lease_duration_seconds: How long the lease is held before expiring
+            retry_period_seconds: How often to retry acquiring the lease
+            acquire_timeout_seconds: Total time to wait for leader election
+        """
+        from kubernetes import client
+
+        self.api_client = api_client
+        self._coordination_api = client.CoordinationV1Api(api_client=api_client)
+        self._holder_identity = holder_identity
+        self.lease_duration_seconds = lease_duration_seconds
+        self.retry_period_seconds = retry_period_seconds
+        self.acquire_timeout_seconds = acquire_timeout_seconds
+
+    def acquire_or_wait(
+        self,
+        lease_name: str,
+        namespace: str,
+        resource_exists_check: Callable[[], bool],
+    ) -> LeaderElectionResult:
+        """Acquire leadership or wait for the resource to be created by the leader.
+
+        Args:
+            lease_name: Name of the lease to use for coordination
+            namespace: Kubernetes namespace
+            resource_exists_check: Callable that returns True if the shared resource exists
+
+        Returns:
+            LeaderElectionResult indicating whether this process should create the resource
+        """
+        start_time = time.time()
+
+        while time.time() - start_time < self.acquire_timeout_seconds:
+            # Check if resource already exists
+            if resource_exists_check():
+                logger.debug("Shared resource already exists, not acquiring leadership")
+                # Try to get the current lease holder for logging purposes
+                current_holder = self._get_lease_holder(lease_name, namespace)
+                return LeaderElectionResult(
+                    is_leader=False,
+                    lease_name=lease_name,
+                    holder_identity=current_holder or "unknown",
+                )
+
+            # Try to become the leader
+            if self._try_create_lease(lease_name, namespace):
+                logger.info(f"Acquired leadership by creating lease {namespace}/{lease_name}")
+                return LeaderElectionResult(
+                    is_leader=True,
+                    lease_name=lease_name,
+                    holder_identity=self._holder_identity,
+                )
+
+            # Someone else is the leader, wait for resource to appear
+            current_holder = self._get_lease_holder(lease_name, namespace)
+            logger.debug(f"Lease {namespace}/{lease_name} held by {current_holder}, waiting for resource")
+            time.sleep(self.retry_period_seconds)
+
+        # Timeout - assume leadership to avoid deadlock
+        logger.warning(f"Leader election timed out after {self.acquire_timeout_seconds}s, assuming leadership")
+        return LeaderElectionResult(
+            is_leader=True,
+            lease_name=lease_name,
+            holder_identity=self._holder_identity,
+        )
+
+    def _try_create_lease(self, name: str, namespace: str) -> bool:
+        """Attempt to create a new lease. Returns True if successful (we're the leader)."""
+        from kubernetes.client import ApiException
+        from kubernetes.client.models import V1Lease, V1LeaseSpec, V1ObjectMeta
+
+        lease = V1Lease(
+            metadata=V1ObjectMeta(name=name, namespace=namespace),
+            spec=V1LeaseSpec(
+                holder_identity=self._holder_identity,
+                lease_duration_seconds=self.lease_duration_seconds,
+                acquire_time=datetime.now(timezone.utc),
+            ),
+        )
+
+        try:
+            self._coordination_api.create_namespaced_lease(namespace, lease)
+            return True
+        except ApiException as e:
+            if e.status == 409:  # Already exists - someone else is the leader
+                return False
+            raise
+
+    def _get_lease_holder(self, name: str, namespace: str) -> str | None:
+        """Get the current holder identity from an existing lease."""
+        from typing import cast
+
+        from kubernetes.client import ApiException
+        from kubernetes.client.models import V1Lease
+
+        try:
+            lease = cast(V1Lease, self._coordination_api.read_namespaced_lease(name=name, namespace=namespace))
+            return lease.spec.holder_identity if lease.spec else None
+        except ApiException as e:
+            if e.status == 404:
+                return None
+            raise
+
+    def _delete_lease(self, name: str, namespace: str) -> None:
+        """Delete a lease."""
+        from kubernetes.client import ApiException
+
+        try:
+            self._coordination_api.delete_namespaced_lease(name=name, namespace=namespace)
+        except ApiException as e:
+            if e.status != 404:
+                raise
+
+    def release(self, lease_name: str, namespace: str) -> None:
+        """Release leadership by deleting the lease."""
+        self._delete_lease(lease_name, namespace)
+        logger.debug(f"Released lease {namespace}/{lease_name}")
diff --git a/src/dagster_ray/kuberay/resources/raycluster.py b/src/dagster_ray/kuberay/resources/raycluster.py
@@ -10,6 +10,7 @@
 from dagster_ray.configs import Lifecycle
 from dagster_ray.kuberay.client import RayClusterClient
 from dagster_ray.kuberay.configs import ClusterSharing, RayClusterConfig
+from dagster_ray.kuberay.leader_election import LeaderElection
 from dagster_ray.kuberay.resources.base import BaseKubeRayResource
 from dagster_ray.kuberay.utils import normalize_k8s_label_values
 from dagster_ray.types import AnyDagsterContext
@@ -103,6 +104,52 @@ def get_image(self, context: AnyDagsterContext) -> str | None:
         assert context.dagster_run is not None
         return self.image or context.dagster_run.tags.get("dagster/image")
 
+    def _find_shared_cluster(self, label_selector: str) -> list[dict]:
+        """Find existing shared clusters matching the label selector."""
+        return self.client.list(
+            label_selector=label_selector,
+            namespace=self.namespace,
+        ).get("items", [])
+
+    def _use_existing_cluster(self, context: AnyDagsterContext, cluster_name: str) -> None:
+        """Mark an existing cluster as being used by this step."""
+        self._name = cluster_name
+        self._creation_verb = "Using"
+
+        # Place a lock on the cluster using JSON patch to avoid overwriting existing annotations
+        lock_annotations = self.get_sharing_lock_annotations(context)
+        patch_operations = [
+            {
+                "op": "add",
+                "path": f"/metadata/annotations/{key.replace('/', '~1')}",
+                "value": value,
+            }
+            for key, value in lock_annotations.items()
+        ]
+
+        self.client.update_json_patch(
+            name=cluster_name,
+            namespace=self.namespace,
+            body=patch_operations,
+        )
+
+    def _get_lease_name(self, context: AnyDagsterContext) -> str:
+        """Generate a deterministic lease name for cluster sharing coordination.
+
+        The lease name must be the same for all steps that should share a cluster,
+        so it's derived from the sharing label selector (which defines "same cluster" semantics).
+        The label selector is sorted to ensure deterministic ordering.
+        """
+        import hashlib
+
+        label_selector = self.get_sharing_label_selector(context)
+        # Sort the label selector components for deterministic ordering
+        # Label selector format is "key1=value1,key2=value2,..."
+        sorted_selector = ",".join(sorted(label_selector.split(",")))
+        # Create a short hash of the sorted label selector for the lease name
+        selector_hash = hashlib.sha256(sorted_selector.encode()).hexdigest()[:12]
+        return f"dagster-ray-{selector_hash}"
+
     @override
     def create(self, context: AnyDagsterContext):
         assert context.log is not None
@@ -116,47 +163,65 @@ def create(self, context: AnyDagsterContext):
             context.log.info(
                 f"RayCluster sharing is enabled. Looking for clusters matching label selector: {label_selector}"
             )
-            # check whether a cluster matching the sharing config already exists
-            matching_clusters = self.client.list(
-                label_selector=label_selector,
-                namespace=self.namespace,
-            ).get("items", [])
+
+            # Check if a shared cluster already exists
+            matching_clusters = self._find_shared_cluster(label_selector)
             if matching_clusters:
                 cluster_name = matching_clusters[0]["metadata"]["name"]
                 context.log.info(
-                    f"Found {len(matching_clusters)} clusters matching the label selector. Using the first one: {cluster_name}"
-                )
-                self._name = cluster_name
-                self._creation_verb = "Using"
-
-                # place a lock on the cluster
-                # Create individual patch operations for each annotation to avoid replacing existing ones
-                lock_annotations = self.get_sharing_lock_annotations(context)
-                patch_operations = [
-                    {
-                        "op": "add",
-                        "path": f"/metadata/annotations/{key.replace('/', '~1')}",
-                        "value": value,
-                    }
-                    for key, value in lock_annotations.items()
-                ]
-
-                self.client.update_json_patch(
-                    name=cluster_name,
-                    namespace=self.namespace,
-                    body=patch_operations,
+                    f"Found {len(matching_clusters)} clusters matching the label selector. "
+                    f"Using the first one: {cluster_name}"
                 )
-
+                self._use_existing_cluster(context, cluster_name)
                 return
-            else:
-                context.log.info("No matching clusters found. Creating a new one.")
 
-                # mark the cluster as being used by this step
-                annotations.update(self.get_sharing_lock_annotations(context))
+            # No cluster exists - use leader election to coordinate creation
+            context.log.info("No matching shared clusters found. Using leader election to coordinate cluster creation.")
+
+            lease_name = self._get_lease_name(context)
+            holder_identity = f"{context.run_id}-{self.resource_uid}"
+            leader_election = LeaderElection(
+                holder_identity=holder_identity,
+                api_client=self.client.api_client,
+                lease_duration_seconds=10,
+                retry_period_seconds=2.0,
+                acquire_timeout_seconds=60.0,
+            )
+
+            def cluster_exists() -> bool:
+                return bool(self._find_shared_cluster(label_selector))
+
+            result = leader_election.acquire_or_wait(
+                lease_name=lease_name,
+                namespace=self.namespace,
+                resource_exists_check=cluster_exists,
+            )
+
+            if not result.is_leader:
+                # Another step created the cluster while we were waiting
+                matching_clusters = self._find_shared_cluster(label_selector)
+                if matching_clusters:
+                    cluster_name = matching_clusters[0]["metadata"]["name"]
+                    context.log.info(f"Using cluster created by leader: {cluster_name}")
+                    self._use_existing_cluster(context, cluster_name)
+                    return
+                else:
+                    # Edge case: leader created then deleted? Fall through to create
+                    context.log.warning(
+                        "Leader election indicated not leader, but no cluster found. Proceeding to create cluster."
+                    )
+
+            context.log.info("Won leader election - this step will create the shared cluster.")
+
+            # Mark the cluster as being used by this step
+            annotations.update(self.get_sharing_lock_annotations(context))
+
+            # Release the lease after we're done creating (cluster existence is the coordination point now)
+            # We'll release after successful creation below
 
         self._name = self.ray_cluster.metadata.get("name") or self._get_step_name(context)
 
-        # just a safety measure, no need to recreate the cluster for step retries or smth
+        # Safety measure: don't recreate the cluster for step retries
         if not self.client.get(
             name=self.name,
             namespace=self.namespace,
@@ -175,6 +240,20 @@ def create(self, context: AnyDagsterContext):
             if not resource:
                 raise RuntimeError(f"Couldn't create {self.display_name}")
 
+            # Release the lease now that the cluster is created
+            if self.cluster_sharing.enabled:
+                try:
+                    lease_name = self._get_lease_name(context)
+                    holder_identity = f"{context.run_id}-{self.resource_uid}"
+                    leader_election = LeaderElection(
+                        holder_identity=holder_identity,
+                        api_client=self.client.api_client,
+                    )
+                    leader_election.release(lease_name, self.namespace)
+                except Exception as e:
+                    # Non-fatal: lease will expire on its own
+                    context.log.debug(f"Failed to release leader election lease: {e}")
+
     @override
     def wait(self, context: AnyDagsterContext):
         assert context.log is not None
