From 9687db48c1d9be9f4675e7f585b960a26f7662d6 Mon Sep 17 00:00:00 2001 From: Weichen Xu Date: Wed, 20 Dec 2023 00:17:00 +0800 Subject: [PATCH 1/5] init Signed-off-by: Weichen Xu --- python/ray/util/spark/cluster_init.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/python/ray/util/spark/cluster_init.py b/python/ray/util/spark/cluster_init.py index ed1ae8324bbf..09501d445b43 100644 --- a/python/ray/util/spark/cluster_init.py +++ b/python/ray/util/spark/cluster_init.py @@ -1202,6 +1202,12 @@ def _setup_ray_cluster_internal( head_ip = cluster.address.split(":")[0] remote_connection_address = f"ray://{head_ip}:{cluster.ray_client_server_port}" + + if is_in_databricks_runtime(): + displayHTML( + "When you are using Ray on spark cluster, " + "you only pay for spark cluster usage." + ) return cluster.address, remote_connection_address From 4e9932ffbd3744a62a22893de288055bc4796fc9 Mon Sep 17 00:00:00 2001 From: Weichen Xu Date: Wed, 20 Dec 2023 09:12:20 +0800 Subject: [PATCH 2/5] address-comment Signed-off-by: Weichen Xu --- python/ray/util/spark/cluster_init.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/ray/util/spark/cluster_init.py b/python/ray/util/spark/cluster_init.py index 09501d445b43..31d9099676e4 100644 --- a/python/ray/util/spark/cluster_init.py +++ b/python/ray/util/spark/cluster_init.py @@ -1205,8 +1205,8 @@ def _setup_ray_cluster_internal( if is_in_databricks_runtime(): displayHTML( - "When you are using Ray on spark cluster, " - "you only pay for spark cluster usage." + "When you are using Ray on Spark cluster, " + "you only pay for Spark cluster usage." ) return cluster.address, remote_connection_address From 16d5f2722bd264e50f2353131081d9986d66e7c7 Mon Sep 17 00:00:00 2001 From: Weichen Xu Date: Wed, 20 Dec 2023 11:11:01 +0800 Subject: [PATCH 3/5] update Signed-off-by: Weichen Xu --- python/ray/util/spark/cluster_init.py | 4 ++-- python/ray/util/spark/databricks_hook.py | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/python/ray/util/spark/cluster_init.py b/python/ray/util/spark/cluster_init.py index 31d9099676e4..d4e8ae1aba89 100644 --- a/python/ray/util/spark/cluster_init.py +++ b/python/ray/util/spark/cluster_init.py @@ -34,7 +34,7 @@ _wait_service_up, ) from .start_hook_base import RayOnSparkStartHook -from .databricks_hook import DefaultDatabricksRayOnSparkStartHook +from .databricks_hook import DefaultDatabricksRayOnSparkStartHook, get_databricks_function from threading import Event @@ -1204,7 +1204,7 @@ def _setup_ray_cluster_internal( remote_connection_address = f"ray://{head_ip}:{cluster.ray_client_server_port}" if is_in_databricks_runtime(): - displayHTML( + get_databricks_function("displayHTML")( "When you are using Ray on Spark cluster, " "you only pay for Spark cluster usage." ) diff --git a/python/ray/util/spark/databricks_hook.py b/python/ray/util/spark/databricks_hook.py index 2d8526c411d1..cbb1401cb58b 100644 --- a/python/ray/util/spark/databricks_hook.py +++ b/python/ray/util/spark/databricks_hook.py @@ -9,6 +9,15 @@ _logger = logging.getLogger(__name__) +def get_databricks_function(func_name): + import IPython + + ip_shell = IPython.get_ipython() + if ip_shell is None: + raise RuntimeError("No IPython environment.") + return ip_shell.ns_table["user_global"][func_name] + + def get_db_entry_point(): """ Return databricks entry_point instance, it is for calling some From 47e0168a9516a04170b83f234261e76c40686877 Mon Sep 17 00:00:00 2001 From: Weichen Xu Date: Thu, 21 Dec 2023 17:24:10 +0800 Subject: [PATCH 4/5] update Signed-off-by: Weichen Xu --- python/ray/tests/spark/test_databricks_hook.py | 4 ++++ python/ray/util/spark/cluster_init.py | 7 +------ python/ray/util/spark/databricks_hook.py | 9 +++++++++ 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/python/ray/tests/spark/test_databricks_hook.py b/python/ray/tests/spark/test_databricks_hook.py index b95b788fe4ba..637cfb29736e 100644 --- a/python/ray/tests/spark/test_databricks_hook.py +++ b/python/ray/tests/spark/test_databricks_hook.py @@ -52,6 +52,10 @@ def test_hook(self, monkeypatch): monkeypatch.setattr( "ray.util.spark.databricks_hook.get_db_entry_point", lambda: db_api_entry ) + monkeypatch.setattr( + "ray.util.spark.databricks_hook.get_databricks_function", + lambda *args, **kwargs: None, + ) try: setup_ray_cluster( max_worker_nodes=2, diff --git a/python/ray/util/spark/cluster_init.py b/python/ray/util/spark/cluster_init.py index d4e8ae1aba89..2dab338a42f0 100644 --- a/python/ray/util/spark/cluster_init.py +++ b/python/ray/util/spark/cluster_init.py @@ -34,7 +34,7 @@ _wait_service_up, ) from .start_hook_base import RayOnSparkStartHook -from .databricks_hook import DefaultDatabricksRayOnSparkStartHook, get_databricks_function +from .databricks_hook import DefaultDatabricksRayOnSparkStartHook from threading import Event @@ -1203,11 +1203,6 @@ def _setup_ray_cluster_internal( head_ip = cluster.address.split(":")[0] remote_connection_address = f"ray://{head_ip}:{cluster.ray_client_server_port}" - if is_in_databricks_runtime(): - get_databricks_function("displayHTML")( - "When you are using Ray on Spark cluster, " - "you only pay for Spark cluster usage." - ) return cluster.address, remote_connection_address diff --git a/python/ray/util/spark/databricks_hook.py b/python/ray/util/spark/databricks_hook.py index cbb1401cb58b..2fb8789cd9f3 100644 --- a/python/ray/util/spark/databricks_hook.py +++ b/python/ray/util/spark/databricks_hook.py @@ -80,6 +80,15 @@ def on_ray_dashboard_created(self, port): def on_cluster_created(self, ray_cluster_handler): db_api_entry = get_db_entry_point() + + try: + get_databricks_function("displayHTML")( + "When you are using Ray on Spark " + "cluster, you only pay for Spark cluster usage." + ) + except Exception: + pass + if ray_cluster_handler.autoscale or self.is_global: # Disable auto shutdown if # 1) autoscaling enabled From 7994f90d4b4b768cc1bbb200b9ced1ea71ae0547 Mon Sep 17 00:00:00 2001 From: Weichen Xu Date: Thu, 21 Dec 2023 17:26:34 +0800 Subject: [PATCH 5/5] update Signed-off-by: Weichen Xu --- python/ray/util/spark/cluster_init.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/ray/util/spark/cluster_init.py b/python/ray/util/spark/cluster_init.py index 2dab338a42f0..ed1ae8324bbf 100644 --- a/python/ray/util/spark/cluster_init.py +++ b/python/ray/util/spark/cluster_init.py @@ -1202,7 +1202,6 @@ def _setup_ray_cluster_internal( head_ip = cluster.address.split(":")[0] remote_connection_address = f"ray://{head_ip}:{cluster.ray_client_server_port}" - return cluster.address, remote_connection_address