From 6f37c9b974a7a24df37749967bda0c86a544b4d8 Mon Sep 17 00:00:00 2001 From: Joyjit Chatterjee Date: Thu, 13 Aug 2026 21:19:17 +0000 Subject: [PATCH] feat: Add hyp create ray-dashboard-connection command --- .../cli/commands/ray_dashboard_connection.py | 110 ++++++++++ .../ray_dashboard_connection_constants.py | 16 ++ src/sagemaker/hyperpod/cli/hyp_cli.py | 2 + .../cli/test_ray_dashboard_connection.py | 189 ++++++++++++++++++ 4 files changed, 317 insertions(+) create mode 100644 src/sagemaker/hyperpod/cli/commands/ray_dashboard_connection.py create mode 100644 src/sagemaker/hyperpod/cli/constants/ray_dashboard_connection_constants.py create mode 100644 test/unit_tests/cli/test_ray_dashboard_connection.py diff --git a/src/sagemaker/hyperpod/cli/commands/ray_dashboard_connection.py b/src/sagemaker/hyperpod/cli/commands/ray_dashboard_connection.py new file mode 100644 index 00000000..0eedb2f2 --- /dev/null +++ b/src/sagemaker/hyperpod/cli/commands/ray_dashboard_connection.py @@ -0,0 +1,110 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. + +import click +from kubernetes import client, config +from kubernetes.client.rest import ApiException + +from sagemaker.hyperpod.cli.constants.ray_dashboard_connection_constants import ( + RAY_DASHBOARD_CONNECTION_GROUP, + RAY_DASHBOARD_CONNECTION_VERSION, + RAY_DASHBOARD_CONNECTION_PLURAL, +) +from sagemaker.hyperpod.common.telemetry.telemetry_logging import ( + _hyperpod_telemetry_emitter, +) +from sagemaker.hyperpod.common.telemetry.constants import Feature +from sagemaker.hyperpod.common.cli_decorators import handle_cli_exceptions + + +def _get_eks_api_client(): + """Load kubeconfig and create an authenticated API client. + + Works around a kubernetes-client issue where exec-based tokens + are not properly forwarded in the Authorization header. + """ + config.load_kube_config() + configuration = client.Configuration.get_default_copy() + + # Extract the token from the exec provider + token = None + if configuration.api_key and "authorization" in configuration.api_key: + token_value = configuration.api_key["authorization"] + prefix = "Bearer " + if token_value.startswith(prefix): + token = token_value.removeprefix(prefix) + else: + token = token_value + + # Clear api_key to avoid double-auth conflicts + configuration.api_key = {} + configuration.api_key_prefix = {} + + if token: + return client.ApiClient( + configuration, + header_name="Authorization", + header_value=f"Bearer {token}", + ) + return client.ApiClient(configuration) + + +@click.command("ray-dashboard-connection") +@click.option("--cluster-name", required=True, help="Name of the RayCluster") +@click.option("--namespace", "-n", required=False, default="default", help="Namespace of the RayCluster") +@_hyperpod_telemetry_emitter(Feature.HYPERPOD_CLI, "create_ray_dashboard_connection") +@handle_cli_exceptions() +def create_ray_dashboard_connection(cluster_name, namespace): + """Create a RayDashboardConnection to get a dashboard URL for a RayCluster.""" + api_client = _get_eks_api_client() + + body = { + "apiVersion": f"{RAY_DASHBOARD_CONNECTION_GROUP}/{RAY_DASHBOARD_CONNECTION_VERSION}", + "kind": "RayDashboardConnection", + "metadata": { + "namespace": namespace, + }, + "spec": { + "clusterName": cluster_name, + }, + } + + api = client.CustomObjectsApi(api_client) + + try: + result = api.create_namespaced_custom_object( + group=RAY_DASHBOARD_CONNECTION_GROUP, + version=RAY_DASHBOARD_CONNECTION_VERSION, + namespace=namespace, + plural=RAY_DASHBOARD_CONNECTION_PLURAL, + body=body, + ) + except ApiException as e: + if e.status == 404: + body_str = e.body or "" + if "raydashboardconnections" in body_str.lower() or RAY_DASHBOARD_CONNECTION_GROUP in body_str: + raise click.ClickException( + "The RayDashboardConnection API is not available on this cluster.\n" + "Please install the hyperpod-ray-endpoint-operator Helm chart.\n" + ) + raise click.ClickException(f"Not found: {body_str}") + raise + + connection_url = result.get("status", {}).get("connectionUrl", "") + if connection_url: + click.echo(connection_url) + else: + raise click.ClickException( + f"Failed to get dashboard URL for RayCluster '{cluster_name}' in namespace '{namespace}'.\n" + "Please contact your cluster administrator." + ) diff --git a/src/sagemaker/hyperpod/cli/constants/ray_dashboard_connection_constants.py b/src/sagemaker/hyperpod/cli/constants/ray_dashboard_connection_constants.py new file mode 100644 index 00000000..2524785b --- /dev/null +++ b/src/sagemaker/hyperpod/cli/constants/ray_dashboard_connection_constants.py @@ -0,0 +1,16 @@ +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. + +RAY_DASHBOARD_CONNECTION_GROUP = "connection.access.sagemaker.amazonaws.com" +RAY_DASHBOARD_CONNECTION_VERSION = "v1alpha1" +RAY_DASHBOARD_CONNECTION_PLURAL = "raydashboardconnections" diff --git a/src/sagemaker/hyperpod/cli/hyp_cli.py b/src/sagemaker/hyperpod/cli/hyp_cli.py index 4b2107c0..8de323cf 100644 --- a/src/sagemaker/hyperpod/cli/hyp_cli.py +++ b/src/sagemaker/hyperpod/cli/hyp_cli.py @@ -54,6 +54,7 @@ space_template_update, ) from sagemaker.hyperpod.cli.commands.space_access import space_access_create +from sagemaker.hyperpod.cli.commands.ray_dashboard_connection import create_ray_dashboard_connection from sagemaker.hyperpod.cli.commands.init import ( init, @@ -216,6 +217,7 @@ def exec(): create.add_command(space_create) create.add_command(space_template_create) create.add_command(space_access_create) +create.add_command(create_ray_dashboard_connection) list.add_command(list_jobs) recipe_list_cmd = copy.copy(list_jobs) diff --git a/test/unit_tests/cli/test_ray_dashboard_connection.py b/test/unit_tests/cli/test_ray_dashboard_connection.py new file mode 100644 index 00000000..2f595d89 --- /dev/null +++ b/test/unit_tests/cli/test_ray_dashboard_connection.py @@ -0,0 +1,189 @@ +import pytest +from click.testing import CliRunner +from unittest.mock import Mock, patch, MagicMock + +from kubernetes.client.rest import ApiException + +from sagemaker.hyperpod.cli.commands.ray_dashboard_connection import create_ray_dashboard_connection + + +class TestRayDashboardConnectionCommand: + """Test cases for ray-dashboard-connection command""" + + def setup_method(self): + self.runner = CliRunner() + + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_success_returns_url(self, mock_custom_objects_api_class, mock_get_client): + """Test successful creation returns the connection URL""" + mock_api = Mock() + mock_api.create_namespaced_custom_object.return_value = { + "status": { + "connectionUrl": "https://my-cluster.spaces.example.com/bearer-auth?token=abc123" + } + } + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + '--namespace', 'team-a', + ]) + + assert result.exit_code == 0 + assert "https://my-cluster.spaces.example.com/bearer-auth?token=abc123" in result.output + mock_api.create_namespaced_custom_object.assert_called_once_with( + group="connection.access.sagemaker.amazonaws.com", + version="v1alpha1", + namespace="team-a", + plural="raydashboardconnections", + body={ + "apiVersion": "connection.access.sagemaker.amazonaws.com/v1alpha1", + "kind": "RayDashboardConnection", + "metadata": {"namespace": "team-a"}, + "spec": {"clusterName": "my-raycluster"}, + }, + ) + + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_default_namespace(self, mock_custom_objects_api_class, mock_get_client): + """Test namespace defaults to 'default' when not specified""" + mock_api = Mock() + mock_api.create_namespaced_custom_object.return_value = { + "status": {"connectionUrl": "https://example.com/dashboard"} + } + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + ]) + + assert result.exit_code == 0 + assert "https://example.com/dashboard" in result.output + call_kwargs = mock_api.create_namespaced_custom_object.call_args[1] + assert call_kwargs["namespace"] == "default" + + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_empty_url_raises_error(self, mock_custom_objects_api_class, mock_get_client): + """Test that empty connectionUrl raises an error""" + mock_api = Mock() + mock_api.create_namespaced_custom_object.return_value = { + "status": {"connectionUrl": ""} + } + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + '--namespace', 'default', + ]) + + assert result.exit_code != 0 + assert "Failed to get dashboard URL" in result.output + assert "contact your cluster administrator" in result.output + + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_no_status_raises_error(self, mock_custom_objects_api_class, mock_get_client): + """Test that missing status raises an error""" + mock_api = Mock() + mock_api.create_namespaced_custom_object.return_value = { + "metadata": {"name": "generated-name"} + } + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + ]) + + assert result.exit_code != 0 + assert "Failed to get dashboard URL" in result.output + + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_404_api_not_installed(self, mock_custom_objects_api_class, mock_get_client): + """Test 404 when operator is not installed shows install instructions""" + mock_api = Mock() + mock_api.create_namespaced_custom_object.side_effect = ApiException( + status=404, + reason="Not Found", + http_resp=Mock( + status=404, + reason="Not Found", + data=b'{"message":"the server could not find the requested resource","details":{"group":"connection.access.sagemaker.amazonaws.com","kind":"raydashboardconnections"}}' + ), + ) + mock_api.create_namespaced_custom_object.side_effect.body = ( + '{"message":"the server could not find the requested resource",' + '"details":{"group":"connection.access.sagemaker.amazonaws.com","kind":"raydashboardconnections"}}' + ) + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + '--namespace', 'default', + ]) + + assert result.exit_code != 0 + assert "RayDashboardConnection API is not available" in result.output + assert "hyperpod-ray-endpoint-operator" in result.output + + @patch('sagemaker.hyperpod.common.cli_decorators._namespace_exists', return_value=True) + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_404_namespace_not_found(self, mock_custom_objects_api_class, mock_get_client, mock_ns_exists): + """Test 404 for missing namespace shows raw error""" + mock_api = Mock() + mock_api.create_namespaced_custom_object.side_effect = ApiException( + status=404, + reason="Not Found", + http_resp=Mock(status=404, reason="Not Found", data=b'{"message":"namespaces not-exists not found"}'), + ) + mock_api.create_namespaced_custom_object.side_effect.body = '{"message":"namespaces not-exists not found"}' + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + '--namespace', 'not-exists', + ]) + + assert result.exit_code != 0 + assert "not found" in result.output.lower() or "not-exists" in result.output + + @patch('sagemaker.hyperpod.common.cli_decorators._namespace_exists', return_value=True) + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection._get_eks_api_client') + @patch('sagemaker.hyperpod.cli.commands.ray_dashboard_connection.client.CustomObjectsApi') + def test_create_403_raises_exception(self, mock_custom_objects_api_class, mock_get_client, mock_ns_exists): + """Test 403 forbidden is propagated as an error""" + mock_api = Mock() + exc = ApiException( + status=403, + reason="Forbidden", + http_resp=Mock(status=403, reason="Forbidden", data=b'{"message":"forbidden"}'), + ) + exc.body = '{"message":"forbidden"}' + mock_api.create_namespaced_custom_object.side_effect = exc + mock_custom_objects_api_class.return_value = mock_api + mock_get_client.return_value = Mock() + + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--cluster-name', 'my-raycluster', + ]) + + assert result.exit_code != 0 + + def test_missing_cluster_name(self): + """Test that --cluster-name is required""" + result = self.runner.invoke(create_ray_dashboard_connection, [ + '--namespace', 'default', + ]) + + assert result.exit_code != 0 + assert "Missing option '--cluster-name'" in result.output