Avoid shutting down Ray driver when running on a Ray worker (PrefectH…

…Q#15552)

src/integrations/prefect-ray/prefect_ray/task_runners.py

@@ -401,7 +401,7 @@ def __enter__(self):
                 "Using existing local instance."
             )
             return self
-        if self.address and self.address != "auto":
+        elif self.address and self.address != "auto":
             self.logger.info(
                 f"Connecting to an existing Ray instance at {self.address}"
             )
@@ -427,8 +427,12 @@ def __enter__(self):
 
     def __exit__(self, *exc_info):
         """
-        Shuts down the cluster.
+        Shuts down the driver/cluster.
         """
-        self.logger.debug("Shutting down Ray cluster...")
-        ray.shutdown()
+        # Check if we are running on the driver. Calling ray.shutdown() when running on a
+        # worker will crash the worker.
+        if ray.get_runtime_context().worker.mode == 0:
+            # Running on the driver. Will shutdown cluster if started by this task runner.
+            self.logger.debug("Shutting down Ray driver...")
+            ray.shutdown()
         super().__exit__(*exc_info)

src/integrations/prefect-ray/tests/test_task_runners.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import random
 import subprocess
 import sys
 import time
@@ -508,3 +509,32 @@ def test_flow():
         test_flow()
 
         assert "A future was garbage collected before it resolved" not in caplog.text
+
+    def test_can_run_many_tasks_without_crashing(self, task_runner):
+        """
+        Regression test for https://github.com/PrefectHQ/prefect/issues/15539
+        """
+
+        @task
+        def random_integer(range_from: int = 0, range_to: int = 100) -> int:
+            """Task that returns a random integer."""
+
+            random_int = random.randint(range_from, range_to)
+
+            return random_int
+
+        @flow(task_runner=task_runner)
+        def add_random_integers(number_tasks: int = 50) -> int:
+            """Flow that submits some random_integer tasks and returns the sum of the results."""
+
+            futures = []
+            for _ in range(number_tasks):
+                futures.append(random_integer.submit())
+
+            sum = 0
+            for future in futures:
+                sum += future.result()
+
+            return sum
+
+        assert add_random_integers() > 0