chore: import upstream snapshot with attribution

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

@@ -0,0 +1,435 @@
+import os
+import sys
+import tempfile
+import time
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
+
+import aiohttp
+import pytest
+
+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,
+    wait_until_server_available,
+)
+from ray._raylet import GcsClient
+from ray.dashboard.consts import (
+    DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
+    GCS_RPC_TIMEOUT_SECONDS,
+    RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR,
+)
+from ray.dashboard.modules.dashboard_sdk import (
+    DEFAULT_DASHBOARD_ADDRESS,
+    ClusterInfo,
+    parse_cluster_info,
+)
+from ray.dashboard.modules.job.pydantic_models import JobType
+from ray.dashboard.modules.job.sdk import JobStatus, JobSubmissionClient
+from ray.dashboard.tests.conftest import *  # noqa
+from ray.runtime_env.runtime_env import RuntimeEnv
+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)
+    if status == JobStatus.FAILED:
+        logs = client.get_job_logs(job_id)
+        raise RuntimeError(f"Job failed\nlogs:\n{logs}")
+    return status == JobStatus.SUCCEEDED
+
+
+def check_internal_kv_gced():
+    return len(kv._internal_kv_list("gcs://")) == 0
+
+
+@pytest.mark.parametrize(
+    "address_param",
+    [
+        ("ray://1.2.3.4:10001", "ray", "1.2.3.4:10001"),
+        ("other_module://", "other_module", ""),
+        ("other_module://address", "other_module", "address"),
+    ],
+)
+@pytest.mark.parametrize("create_cluster_if_needed", [True, False])
+@pytest.mark.parametrize("cookies", [None, {"test_cookie_key": "test_cookie_val"}])
+@pytest.mark.parametrize("metadata", [None, {"test_metadata_key": "test_metadata_val"}])
+@pytest.mark.parametrize("headers", [None, {"test_headers_key": "test_headers_val"}])
+@pytest.mark.parametrize("extra_kwargs", [{}, {"cloud": "my-cloud"}])
+def test_parse_cluster_info(
+    address_param: Tuple[str, str, str],
+    create_cluster_if_needed: bool,
+    cookies: Optional[Dict[str, str]],
+    metadata: Optional[Dict[str, str]],
+    headers: Optional[Dict[str, str]],
+    extra_kwargs: Dict[str, str],
+):
+    """
+    Test ray.dashboard.modules.dashboard_sdk.parse_cluster_info for different
+    format of addresses.
+    """
+    mock_get_job_submission_client_cluster = Mock(return_value="Ray ClusterInfo")
+    mock_module = Mock()
+    mock_module.get_job_submission_client_cluster_info = Mock(
+        return_value="Other module ClusterInfo"
+    )
+    mock_import_module = Mock(return_value=mock_module)
+
+    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),
+    ):
+        if module_string == "ray":
+            with pytest.raises(ValueError, match="ray://"):
+                parse_cluster_info(
+                    address,
+                    create_cluster_if_needed=create_cluster_if_needed,
+                    cookies=cookies,
+                    metadata=metadata,
+                    headers=headers,
+                    **extra_kwargs,
+                )
+        elif module_string == "other_module":
+            assert (
+                parse_cluster_info(
+                    address,
+                    create_cluster_if_needed=create_cluster_if_needed,
+                    cookies=cookies,
+                    metadata=metadata,
+                    headers=headers,
+                    **extra_kwargs,
+                )
+                == "Other module ClusterInfo"
+            )
+            mock_import_module.assert_called_once_with(module_string)
+            mock_module.get_job_submission_client_cluster_info.assert_called_once_with(
+                inner_address,
+                create_cluster_if_needed=create_cluster_if_needed,
+                cookies=cookies,
+                metadata=metadata,
+                headers=headers,
+                **extra_kwargs,
+            )
+
+
+def test_parse_cluster_info_default_address():
+    assert parse_cluster_info(
+        address=None,
+    ) == ClusterInfo(address=DEFAULT_DASHBOARD_ADDRESS)
+
+
+def test_submit_job_does_not_mutate_runtime_env():
+    class TestClient(JobSubmissionClient):
+        def __init__(self):
+            self._default_metadata = {}
+
+        def _upload_working_dir_if_needed(self, runtime_env):
+            runtime_env["working_dir"] = "gcs://test.zip"
+
+        def _upload_py_modules_if_needed(self, runtime_env):
+            runtime_env["py_modules"] = ["gcs://test_module.zip"]
+
+        def _do_request(self, method, endpoint, **kwargs):
+            return MagicMock(
+                status_code=200,
+                json=lambda: {"job_id": "test_job", "submission_id": "test_job"},
+            )
+
+    runtime_env = {"working_dir": "/tmp/test", "py_modules": ["/tmp/test_module"]}
+    original_runtime_env = {
+        "working_dir": runtime_env["working_dir"],
+        "py_modules": list(runtime_env["py_modules"]),
+    }
+
+    assert (
+        TestClient().submit_job(entrypoint="echo hi", runtime_env=runtime_env)
+        == "test_job"
+    )
+    assert runtime_env == original_runtime_env
+
+
+@pytest.mark.parametrize("expiration_s", [0, 10])
+def test_temporary_uri_reference(monkeypatch, expiration_s):
+    """Test that temporary GCS URI references are deleted after expiration_s."""
+    monkeypatch.setenv(
+        "RAY_RUNTIME_ENV_TEMPORARY_REFERENCE_EXPIRATION_S", str(expiration_s)
+    )
+    # We can't use a fixture with a shared Ray runtime because we need to set the
+    # expiration_s env var before Ray starts.
+    with _ray_start(include_dashboard=True, num_cpus=1) as ctx:
+        headers = {"Connection": "keep-alive", "Authorization": "TOK:<MY_TOKEN>"}
+        address = ctx.address_info["webui_url"]
+        assert wait_until_server_available(address)
+        client = JobSubmissionClient(format_web_url(address), headers=headers)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = Path(tmp_dir)
+
+            hello_file = path / "hi.txt"
+            with hello_file.open(mode="w") as f:
+                f.write("hi\n")
+
+            start = time.time()
+
+            runtime_env = {"working_dir": tmp_dir}
+            job_id = client.submit_job(entrypoint="echo hi", runtime_env=runtime_env)
+            assert runtime_env == {"working_dir": tmp_dir}
+            wait_for_condition(
+                _check_job_succeeded, client=client, job_id=job_id, timeout=30
+            )
+
+            # Give time for deletion to occur if expiration_s is 0.
+            time.sleep(2)
+            # Need to connect to Ray to check internal_kv.
+            # ray.init(address="auto")
+
+            print("Starting Internal KV checks at time ", time.time() - start)
+            if expiration_s > 0:
+                assert not check_internal_kv_gced()
+                wait_for_condition(check_internal_kv_gced, timeout=2 * expiration_s)
+                assert expiration_s < time.time() - start < 2 * expiration_s
+                print("Internal KV was GC'ed at time ", time.time() - start)
+            else:
+                wait_for_condition(check_internal_kv_gced)
+                print("Internal KV was GC'ed at time ", time.time() - start)
+
+                # Regression test for #46625: reusing the same runtime_env after
+                # the package has been GC'ed should re-upload the local working_dir.
+                job_id = client.submit_job(
+                    entrypoint="echo hi", runtime_env=runtime_env
+                )
+                wait_for_condition(
+                    _check_job_succeeded, client=client, job_id=job_id, timeout=30
+                )
+
+
+def get_register_agents_number(gcs_client):
+    keys = gcs_client.internal_kv_keys(
+        prefix=DASHBOARD_AGENT_ADDR_NODE_ID_PREFIX,
+        namespace=KV_NAMESPACE_DASHBOARD,
+        timeout=GCS_RPC_TIMEOUT_SECONDS,
+    )
+    return len(keys)
+
+
+@pytest.mark.parametrize(
+    "ray_start_cluster_head_with_env_vars",
+    [
+        {
+            "include_dashboard": True,
+            "env_vars": {RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR: "1"},
+        },
+        {
+            "include_dashboard": True,
+            "env_vars": {RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR: "0"},
+        },
+    ],
+    indirect=True,
+)
+def test_jobs_run_on_head_by_default_E2E(ray_start_cluster_head_with_env_vars):
+    allow_driver_on_worker_nodes = (
+        os.environ.get(RAY_JOB_ALLOW_DRIVER_ON_WORKER_NODES_ENV_VAR) == "1"
+    )
+    # Cluster setup
+    cluster = ray_start_cluster_head_with_env_vars
+    cluster.add_node(dashboard_agent_listen_port=52366)
+    cluster.add_node(dashboard_agent_listen_port=52367)
+    assert wait_until_server_available(cluster.webui_url) is True
+    webui_url = cluster.webui_url
+    webui_url = format_web_url(webui_url)
+    client = JobSubmissionClient(webui_url)
+    gcs_client = GcsClient(address=cluster.gcs_address)
+
+    def _check_nodes(num_nodes):
+        try:
+            assert len(list_nodes()) == num_nodes
+            return True
+        except Exception as ex:
+            print(ex)
+            return False
+
+    wait_for_condition(lambda: _check_nodes(num_nodes=3), timeout=15)
+    wait_for_condition(lambda: get_register_agents_number(gcs_client) == 3, timeout=20)
+
+    # Submit 20 simple jobs.
+    for i in range(20):
+        client.submit_job(entrypoint="echo hi", submission_id=f"job_{i}")
+    import pprint
+
+    def check_all_jobs_succeeded():
+        submission_jobs = [
+            job for job in client.list_jobs() if job.type == JobType.SUBMISSION
+        ]
+        for job in submission_jobs:
+            pprint.pprint(job)
+            if job.status != JobStatus.SUCCEEDED:
+                return False
+        return True
+
+    # Wait until all jobs have finished.
+    wait_for_condition(check_all_jobs_succeeded, timeout=60, retry_interval_ms=1000)
+
+    # Check driver_node_id of all jobs.
+    submission_jobs = [
+        job for job in client.list_jobs() if job.type == JobType.SUBMISSION
+    ]
+    driver_node_ids = [job.driver_node_id for job in submission_jobs]
+
+    # Spuriously fails with probability (1/3)^20.
+    pprint.pprint(driver_node_ids)
+    num_ids = len(set(driver_node_ids))
+    assert (num_ids > 1) if allow_driver_on_worker_nodes else (num_ids == 1), [
+        id[:5] for id in driver_node_ids
+    ]
+
+
+@pytest.fixture
+def runtime_env_working_dir():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        path = Path(tmp_dir)
+        working_dir = path / "working_dir"
+        working_dir.mkdir(parents=True)
+        yield working_dir
+
+
+@pytest.fixture
+def py_module_whl():
+    with tempfile.NamedTemporaryFile(suffix=".whl") as tmp_file:
+        yield tmp_file.name
+
+
+def test_job_submission_with_runtime_env_as_dict(
+    runtime_env_working_dir, py_module_whl
+):
+    working_dir_str = str(runtime_env_working_dir)
+    with _ray_start(num_cpus=1):
+        client = JobSubmissionClient()
+        runtime_env = {"working_dir": working_dir_str, "py_modules": [py_module_whl]}
+        job_id = client.submit_job(entrypoint="echo hi", runtime_env=runtime_env)
+        job_details = client.get_job_info(job_id)
+        parsed_runtime_env = job_details.runtime_env
+        assert "gcs://" in parsed_runtime_env["working_dir"]
+        assert len(parsed_runtime_env["py_modules"]) == 1
+        assert "gcs://" in parsed_runtime_env["py_modules"][0]
+
+
+def test_job_submission_with_runtime_env_as_object(
+    runtime_env_working_dir, py_module_whl
+):
+    working_dir_str = str(runtime_env_working_dir)
+    with _ray_start(num_cpus=1):
+        client = JobSubmissionClient()
+        runtime_env = RuntimeEnv(
+            working_dir=working_dir_str, py_modules=[py_module_whl]
+        )
+        job_id = client.submit_job(entrypoint="echo hi", runtime_env=runtime_env)
+        job_details = client.get_job_info(job_id)
+        parsed_runtime_env = job_details.runtime_env
+        assert "gcs://" in parsed_runtime_env["working_dir"]
+        assert len(parsed_runtime_env["py_modules"]) == 1
+        assert "gcs://" in parsed_runtime_env["py_modules"][0]
+
+
+@pytest.mark.asyncio
+async def test_tail_job_logs_passes_headers_to_websocket(ray_start_regular):
+    """
+    Test that authentication headers are passed to WebSocket connections.
+
+    This test verifies that headers provided to JobSubmissionClient are
+    explicitly passed to the ws_connect() method, not just to the ClientSession.
+    This is required because aiohttp's ClientSession does not automatically
+    include session headers in WebSocket upgrade requests.
+    """
+    dashboard_url = ray_start_regular.dashboard_url
+    test_headers = {"Authorization": "Bearer test-token"}
+    client = JobSubmissionClient(format_web_url(dashboard_url), headers=test_headers)
+
+    # Submit a simple job
+    job_id = client.submit_job(entrypoint="echo hello")
+
+    # Mock the aiohttp ClientSession and WebSocket
+    mock_ws = AsyncMock()
+    mock_ws.receive = AsyncMock()
+    mock_ws.receive.side_effect = [
+        # First call returns a text message
+        MagicMock(type=aiohttp.WSMsgType.TEXT, data="test log line\n"),
+        # Second call indicates WebSocket is closed
+        MagicMock(type=aiohttp.WSMsgType.CLOSED),
+    ]
+    mock_ws.close_code = 1000  # Normal closure
+
+    mock_session = AsyncMock()
+    mock_session.ws_connect = AsyncMock(return_value=mock_ws)
+    mock_session.__aenter__ = AsyncMock(return_value=mock_session)
+    mock_session.__aexit__ = AsyncMock(return_value=None)
+
+    # Patch ClientSession to use our mock
+    with patch("aiohttp.ClientSession", return_value=mock_session):
+        # Tail logs
+        log_lines = []
+        async for lines in client.tail_job_logs(job_id):
+            log_lines.append(lines)
+
+        # Verify ws_connect was called with headers
+        mock_session.ws_connect.assert_called_once()
+        call_args = mock_session.ws_connect.call_args
+
+        assert "headers" in call_args.kwargs
+        assert call_args.kwargs["headers"] == test_headers
+
+
+@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__]))