[FIX] raise error if job does not terminate in tail_job_logs()

python/ray/dashboard/modules/job/sdk.py

@@ -482,8 +482,9 @@ async def tail_job_logs(self, job_id: str) -> AsyncIterator[str]:
             The iterator.
 
         Raises:
-            RuntimeError: If the job does not exist or if the request to the
-                job server fails.
+            RuntimeError: If the job does not exist, if the request to the
+                job server fails, or if the connection closes unexpectedly
+                before the job reaches a terminal state.
         """
         async with aiohttp.ClientSession(
             cookies=self._cookies, headers=self._headers
@@ -498,6 +499,17 @@ async def tail_job_logs(self, job_id: str) -> AsyncIterator[str]:
                 if msg.type == aiohttp.WSMsgType.TEXT:
                     yield msg.data
                 elif msg.type == aiohttp.WSMsgType.CLOSED:
+                    logger.debug(
+                        f"WebSocket closed for job {job_id} with close code {ws.close_code}"
+                    )
+                    if ws.close_code == aiohttp.WSCloseCode.ABNORMAL_CLOSURE:
+                        raise RuntimeError(
+                            f"WebSocket connection closed unexpectedly with close code {ws.close_code}"
+                        )
                     break
                 elif msg.type == aiohttp.WSMsgType.ERROR:
-                    pass
+                    # Old Ray versions may send ERROR on connection close
+                    logger.debug(
+                        f"WebSocket error for job {job_id}, treating as normal close. Err: {ws.exception()}"
+                    )
+                    break

python/ray/dashboard/modules/job/tests/test_http_job_server.py

@@ -268,8 +268,7 @@ def f():
                 yield {
                     "runtime_env": {"py_modules": [str(Path(tmp_dir) / "test_module")]},
                     "entrypoint": (
-                        "python -c 'import test_module;"
-                        "print(test_module.run_test())'"
+                        "python -c 'import test_module;print(test_module.run_test())'"
                     ),
                     "expected_logs": "Hello from test_module!\n",
                 }

python/ray/dashboard/modules/job/tests/test_sdk.py

@@ -10,8 +10,10 @@
 
 import ray.experimental.internal_kv as kv
 from ray._common.test_utils import wait_for_condition
+from ray._private import worker
 from ray._private.ray_constants import (
     KV_NAMESPACE_DASHBOARD,
+    PROCESS_TYPE_DASHBOARD,
 )
 from ray._private.test_utils import (
     format_web_url,
@@ -35,6 +37,8 @@
 from ray.tests.conftest import _ray_start
 from ray.util.state import list_nodes
 
+import psutil
+
 
 def _check_job_succeeded(client: JobSubmissionClient, job_id: str) -> bool:
     status = client.get_job_status(job_id)
@@ -80,10 +84,13 @@ def test_parse_cluster_info(
 
     address, module_string, inner_address = address_param
 
-    with patch.multiple(
-        "ray.dashboard.modules.dashboard_sdk",
-        get_job_submission_client_cluster_info=mock_get_job_submission_client_cluster,
-    ), patch.multiple("importlib", import_module=mock_import_module):
+    with (
+        patch.multiple(
+            "ray.dashboard.modules.dashboard_sdk",
+            get_job_submission_client_cluster_info=mock_get_job_submission_client_cluster,
+        ),
+        patch.multiple("importlib", import_module=mock_import_module),
+    ):
         if module_string == "ray":
             with pytest.raises(ValueError, match="ray://"):
                 parse_cluster_info(
@@ -289,5 +296,44 @@ def test_job_submission_with_runtime_env_as_object(
         assert "gcs://" in parsed_runtime_env["py_modules"][0]
 
 
+@pytest.mark.asyncio
+async def test_tail_job_logs_websocket_abnormal_closure(ray_start_regular):
+    """
+    Test that ABNORMAL_CLOSURE raises RuntimeError when tailing logs.
+
+    This test uses its own Ray cluster and kills the dashboard while tailing logs
+    to simulate an abnormal WebSocket closure.
+    """
+    dashboard_url = ray_start_regular.dashboard_url
+    client = JobSubmissionClient(format_web_url(dashboard_url))
+
+    # Submit a long-running job
+    driver_script = """
+import time
+for i in range(100):
+    print("Hello", i)
+    time.sleep(0.5)
+"""
+    entrypoint = f"python -c '{driver_script}'"
+    job_id = client.submit_job(entrypoint=entrypoint)
+
+    # Start tailing logs and stop Ray while tailing
+    # Expect RuntimeError when WebSocket closes abnormally
+    with pytest.raises(
+        RuntimeError,
+        match="WebSocket connection closed unexpectedly with close code",
+    ):
+        i = 0
+        async for lines in client.tail_job_logs(job_id):
+            print(lines, end="")
+            i += 1
+
+            # Kill the dashboard after receiving a few log lines
+            if i == 3:
+                print("\nKilling the dashboard to close websocket abnormally...")
+                dash_info = worker._global_node.all_processes[PROCESS_TYPE_DASHBOARD][0]
+                psutil.Process(dash_info.process.pid).kill()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))