[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",
                 }
@@ -711,6 +710,52 @@ async def test_tail_job_logs(job_sdk_client):
         wait_for_condition(_check_job_succeeded, client=client, job_id=job_id)
 
 
+@pytest.mark.asyncio
+async def test_tail_job_logs_websocket_abnormal_closure():
+    """
+    Test that ABNORMAL_CLOSURE raises RuntimeError when tailing logs.
+
+    This test starts its own Ray cluster using default ports,
+    then forcefully stops Ray while tailing logs to simulate an abnormal WebSocket closure.
+    """
+    subprocess.check_output(["ray", "start", "--head"])
+    address = "127.0.0.1:8265"
+
+    try:
+        assert wait_until_server_available(address)
+        client = JobSubmissionClient(format_web_url(address))
+
+        # 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
+
+                # Run ray stop to terminate Ray after receiving a few log lines
+                if i == 3:
+                    print("\nStopping Ray cluster forcefully...")
+                    subprocess.check_output(["ray", "stop", "--force"])
+
+    finally:
+        # Ensure Ray is stopped even if test fails
+        subprocess.check_output(["ray", "stop", "--force"])
+
+
 def _hook(env):
     with open(env["env_vars"]["TEMPPATH"], "w+") as f:
         f.write(env["env_vars"]["TOKEN"])