anyscale 0.26.67__py3-none-any.whl → 0.26.69__py3-none-any.whl

anyscale/commands/setup_k8s.py

@@ -42,11 +42,7 @@ class ClusterInfo:
     region: str
     cluster_name: str
     project_id: Optional[str] = None
-    cluster_arn: Optional[str] = None
     oidc_provider: Optional[str] = None
-    cluster_location: Optional[str] = None
-    workload_identity_pool: Optional[str] = None
-    cluster_version: Optional[str] = None
 
 
 @dataclass
@@ -79,6 +75,7 @@ class KubernetesCloudSetupCommand:
         functional_verify: bool,
         yes: bool,
         values_file: Optional[str] = None,
+        operator_chart: Optional[str] = None,
     ) -> None:
         """
         Main entry point for Kubernetes cloud setup.
@@ -93,6 +90,7 @@ class KubernetesCloudSetupCommand:
             functional_verify: Whether to run functional verification
             yes: Skip confirmation prompts
             values_file: Optional custom path for Helm values file
+            operator_chart: Optional path to operator chart (skips helm repo add/update)
         """
         self.log.open_block(
             "Setup", f"Setting up Kubernetes cloud '{name}' on {provider.upper()}"
@@ -134,6 +132,7 @@ class KubernetesCloudSetupCommand:
                 final_namespace,
                 infrastructure,
                 values_file,
+                operator_chart,
             )
 
             # Step 6: Verify installation
@@ -194,7 +193,7 @@ class KubernetesCloudSetupCommand:
         namespace: str,
         provider: str,
         region: str,
-        project_id: Optional[str],  # noqa: ARG002
+        project_id: Optional[str],
     ) -> ClusterInfo:
         """Discover and validate the target Kubernetes cluster using cloud provider APIs."""
         self.log.info(
@@ -205,8 +204,12 @@ class KubernetesCloudSetupCommand:
         if provider == "aws":
             return self._discover_aws_cluster(cluster_name, namespace, region)
         elif provider == "gcp":
-            raise click.ClickException(
-                "GCP support is not yet implemented. Please use AWS for now."
+            if not project_id:
+                raise click.ClickException(
+                    "GCP project ID is required. Please provide --project-id"
+                )
+            return self._discover_gcp_cluster(
+                cluster_name, namespace, region, project_id
             )
         else:
             raise click.ClickException(f"Unsupported provider: {provider}")
@@ -215,18 +218,6 @@ class KubernetesCloudSetupCommand:
         self, cluster_name: str, namespace: str, region: str
     ) -> ClusterInfo:
         """Discover AWS EKS cluster details and configure kubeconfig."""
-        try:
-            self._debug(f"Fetching EKS cluster info for {cluster_name} in {region}...")
-            cluster_info = self._get_eks_cluster_info(cluster_name, region)
-            self._debug(f"EKS Cluster ARN: {cluster_info.get('arn', 'Unknown')}")
-            self._debug(
-                f"EKS Cluster Version: {cluster_info.get('version', 'Unknown')}"
-            )
-        except Exception as e:  # noqa: BLE001
-            self.log.error(f"Failed to get EKS cluster info: {e}")
-            raise click.ClickException(
-                f"Failed to discover EKS cluster {cluster_name}: {e}"
-            )
 
         try:
             self._debug("Fetching OIDC provider information...")
@@ -262,9 +253,39 @@ class KubernetesCloudSetupCommand:
             provider="aws",
             region=region,
             cluster_name=cluster_name,
-            cluster_arn=cluster_info.get("arn"),
             oidc_provider=oidc_provider,
-            cluster_version=cluster_info.get("version"),
+        )
+
+    def _discover_gcp_cluster(
+        self, cluster_name: str, namespace: str, region: str, project_id: str
+    ) -> ClusterInfo:
+        """Discover GCP GKE cluster details and configure kubeconfig."""
+
+        try:
+            self._debug("Configuring kubeconfig for GKE cluster...")
+            self._configure_gcp_kubeconfig(cluster_name, region, project_id)
+        except Exception as e:  # noqa: BLE001
+            self.log.error(f"Failed to configure kubeconfig: {e}")
+            raise click.ClickException(
+                f"Failed to configure kubeconfig for GKE cluster: {e}"
+            )
+
+        try:
+            self._debug("Verifying kubeconfig configuration...")
+            self._verify_kubeconfig()
+            current_context = self._get_current_kubectl_context()
+            self.log.info(f"Cluster discovered: {current_context}", block_label="Setup")
+        except Exception as e:  # noqa: BLE001
+            self.log.error(f"Failed to verify kubeconfig: {e}")
+            raise click.ClickException(f"Failed to verify kubeconfig: {e}")
+
+        return ClusterInfo(
+            context=current_context,
+            namespace=namespace,
+            provider="gcp",
+            region=region,
+            cluster_name=cluster_name,
+            project_id=project_id,
         )
 
     def _setup_infrastructure(
@@ -278,9 +299,7 @@ class KubernetesCloudSetupCommand:
         if provider == "aws":
             return self._setup_aws_infrastructure(region, name, cluster_info)
         elif provider == "gcp":
-            raise click.ClickException(
-                "GCP support is not yet implemented. Please use AWS for now."
-            )
+            return self._setup_gcp_infrastructure(region, name, cluster_info)
         else:
             raise click.ClickException(f"Unsupported provider: {provider}")
 
@@ -502,6 +521,344 @@ class KubernetesCloudSetupCommand:
 
         return json.dumps(template, indent=2)
 
+    def _setup_gcp_infrastructure(  # noqa: PLR0912
+        self, region: str, name: str, cluster_info: ClusterInfo,
+    ) -> InfrastructureResources:
+        """Set up GCP infrastructure for Kubernetes using GCP Python SDK.
+
+        Note: Deployment Manager is deprecated so it is unused here.
+        Infrastructure Manager was tried but did not work well, so we rely
+        on the GCP Python SDK instead.
+        """
+        try:
+            from anyscale.utils.gcp_utils import get_google_cloud_client_factory
+        except ImportError as e:
+            self.log.error(f"Failed to import required modules: {e}")
+            raise click.ClickException(f"Failed to import required modules: {e}")
+
+        try:
+            # Generate a unique cloud ID
+            cloud_id = f"k8s-{name}-{os.urandom(4).hex()}"
+            deployment_name = cloud_id.replace("_", "-").lower()
+            self._debug(f"Generated cloud ID: {cloud_id}")
+            self._debug(f"Infrastructure Manager deployment name: {deployment_name}")
+        except Exception as e:  # noqa: BLE001
+            self.log.error(f"Failed to generate cloud ID: {e}")
+            raise click.ClickException(f"Failed to generate cloud ID: {e}")
+
+        try:
+            # Get Google Cloud client factory
+            factory = get_google_cloud_client_factory(self.log, cluster_info.project_id)
+        except Exception as e:  # noqa: BLE001
+            self.log.error(f"Failed to initialize GCP client: {e}")
+            raise click.ClickException(f"Failed to initialize GCP client: {e}")
+
+        try:
+            with self.log.indent():
+                self.log.warning(
+                    "NOTE: GCP resources (bucket and service account) created by this command are not managed by Anyscale.",
+                )
+                self.log.warning(
+                    "You will need to manually delete these resources when the cloud is no longer needed.",
+                )
+                self.log.info(
+                    "Creating GCP resources (bucket, service account, IAM bindings)...",
+                )
+
+                # Calculate resource names
+                # Service account name: anyscale-operator-<random 8 chars>
+                # Max length for GCP service account is 30 characters
+                random_suffix = os.urandom(4).hex()  # 8 hex chars
+                anyscale_service_account_name = f"anyscale-operator-{random_suffix}"
+                bucket_name = f"anyscale-{cloud_id.replace('_', '-').lower()}"
+
+                # Create GCS bucket
+                self._debug(f"Creating GCS bucket: {bucket_name}")
+                storage_client = factory.storage.Client()
+                bucket = storage_client.bucket(bucket_name)
+                bucket.location = region
+                bucket.storage_class = "REGIONAL"
+                bucket.iam_configuration.uniform_bucket_level_access_enabled = True
+                bucket.iam_configuration.public_access_prevention = "enforced"
+                bucket.versioning_enabled = True
+                bucket.labels = {"anyscale-cloud-id": cloud_id.replace("-", "_")}
+
+                # Set CORS
+                # Use ANYSCALE_CORS_ORIGIN from shared config
+                # This respects the ANYSCALE_HOST environment variable
+                allowed_origin = ANYSCALE_CORS_ORIGIN
+                bucket.cors = [
+                    {
+                        "origin": [allowed_origin],
+                        "responseHeader": ["*"],
+                        "method": ["GET", "PUT", "POST", "HEAD", "DELETE"],
+                        "maxAgeSeconds": 3600,
+                    }
+                ]
+
+                storage_client.create_bucket(bucket, location=region)
+                self.log.info(f"Created GCS bucket: {bucket_name}", block_label="Setup")
+
+                # Create service account
+                self._debug(
+                    f"Creating service account: {anyscale_service_account_name}"
+                )
+                iam_client = factory.build("iam", "v1")
+                service_account_body = {
+                    "accountId": anyscale_service_account_name,
+                    "serviceAccount": {
+                        "displayName": f"{cloud_id} Anyscale operator service account",
+                        "description": "Service account for Anyscale Kubernetes operator",
+                    },
+                }
+
+                service_account = (
+                    iam_client.projects()
+                    .serviceAccounts()
+                    .create(
+                        name=f"projects/{cluster_info.project_id}",
+                        body=service_account_body,
+                    )
+                    .execute()
+                )
+
+                service_account_email = service_account["email"]
+                self.log.info(
+                    f"Created service account: {service_account_email}",
+                    block_label="Setup",
+                )
+
+                # Wait for service account to propagate through GCP systems
+                import time
+
+                self._debug("Waiting 10 seconds for service account to propagate...")
+                time.sleep(10)
+
+                # Grant Workload Identity binding
+                self._debug("Setting up Workload Identity binding")
+
+                # The K8s service account needs:
+                # 1. workloadIdentityUser role - to impersonate the GCP service account
+                # 2. serviceAccountTokenCreator - to generate tokens (for getOpenIdToken)
+
+                policy_body = {
+                    "policy": {
+                        "bindings": [
+                            {
+                                "role": "roles/iam.workloadIdentityUser",
+                                "members": [
+                                    f"serviceAccount:{cluster_info.project_id}.svc.id.goog[{cluster_info.namespace}/anyscale-operator]"
+                                ],
+                            },
+                            {
+                                "role": "roles/iam.serviceAccountTokenCreator",
+                                "members": [f"serviceAccount:{service_account_email}"],
+                            },
+                        ]
+                    }
+                }
+
+                iam_client.projects().serviceAccounts().setIamPolicy(
+                    resource=f"projects/{cluster_info.project_id}/serviceAccounts/{service_account_email}",
+                    body=policy_body,
+                ).execute()
+
+                self.log.info(
+                    "Configured Workload Identity binding", block_label="Setup"
+                )
+
+                # Grant storage admin role to service account for the bucket
+                # Note: There's often a propagation delay after service account creation
+                # We need to retry with exponential backoff
+                self._debug("Granting storage permissions")
+
+                import time
+
+                max_retries = 5
+                retry_delay = 2  # Start with 2 seconds
+
+                for attempt in range(max_retries):
+                    try:
+                        bucket_policy = bucket.get_iam_policy(
+                            requested_policy_version=3
+                        )
+                        bucket_policy.bindings.append(
+                            {
+                                "role": "roles/storage.admin",
+                                "members": {f"serviceAccount:{service_account_email}"},
+                            }
+                        )
+                        bucket.set_iam_policy(bucket_policy)
+                        break  # Success!
+                    except Exception as e:  # noqa: BLE001
+                        if "does not exist" in str(e) and attempt < max_retries - 1:
+                            self._debug(
+                                f"Service account not yet propagated, retrying in {retry_delay}s... (attempt {attempt + 1}/{max_retries})"
+                            )
+                            time.sleep(retry_delay)
+                            retry_delay *= 2  # Exponential backoff
+                        else:
+                            raise  # Re-raise if it's not a propagation issue or we're out of retries
+
+                self.log.info(
+                    "Granted storage permissions to service account",
+                    block_label="Setup",
+                )
+
+            self.log.info("GCP resources created successfully", block_label="Setup")
+            self.log.warning(
+                f"REMINDER: To clean up when no longer needed, delete GCS bucket '{bucket_name}' and service account '{service_account_email}'"
+            )
+        except Exception as e:  # noqa: BLE001
+            self.log.error(f"Failed to create GCP resources: {e}")
+            raise click.ClickException(f"Failed to create GCP resources: {e}")
+
+        # Resources were created in the try block above
+        # bucket_name and service_account_email are already set
+        self._debug(f"GCS Bucket: {bucket_name}")
+        self._debug(f"Service Account Email: {service_account_email}")
+
+        return InfrastructureResources(
+            bucket_name=bucket_name,
+            iam_role_arn=service_account_email,  # For GCP, we use service account email
+            region=region,
+            project_id=cluster_info.project_id,
+        )
+
+    def _get_gke_cluster_info(
+        self, cluster_name: str, region: str, project_id: str
+    ) -> Dict[str, Any]:
+        """Get GKE cluster information using gcloud CLI."""
+        try:
+            # Try regional cluster first
+            result = subprocess.run(
+                [
+                    "gcloud",
+                    "container",
+                    "clusters",
+                    "describe",
+                    cluster_name,
+                    f"--region={region}",
+                    f"--project={project_id}",
+                    "--format=json",
+                ],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+            if result.returncode == 0:
+                return json.loads(result.stdout)
+
+            # Try zonal cluster
+            # Assuming zone 'a' if regional fails
+            zone = f"{region}-a"
+            result = subprocess.run(
+                [
+                    "gcloud",
+                    "container",
+                    "clusters",
+                    "describe",
+                    cluster_name,
+                    f"--zone={zone}",
+                    f"--project={project_id}",
+                    "--format=json",
+                ],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            return json.loads(result.stdout)
+        except subprocess.CalledProcessError as e:
+            raise click.ClickException(f"Failed to get GKE cluster info: {e.stderr}")
+        except json.JSONDecodeError as e:
+            raise click.ClickException(f"Failed to parse GKE cluster info: {e}")
+
+    def _get_gke_zones(
+        self, cluster_name: str, region: str, project_id: str
+    ) -> List[str]:
+        """Get zones where the GKE cluster's node pools are located."""
+        try:
+            cluster_info = self._get_gke_cluster_info(cluster_name, region, project_id)
+
+            # Extract zones from node pools
+            zones = []
+            node_pools = cluster_info.get("nodePools", [])
+
+            for pool in node_pools:
+                # For zonal clusters, each node pool has locations
+                pool_locations = pool.get("locations", [])
+                zones.extend(pool_locations)
+
+            # If no zones found from node pools, try cluster-level locations
+            if not zones:
+                cluster_locations = cluster_info.get("locations", [])
+                if cluster_locations:
+                    zones = cluster_locations
+
+            # Remove duplicates and sort
+            if zones:
+                unique_zones = sorted(set(zones))
+                self._debug(f"Discovered zones: {', '.join(unique_zones)}")
+                return unique_zones
+            else:
+                # Fallback to default zones
+                self._debug(
+                    "No zones found in cluster info, falling back to default zones"
+                )
+                return [region + "-a", region + "-b", region + "-c"]
+
+        except Exception as e:  # noqa: BLE001
+            self._debug(f"Failed to get zones: {e}, using default zones")
+            return [region + "-a", region + "-b", region + "-c"]
+
+    def _configure_gcp_kubeconfig(
+        self, cluster_name: str, region: str, project_id: str
+    ) -> None:
+        """Configure kubeconfig for GCP GKE cluster."""
+        self.log.info(f"Configuring kubeconfig for GKE cluster: {cluster_name}")
+
+        try:
+            # Try regional cluster first
+            result = subprocess.run(
+                [
+                    "gcloud",
+                    "container",
+                    "clusters",
+                    "get-credentials",
+                    cluster_name,
+                    f"--region={region}",
+                    f"--project={project_id}",
+                ],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+            if result.returncode == 0:
+                self.log.info("GKE kubeconfig configured successfully")
+                return
+
+            # Try zonal cluster
+            zone = f"{region}-a"
+            subprocess.run(
+                [
+                    "gcloud",
+                    "container",
+                    "clusters",
+                    "get-credentials",
+                    cluster_name,
+                    f"--zone={zone}",
+                    f"--project={project_id}",
+                ],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            self.log.info("GKE kubeconfig configured successfully")
+        except subprocess.CalledProcessError as e:
+            raise click.ClickException(
+                f"Failed to configure GKE kubeconfig: {e.stderr}"
+            )
+
     def _get_eks_cluster_info(self, cluster_name: str, region: str) -> Dict[str, Any]:
         """Get EKS cluster information using AWS CLI."""
         try:
@@ -698,11 +1055,33 @@ class KubernetesCloudSetupCommand:
                     zones=zones,
                 ),
             )
-        else:
-            raise click.ClickException(
-                "GCP support is not yet implemented. Please use AWS for now."
+        elif provider == "gcp":
+            assert infrastructure.project_id, "Project ID is required for GCP"
+
+            from anyscale.client.openapi_client.models import GCPConfig
+
+            # Dynamically determine zones from the GKE cluster
+            zones = self._get_gke_zones(
+                cluster_info.cluster_name, region, infrastructure.project_id
             )
 
+            cloud_deployment = CloudDeployment(
+                name=name,
+                provider=CloudProviders.GCP,
+                region=region,
+                compute_stack=ComputeStack.K8S,
+                object_storage=ObjectStorage(
+                    bucket_name=infrastructure.bucket_name, region=region
+                ),
+                gcp_config=GCPConfig(project_id=infrastructure.project_id,),
+                kubernetes_config=OpenAPIKubernetesConfig(
+                    anyscale_operator_iam_identity=infrastructure.iam_role_arn,
+                    zones=zones,
+                ),
+            )
+        else:
+            raise click.ClickException(f"Unsupported provider: {provider}")
+
         # Register the cloud
         try:
             self._debug("Cloud deployment details:")
@@ -735,10 +1114,18 @@ class KubernetesCloudSetupCommand:
                         skip_verifications=True,
                         auto_add_user=True,
                     )
-                else:
-                    raise click.ClickException(
-                        "GCP support is not yet implemented. Please use AWS for now."
+                elif provider == "gcp":
+                    self.log.info("Calling register_gcp_cloud...")
+                    self.cloud_controller.register_gcp_cloud(
+                        name=name,
+                        cloud_resource=cloud_deployment,
+                        functional_verify=None,
+                        yes=True,
+                        skip_verifications=True,
+                        auto_add_user=True,
                     )
+                else:
+                    raise click.ClickException(f"Unsupported provider: {provider}")
             finally:
                 # Restore the original log.info method
                 self.cloud_controller.log.info = original_log_info
@@ -781,6 +1168,7 @@ class KubernetesCloudSetupCommand:
         namespace: str,
         infrastructure: InfrastructureResources,
         values_file: Optional[str] = None,
+        operator_chart: Optional[str] = None,
     ) -> None:
         """Install the Anyscale operator using Helm."""
         self.log.info("Installing Anyscale operator...", block_label="Setup")
@@ -817,10 +1205,21 @@ class KubernetesCloudSetupCommand:
             additional_values=set_string_values,
         )
 
+        # Determine chart reference based on operator_chart parameter
+        if operator_chart:
+            # Use the provided chart path directly
+            self._debug(f"Using operator chart from: {operator_chart}")
+            chart_reference = operator_chart
+        else:
+            # Add Helm repo before installing
+            self._debug("Adding Anyscale Helm repository...")
+            self._add_helm_repo()
+            chart_reference = "anyscale/anyscale-operator"
+
         # Build a simple Helm command that only uses the values file
         self._debug("Generating Helm command...")
         helm_command = (
-            f"helm upgrade {release_name} anyscale/anyscale-operator "
+            f"helm upgrade {release_name} {chart_reference} "
             f"--values {values_file_path} "
             f"--namespace {namespace} "
             f"--create-namespace "
@@ -830,6 +1229,41 @@ class KubernetesCloudSetupCommand:
 
         self._execute_helm_command(helm_command)
 
+    def _add_helm_repo(self) -> None:
+        """Add and update the Anyscale Helm repository."""
+        try:
+            # Add the Anyscale Helm repository
+            self.log.info("Adding Anyscale Helm repository...", block_label="Setup")
+            subprocess.run(
+                [
+                    "helm",
+                    "repo",
+                    "add",
+                    "anyscale",
+                    "https://anyscale.github.io/helm-charts",
+                ],
+                capture_output=True,
+                text=True,
+                check=False,  # Don't fail if repo already exists
+            )
+
+            # Update the Helm repository
+            self.log.info("Updating Helm repositories...", block_label="Setup")
+            subprocess.run(
+                ["helm", "repo", "update", "anyscale"],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            self.log.info(
+                "Helm repository configured successfully", block_label="Setup"
+            )
+        except subprocess.CalledProcessError as e:
+            self.log.error(f"Failed to configure Helm repository: {e.stderr}")
+            raise click.ClickException(
+                f"Failed to configure Helm repository: {e.stderr}"
+            )
+
     def _extract_set_string_values(self, helm_command: str) -> Dict[str, str]:
         """
         Extract all --set-string key=value pairs from a Helm command.
@@ -853,6 +1287,35 @@ class KubernetesCloudSetupCommand:
 
         return set_string_values
 
+    def _set_nested_value(self, d: Dict[str, Any], key_path: str, value: Any) -> None:
+        """
+        Set a value in a nested dictionary using a dotted key path.
+
+        Args:
+            d: The dictionary to modify
+            key_path: Dotted key path (e.g., "workloads.serviceaccount.name")
+            value: The value to set
+
+        Example:
+            _set_nested_value({}, "workloads.serviceaccount.name", "my-sa")
+            # Results in: {"workloads": {"serviceaccount": {"name": "my-sa"}}}
+        """
+        keys = key_path.split(".")
+        current = d
+
+        # Navigate/create the nested structure
+        for key in keys[:-1]:
+            if key not in current:
+                current[key] = {}
+            elif not isinstance(current[key], dict):
+                # If the key exists but isn't a dict, we have a conflict
+                # In this case, we'll overwrite it with a dict
+                current[key] = {}
+            current = current[key]
+
+        # Set the final value
+        current[keys[-1]] = value
+
     def _prompt_for_namespace(
         self, default_namespace: str, skip_confirmation: bool = False
     ) -> str:
@@ -864,10 +1327,12 @@ class KubernetesCloudSetupCommand:
             return final_namespace
 
         self.log.info("Configuring Kubernetes namespace...")
-
         self.log.info(
-            f"Enter the namespace to use for the Anyscale operator (default: {final_namespace}):"
+            f"Specify the namespace to use for the Anyscale operator (leave blank for default: {final_namespace})."
         )
+        self.log.info("If the namespace does not exist, it will be created.")
+        self.log.info("Enter your namespace:")
+
         final_namespace = click.prompt("", default=final_namespace, show_default=True)
 
         # Validate namespace (Kubernetes DNS-1123 label requirements)
@@ -900,24 +1365,32 @@ class KubernetesCloudSetupCommand:
         """Generate Helm values file and save it locally."""
         self.log.info("Generating Helm values file...")
 
-        # Create values dictionary starting with base values
-        values: Dict[str, Any] = {
-            "cloudProvider": provider,
-            "cloudDeploymentId": cloud_deployment_id,
-            "region": region,
-            "operatorIamIdentity": infrastructure.iam_role_arn,
-            "ingress-nginx": {"enabled": True},
-        }
+        # Start with an empty dictionary to build up values
+        values: Dict[str, Any] = {}
 
+        # First, parse and merge additional_values with nested keys
         if additional_values:
             for key, value in additional_values.items():
-                if key not in values:
-                    values[key] = value
+                self._set_nested_value(values, key, value)
+
+        # Now overlay our constants on top (these take precedence)
+        # Use _set_nested_value to ensure proper nesting
+        self._set_nested_value(values, "global.cloudDeploymentId", cloud_deployment_id)
+        self._set_nested_value(values, "global.cloudProvider", provider)
+        self._set_nested_value(
+            values, "global.auth.iamIdentity", infrastructure.iam_role_arn
+        )
+        self._set_nested_value(values, "ingress-nginx.enabled", True)
+
+        # Add region for AWS only (using global.aws.region)
+        # Region field is deprecated for other providers
+        if provider == "aws":
+            self._set_nested_value(values, "global.aws.region", region)
 
         # Add control plane URL from ANYSCALE_HOST environment variable
         if ANYSCALE_HOST:
-            values["controlPlaneURL"] = ANYSCALE_HOST
-            self._debug(f"Using control plane URL: {ANYSCALE_HOST}")
+            self._set_nested_value(values, "global.controlPlaneURL", ANYSCALE_HOST)
+            self.log.info(f"Using control plane URL: {ANYSCALE_HOST}")
 
         if custom_path:
             values_file_path = custom_path
@@ -983,11 +1456,6 @@ class KubernetesCloudSetupCommand:
             operator_namespace=namespace,
         )
 
-        # Sleep to avoid race condition where operator has not loaded its IAM identity
-        import time
-
-        time.sleep(5)
-
         # Run verification
         success = verifier.verify(cloud_deployment)
 
@@ -1009,6 +1477,7 @@ def setup_kubernetes_cloud(  # noqa: PLR0913
     yes: bool = False,
     values_file: Optional[str] = None,
     debug: bool = False,
+    operator_chart: Optional[str] = None,
 ) -> None:
     """
     Set up Anyscale on a Kubernetes cluster.
@@ -1027,6 +1496,7 @@ def setup_kubernetes_cloud(  # noqa: PLR0913
         yes: Skip confirmation prompts
         values_file: Optional path for Helm values file
         debug: Enable debug logging
+        operator_chart: Optional path to operator chart (skips helm repo add/update)
     """
     cmd = KubernetesCloudSetupCommand(debug=debug)
 
@@ -1041,6 +1511,7 @@ def setup_kubernetes_cloud(  # noqa: PLR0913
             functional_verify=functional_verify,
             yes=yes,
             values_file=values_file,
+            operator_chart=operator_chart,
         )
     except Exception as e:  # noqa: BLE001
         click.echo(f"Setup failed: {e}", err=True)