chore: import upstream snapshot with attribution

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

@@ -0,0 +1,177 @@
+import json
+import os
+import pprint
+import sys
+
+import jsonschema
+import pytest
+import requests
+
+from ray._common.test_utils import (
+    run_string_as_driver,
+    wait_for_condition,
+)
+from ray._private.test_utils import (
+    format_web_url,
+    run_string_as_driver_nonblocking,
+)
+from ray.dashboard import dashboard
+from ray.dashboard.consts import RAY_CLUSTER_ACTIVITY_HOOK
+from ray.dashboard.modules.job.job_head import RayActivityResponse
+from ray.dashboard.tests.conftest import *  # noqa
+
+
+@pytest.fixture
+def set_ray_cluster_activity_hook(request):
+    """
+    Fixture that sets RAY_CLUSTER_ACTIVITY_HOOK environment variable
+    for test_e2e_component_activities_hook.
+    """
+    external_hook = request.param
+    assert (
+        external_hook
+    ), "Please pass value of RAY_CLUSTER_ACTIVITY_HOOK env var to this fixture"
+    old_hook = os.environ.get(RAY_CLUSTER_ACTIVITY_HOOK)
+    os.environ[RAY_CLUSTER_ACTIVITY_HOOK] = external_hook
+
+    yield external_hook
+
+    if old_hook is not None:
+        os.environ[RAY_CLUSTER_ACTIVITY_HOOK] = old_hook
+    else:
+        del os.environ[RAY_CLUSTER_ACTIVITY_HOOK]
+
+
+@pytest.mark.parametrize(
+    "set_ray_cluster_activity_hook",
+    [
+        "ray._private.test_utils.external_ray_cluster_activity_hook1",
+        "ray._private.test_utils.external_ray_cluster_activity_hook2",
+        "ray._private.test_utils.external_ray_cluster_activity_hook3",
+        "ray._private.test_utils.external_ray_cluster_activity_hook4",
+        "ray._private.test_utils.external_ray_cluster_activity_hook5",
+    ],
+    indirect=True,
+)
+def test_component_activities_hook(set_ray_cluster_activity_hook, call_ray_start):
+    """
+    Tests /api/component_activities returns correctly for various
+    responses of RAY_CLUSTER_ACTIVITY_HOOK.
+
+    Verify no active drivers are correctly reflected in response.
+    """
+    external_hook = set_ray_cluster_activity_hook
+
+    response = requests.get("http://127.0.0.1:8265/api/component_activities")
+    response.raise_for_status()
+
+    # Validate schema of response
+    data = response.json()
+    schema_path = os.path.join(
+        os.path.dirname(dashboard.__file__),
+        "modules/job/component_activities_schema.json",
+    )
+    pprint.pprint(data)
+    jsonschema.validate(instance=data, schema=json.load(open(schema_path)))
+
+    # Validate driver response can be cast to RayActivityResponse object
+    # and that there are no active drivers.
+    driver_ray_activity_response = RayActivityResponse(**data["driver"])
+    assert driver_ray_activity_response.is_active == "INACTIVE"
+    assert driver_ray_activity_response.reason is None
+
+    # Validate external component response can be cast to RayActivityResponse object
+    if external_hook[-1] == "5":
+        external_activity_response = RayActivityResponse(**data["test_component5"])
+        assert external_activity_response.is_active == "ACTIVE"
+        assert external_activity_response.reason == "Counter: 1"
+    elif external_hook[-1] == "4":
+        external_activity_response = RayActivityResponse(**data["external_component"])
+        assert external_activity_response.is_active == "ERROR"
+        assert (
+            "'Error in external cluster activity hook'"
+            in external_activity_response.reason
+        )
+    elif external_hook[-1] == "3":
+        external_activity_response = RayActivityResponse(**data["external_component"])
+        assert external_activity_response.is_active == "ERROR"
+    elif external_hook[-1] == "2":
+        external_activity_response = RayActivityResponse(**data["test_component2"])
+        assert external_activity_response.is_active == "ERROR"
+    elif external_hook[-1] == "1":
+        external_activity_response = RayActivityResponse(**data["test_component1"])
+        assert external_activity_response.is_active == "ACTIVE"
+        assert external_activity_response.reason == "Counter: 1"
+
+        # Call endpoint again to validate different response
+        response = requests.get("http://127.0.0.1:8265/api/component_activities")
+        response.raise_for_status()
+        data = response.json()
+        jsonschema.validate(instance=data, schema=json.load(open(schema_path)))
+
+        external_activity_response = RayActivityResponse(**data["test_component1"])
+        assert external_activity_response.is_active == "ACTIVE"
+        assert external_activity_response.reason == "Counter: 2"
+
+
+def test_active_component_activities(ray_start_with_dashboard):
+    # Verify drivers which don't have namespace starting with _ray_internal_
+    # are considered active.
+
+    webui_url = ray_start_with_dashboard["webui_url"]
+    webui_url = format_web_url(webui_url)
+
+    driver_template = """
+import ray
+
+ray.init(address="auto", namespace="{namespace}")
+import time
+time.sleep({sleep_time_s})
+    """
+    run_string_as_driver(
+        driver_template.format(namespace="my_namespace", sleep_time_s=0)
+    )
+
+    run_string_as_driver_nonblocking(
+        driver_template.format(namespace="my_namespace", sleep_time_s=5)
+    )
+    run_string_as_driver_nonblocking(
+        driver_template.format(namespace="_ray_internal_job_info_id1", sleep_time_s=5)
+    )
+    # Simulate the default driver that gets created by dashboard
+    run_string_as_driver_nonblocking(
+        driver_template.format(namespace="_ray_internal_dashboard", sleep_time_s=5)
+    )
+
+    def verify_driver_response():
+        # Verify drivers are considered active after running script
+        response = requests.get(f"{webui_url}/api/component_activities")
+        response.raise_for_status()
+
+        # Validate schema of response
+        data = response.json()
+        schema_path = os.path.join(
+            os.path.dirname(dashboard.__file__),
+            "modules/job/component_activities_schema.json",
+        )
+
+        jsonschema.validate(instance=data, schema=json.load(open(schema_path)))
+
+        # Validate ray_activity_response field can be cast to RayActivityResponse object
+        driver_ray_activity_response = RayActivityResponse(**data["driver"])
+        print(driver_ray_activity_response)
+
+        assert driver_ray_activity_response.is_active == "ACTIVE"
+        # Drivers with namespace starting with "_ray_internal" are not
+        # considered active drivers. Two active drivers are the second one
+        # run with namespace "my_namespace" and the one started
+        # from ray_start_with_dashboard
+        assert driver_ray_activity_response.reason == "Number of active drivers: 2"
+
+        return True
+
+    wait_for_condition(verify_driver_response)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))