wehub-resource-sync
126 lines
4.2 KiB
Python
126 lines
4.2 KiB
Python
"""This file responds to log stream requests and forwards logs
|
|
with its handler.
|
|
"""
|
|
import io
|
|
import logging
|
|
import queue
|
|
import threading
|
|
import uuid
|
|
|
|
import grpc
|
|
|
|
import ray.core.generated.ray_client_pb2 as ray_client_pb2
|
|
import ray.core.generated.ray_client_pb2_grpc as ray_client_pb2_grpc
|
|
from ray._private.ray_logging import global_worker_stdstream_dispatcher
|
|
from ray._private.worker import print_worker_logs
|
|
from ray.util.client.common import CLIENT_SERVER_MAX_THREADS
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class LogstreamHandler(logging.Handler):
|
|
def __init__(self, queue, level):
|
|
super().__init__()
|
|
self.queue = queue
|
|
self.level = level
|
|
|
|
def emit(self, record: logging.LogRecord):
|
|
logdata = ray_client_pb2.LogData()
|
|
logdata.msg = record.getMessage()
|
|
logdata.level = record.levelno
|
|
logdata.name = record.name
|
|
self.queue.put(logdata)
|
|
|
|
|
|
class StdStreamHandler:
|
|
def __init__(self, queue):
|
|
self.queue = queue
|
|
self.id = str(uuid.uuid4())
|
|
|
|
def handle(self, data):
|
|
logdata = ray_client_pb2.LogData()
|
|
logdata.level = -2 if data["is_err"] else -1
|
|
logdata.name = "stderr" if data["is_err"] else "stdout"
|
|
with io.StringIO() as file:
|
|
print_worker_logs(data, file)
|
|
logdata.msg = file.getvalue()
|
|
self.queue.put(logdata)
|
|
|
|
def register_global(self):
|
|
global_worker_stdstream_dispatcher.add_handler(self.id, self.handle)
|
|
|
|
def unregister_global(self):
|
|
global_worker_stdstream_dispatcher.remove_handler(self.id)
|
|
|
|
|
|
def log_status_change_thread(log_queue, request_iterator):
|
|
std_handler = StdStreamHandler(log_queue)
|
|
current_handler = None
|
|
root_logger = logging.getLogger("ray")
|
|
default_level = root_logger.getEffectiveLevel()
|
|
try:
|
|
for req in request_iterator:
|
|
if current_handler is not None:
|
|
root_logger.setLevel(default_level)
|
|
root_logger.removeHandler(current_handler)
|
|
std_handler.unregister_global()
|
|
if not req.enabled:
|
|
current_handler = None
|
|
continue
|
|
current_handler = LogstreamHandler(log_queue, req.loglevel)
|
|
std_handler.register_global()
|
|
root_logger.addHandler(current_handler)
|
|
root_logger.setLevel(req.loglevel)
|
|
except grpc.RpcError as e:
|
|
logger.debug(f"closing log thread " f"grpc error reading request_iterator: {e}")
|
|
finally:
|
|
if current_handler is not None:
|
|
root_logger.setLevel(default_level)
|
|
root_logger.removeHandler(current_handler)
|
|
std_handler.unregister_global()
|
|
log_queue.put(None)
|
|
|
|
|
|
class LogstreamServicer(ray_client_pb2_grpc.RayletLogStreamerServicer):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.num_clients = 0
|
|
self.client_lock = threading.Lock()
|
|
|
|
def Logstream(self, request_iterator, context):
|
|
initialized = False
|
|
with self.client_lock:
|
|
threshold = CLIENT_SERVER_MAX_THREADS / 2
|
|
if self.num_clients + 1 >= threshold:
|
|
context.set_code(grpc.StatusCode.RESOURCE_EXHAUSTED)
|
|
logger.warning(
|
|
f"Logstream: Num clients {self.num_clients} has reached "
|
|
f"the threshold {threshold}. Rejecting new connection."
|
|
)
|
|
return
|
|
self.num_clients += 1
|
|
initialized = True
|
|
logger.info(
|
|
"New logs connection established. " f"Total clients: {self.num_clients}"
|
|
)
|
|
log_queue = queue.Queue()
|
|
thread = threading.Thread(
|
|
target=log_status_change_thread,
|
|
args=(log_queue, request_iterator),
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
try:
|
|
queue_iter = iter(log_queue.get, None)
|
|
for record in queue_iter:
|
|
if record is None:
|
|
break
|
|
yield record
|
|
except grpc.RpcError as e:
|
|
logger.debug(f"Closing log channel: {e}")
|
|
finally:
|
|
thread.join()
|
|
with self.client_lock:
|
|
if initialized:
|
|
self.num_clients -= 1
|