import base64
import logging
from collections import defaultdict
from enum import Enum
from typing import List
import ray
from ray._private.internal_api import node_stats
from ray._raylet import ActorID, JobID, TaskID
logger = logging.getLogger(__name__)
# These values are used to calculate if objectRefs are actor handles.
TASKID_BYTES_SIZE = TaskID.size()
ACTORID_BYTES_SIZE = ActorID.size()
JOBID_BYTES_SIZE = JobID.size()
# We need to multiply 2 because we need bits size instead of bytes size.
TASKID_RANDOM_BITS_SIZE = (TASKID_BYTES_SIZE - ACTORID_BYTES_SIZE) * 2
ACTORID_RANDOM_BITS_SIZE = (ACTORID_BYTES_SIZE - JOBID_BYTES_SIZE) * 2
def decode_object_ref_if_needed(object_ref: str) -> bytes:
"""Decode objectRef bytes string.
gRPC reply contains an objectRef that is encodded by Base64.
This function is used to decode the objectRef.
Note that there are times that objectRef is already decoded as
a hex string. In this case, just convert it to a binary number.
"""
if object_ref.endswith("="):
# If the object ref ends with =, that means it is base64 encoded.
async def test_logs_manager_resolve_file(logs_manager):
node_id = NodeID(b"1" * 28)
"""
Test filename is given.
"""
logs_client = logs_manager.data_source_client
logs_client.get_all_registered_agent_ids = MagicMock()
logs_client.get_all_registered_agent_ids.return_value = [node_id.hex()]
expected_filename = "filename"
log_file_name, n = await logs_manager.resolve_filename(
node_id=node_id,
log_filename=expected_filename,
actor_id=None,
task_id=None,
pid=None,
get_actor_fn=lambda _: True,
timeout=10,
)
assert log_file_name == expected_filename
assert n == node_id
"""
Test actor id is given.
"""
# Actor doesn't exist.
with pytest.raises(ValueError):
actor_id = ActorID(b"2" * 16)
def get_actor_fn(id):
if id == actor_id:
return None
assert False, "Not reachable."
log_file_name, n = await logs_manager.resolve_filename(
node_id=node_id,
log_filename=None,
actor_id=actor_id,
task_id=None,
pid=None,
get_actor_fn=get_actor_fn,
timeout=10,
)
# Actor exists, but it is not scheduled yet.
actor_id = ActorID(b"2" * 16)
with pytest.raises(ValueError):
log_file_name, n = await logs_manager.resolve_filename(
node_id=node_id,
log_filename=None,
actor_id=actor_id,
task_id=None,
pid=None,
get_actor_fn=lambda _: generate_actor_data(actor_id, node_id, None),
timeout=10,
)
# Actor exists.
actor_id = ActorID(b"2" * 16)
worker_id = WorkerID(b"3" * 28)
logs_manager.list_logs = AsyncMock()
logs_manager.list_logs.return_value = {
"worker_out": [f"worker-{worker_id.hex()}-123-123.out"]
}
log_file_name, n = await logs_manager.resolve_filename(
node_id=node_id.hex(),
log_filename=None,
actor_id=actor_id,
task_id=None,
pid=None,
get_actor_fn=lambda _: generate_actor_data(actor_id, node_id, worker_id),
timeout=10,
)
logs_manager.list_logs.assert_awaited_with(
node_id.hex(), 10, glob_filter=f"*{worker_id.hex()}*"
)
assert log_file_name == f"worker-{worker_id.hex()}-123-123.out"
assert n == node_id.hex()
"""
Test task id is given.
"""
with pytest.raises(NotImplementedError):
task_id = TaskID(b"2" * 24)
log_file_name, n = await logs_manager.resolve_filename(
node_id=node_id.hex(),
log_filename=None,
actor_id=None,
task_id=task_id,
pid=None,
get_actor_fn=lambda _: generate_actor_data(actor_id, node_id, worker_id),
timeout=10,
)
"""
Test pid is given.
"""
# Pid doesn't exist.
with pytest.raises(FileNotFoundError):
pid = 456
logs_manager.list_logs = AsyncMock()
# Provide the wrong pid.
logs_manager.list_logs.return_value = {"worker_out": ["worker-123-123-123.out"]}
log_file_name = await logs_manager.resolve_filename(
node_id=node_id.hex(),
log_filename=None,
actor_id=None,
task_id=None,
pid=pid,
get_actor_fn=lambda _: generate_actor_data(actor_id, node_id, worker_id),
timeout=10,
)
# Pid exists.
pid = 123
logs_manager.list_logs = AsyncMock()
# Provide the wrong pid.
logs_manager.list_logs.return_value = {"worker_out": [f"worker-123-123-{pid}.out"]}
log_file_name, n = await logs_manager.resolve_filename(
node_id=node_id.hex(),
log_filename=None,
actor_id=None,
task_id=None,
pid=pid,
get_actor_fn=lambda _: generate_actor_data(actor_id, node_id, worker_id),
timeout=10,
)
logs_manager.list_logs.assert_awaited_with(
node_id.hex(), 10, glob_filter=f"*{pid}*"
)
assert log_file_name == f"worker-123-123-{pid}.out"
"""
Test nothing is given.
"""
with pytest.raises(FileNotFoundError):
log_file_name = await logs_manager.resolve_filename(
node_id=node_id.hex(),
log_filename=None,
actor_id=None,
task_id=None,
pid=None,
get_actor_fn=lambda _: generate_actor_data(actor_id, node_id, worker_id),
timeout=10,
)
async def test_api_manager_summary_tasks(state_api_manager):
data_source_client = state_api_manager.data_source_client
data_source_client.get_all_registered_raylet_ids = MagicMock()
data_source_client.get_all_registered_raylet_ids.return_value = ["1", "2"]
first_task_name = "1"
second_task_name = "2"
data_source_client.get_task_info = AsyncMock()
ids = [TaskID((f"{i}" * 24).encode()) for i in range(5)]
# 1: {SCHEDULED:3, RUNNING:1}, 2:{SCHEDULED: 1}
data_source_client.get_task_info.side_effect = [
GetTasksInfoReply(owned_task_info_entries=[
generate_task_entry(
id=ids[0].binary(),
func_or_class=first_task_name,
state=TaskStatus.SCHEDULED,
type=TaskType.NORMAL_TASK,
),
generate_task_entry(
id=ids[1].binary(),
func_or_class=first_task_name,
state=TaskStatus.SCHEDULED,
type=TaskType.NORMAL_TASK,
),
generate_task_entry(
id=ids[2].binary(),
func_or_class=first_task_name,
state=TaskStatus.SCHEDULED,
type=TaskType.NORMAL_TASK,
),
]),
GetTasksInfoReply(owned_task_info_entries=[
generate_task_entry(
id=ids[3].binary(),
func_or_class=first_task_name,
state=TaskStatus.RUNNING,
type=TaskType.NORMAL_TASK,
),
generate_task_entry(
id=ids[4].binary(),
func_or_class=second_task_name,
state=TaskStatus.SCHEDULED,
type=TaskType.ACTOR_TASK,
),
]),
]
"""
Test cluster summary.
"""
result = await state_api_manager.summarize_tasks(
option=create_summary_options())
assert "cluster" in result.result.node_id_to_summary
data = result.result.node_id_to_summary["cluster"]
assert data.summary[first_task_name].type == "NORMAL_TASK"
assert data.summary[first_task_name].func_or_class_name == first_task_name
assert data.summary[first_task_name].state_counts["SCHEDULED"] == 3
assert data.summary[first_task_name].state_counts["RUNNING"] == 1
assert data.summary[second_task_name].type == "ACTOR_TASK"
assert data.summary[
second_task_name].func_or_class_name == second_task_name
assert data.summary[second_task_name].state_counts["SCHEDULED"] == 1
assert data.total_tasks == 4
assert data.total_actor_tasks == 1
assert data.total_actor_scheduling_tasks == 0
"""
Test if it can be correctly modified to a dictionary.
"""
print(result.result)
result_in_dict = asdict(result.result)
assert json.loads(json.dumps(result_in_dict)) == result_in_dict