import logging import os import sys import time from typing import List, Tuple import pytest import ray from ray._common.test_utils import wait_for_condition from ray._private import ray_constants from ray._private.runtime_env.agent.runtime_env_agent import ReferenceTable, UriType from ray._private.test_utils import ( get_error_message, init_error_pubsub, ) from ray.core.generated import common_pb2 from ray.runtime_env import RuntimeEnv import psutil logger = logging.getLogger(__name__) def test_reference_table(): expected_unused_uris = [] expected_unused_runtime_env = str() def uris_parser(runtime_env) -> Tuple[str, UriType]: result = list() result.append((runtime_env.working_dir(), "working_dir")) py_module_uris = runtime_env.py_modules() for uri in py_module_uris: result.append((uri, "py_modules")) return result def unused_uris_processor(unused_uris: List[Tuple[str, UriType]]) -> None: nonlocal expected_unused_uris assert expected_unused_uris for unused in unused_uris: assert unused in expected_unused_uris expected_unused_uris.remove(unused) assert not expected_unused_uris def unused_runtime_env_processor(unused_runtime_env: str) -> None: nonlocal expected_unused_runtime_env assert expected_unused_runtime_env assert expected_unused_runtime_env == unused_runtime_env expected_unused_runtime_env = None reference_table = ReferenceTable( uris_parser, unused_uris_processor, unused_runtime_env_processor ) runtime_env_1 = RuntimeEnv( working_dir="s3://working_dir_1.zip", py_modules=["s3://py_module_A.zip", "s3://py_module_B.zip"], ) runtime_env_2 = RuntimeEnv( working_dir="s3://working_dir_2.zip", py_modules=["s3://py_module_A.zip", "s3://py_module_C.zip"], ) # Add runtime env 1 reference_table.increase_reference( runtime_env_1, runtime_env_1.serialize(), "raylet" ) # Add runtime env 2 reference_table.increase_reference( runtime_env_2, runtime_env_2.serialize(), "raylet" ) # Add runtime env 1 by `client_server`, this will be skipped by reference table. reference_table.increase_reference( runtime_env_1, runtime_env_1.serialize(), "client_server" ) # Remove runtime env 1 expected_unused_uris.append(("s3://working_dir_1.zip", "working_dir")) expected_unused_uris.append(("s3://py_module_B.zip", "py_modules")) expected_unused_runtime_env = runtime_env_1.serialize() reference_table.decrease_reference( runtime_env_1, runtime_env_1.serialize(), "raylet" ) assert not expected_unused_uris assert not expected_unused_runtime_env # Remove runtime env 2 expected_unused_uris.append(("s3://working_dir_2.zip", "working_dir")) expected_unused_uris.append(("s3://py_module_A.zip", "py_modules")) expected_unused_uris.append(("s3://py_module_C.zip", "py_modules")) expected_unused_runtime_env = runtime_env_2.serialize() reference_table.decrease_reference( runtime_env_2, runtime_env_2.serialize(), "raylet" ) assert not expected_unused_uris assert not expected_unused_runtime_env def search_agent(processes): for p in processes: try: for c in p.cmdline(): # in case linux truncates the proctitle if ray_constants.AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT[:15] in c: return p except Exception: pass def check_agent_register(raylet_proc, agent_pid): # Check if agent register is OK. for x in range(5): logger.info("Check agent is alive.") agent_proc = search_agent(raylet_proc.children()) assert agent_proc.pid == agent_pid time.sleep(1) @pytest.mark.skipif(sys.platform == "win32", reason="no fate sharing for windows") def test_raylet_and_agent_share_fate(shutdown_only): """Test raylet and agent share fate.""" ray.init() p = init_error_pubsub() node = ray._private.worker._global_node all_processes = node.all_processes raylet_proc_info = all_processes[ray_constants.PROCESS_TYPE_RAYLET][0] raylet_proc = psutil.Process(raylet_proc_info.process.pid) wait_for_condition(lambda: search_agent(raylet_proc.children())) agent_proc = search_agent(raylet_proc.children()) agent_pid = agent_proc.pid check_agent_register(raylet_proc, agent_pid) # The agent should be dead if raylet exits. raylet_proc.terminate() raylet_proc.wait() agent_proc.wait(15) # No error should be reported for graceful termination. errors = get_error_message(p, 1, ray_constants.RAYLET_DIED_ERROR) assert len(errors) == 0, errors ray.shutdown() ray_context = ray.init() all_processes = ray._private.worker._global_node.all_processes raylet_proc_info = all_processes[ray_constants.PROCESS_TYPE_RAYLET][0] raylet_proc = psutil.Process(raylet_proc_info.process.pid) wait_for_condition(lambda: search_agent(raylet_proc.children())) agent_proc = search_agent(raylet_proc.children()) agent_pid = agent_proc.pid check_agent_register(raylet_proc, agent_pid) # The raylet should be dead if agent exits. agent_proc.kill() agent_proc.wait() raylet_proc.wait(15) worker_node_id = ray_context.address_info["node_id"] worker_node_info = [ node for node in ray.nodes() if node["NodeID"] == worker_node_id ][0] assert not worker_node_info["Alive"] assert worker_node_info["DeathReason"] == common_pb2.NodeDeathInfo.Reason.Value( "UNEXPECTED_TERMINATION" ) assert ( "failed and raylet fate-shares with it." in worker_node_info["DeathReasonMessage"] ) @pytest.mark.skipif(sys.platform == "win32", reason="no fate sharing for windows") def test_agent_report_unexpected_raylet_death(shutdown_only): """Test agent reports Raylet death if it is not SIGTERM.""" ray.init() p = init_error_pubsub() node = ray._private.worker._global_node all_processes = node.all_processes raylet_proc_info = all_processes[ray_constants.PROCESS_TYPE_RAYLET][0] raylet_proc = psutil.Process(raylet_proc_info.process.pid) wait_for_condition(lambda: search_agent(raylet_proc.children())) agent_proc = search_agent(raylet_proc.children()) agent_pid = agent_proc.pid check_agent_register(raylet_proc, agent_pid) # The agent should be dead if raylet exits. raylet_proc.kill() raylet_proc.wait() agent_proc.wait(15) errors = get_error_message(p, 1, ray_constants.RAYLET_DIED_ERROR) assert len(errors) == 1, errors err = errors[0] assert err["type"] == ray_constants.RAYLET_DIED_ERROR assert "Termination is unexpected." in err["error_message"], err["error_message"] assert "Raylet logs:" in err["error_message"], err["error_message"] assert ( os.path.getsize(os.path.join(node.get_session_dir_path(), "logs", "raylet.out")) < 1 * 1024**2 ) @pytest.mark.skipif(sys.platform == "win32", reason="no fate sharing for windows") def test_agent_report_unexpected_raylet_death_large_file(shutdown_only): """Test agent reports Raylet death if it is not SIGTERM.""" ray.init() p = init_error_pubsub() node = ray._private.worker._global_node all_processes = node.all_processes raylet_proc_info = all_processes[ray_constants.PROCESS_TYPE_RAYLET][0] raylet_proc = psutil.Process(raylet_proc_info.process.pid) wait_for_condition(lambda: search_agent(raylet_proc.children())) agent_proc = search_agent(raylet_proc.children()) agent_pid = agent_proc.pid check_agent_register(raylet_proc, agent_pid) # Append to the Raylet log file with data >> 1 MB. with open( os.path.join(node.get_session_dir_path(), "logs", "raylet.out"), "a" ) as f: f.write("test data\n" * 1024**2) # The agent should be dead if raylet exits. raylet_proc.kill() raylet_proc.wait() agent_proc.wait(15) # Reading and publishing logs should still work. errors = get_error_message(p, 1, ray_constants.RAYLET_DIED_ERROR) assert len(errors) == 1, errors err = errors[0] assert err["type"] == ray_constants.RAYLET_DIED_ERROR assert "Termination is unexpected." in err["error_message"], err["error_message"] assert "Raylet logs:" in err["error_message"], err["error_message"] if __name__ == "__main__": sys.exit(pytest.main(["-v", __file__]))