chore: import upstream snapshot with attribution

python/ray/_private/runtime_env/agent/main.py

@@ -0,0 +1,266 @@
+import argparse
+import logging
+import os
+import socket
+import sys
+
+import ray
+import ray._private.ray_constants as ray_constants
+from ray._common.utils import (
+    get_or_create_event_loop,
+)
+from ray._private import logging_utils
+from ray._private.authentication.http_token_authentication import (
+    get_token_auth_middleware,
+)
+from ray._private.process_watcher import create_check_raylet_task
+from ray._raylet import RUNTIME_ENV_AGENT_PORT_NAME, GcsClient, persist_port
+from ray.core.generated import (
+    runtime_env_agent_pb2,
+)
+
+
+def import_libs():
+    my_dir = os.path.abspath(os.path.dirname(__file__))
+    sys.path.insert(0, os.path.join(my_dir, "thirdparty_files"))  # for aiohttp
+    sys.path.insert(0, my_dir)  # for runtime_env_agent and runtime_env_consts
+
+
+import_libs()
+
+import aiohttp  # noqa: E402
+import runtime_env_consts  # noqa: E402
+from aiohttp import web  # noqa: E402
+from runtime_env_agent import RuntimeEnvAgent  # noqa: E402
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Runtime env agent.")
+    parser.add_argument(
+        "--node-id",
+        required=True,
+        type=str,
+        help="the unique ID of this node.",
+    )
+    parser.add_argument(
+        "--node-ip-address",
+        required=True,
+        type=str,
+        help="the IP address of this node.",
+    )
+    parser.add_argument(
+        "--runtime-env-agent-port",
+        required=True,
+        type=int,
+        default=None,
+        help="The port on which the runtime env agent will receive HTTP requests.",
+    )
+    parser.add_argument(
+        "--session-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="The path of this ray session directory.",
+    )
+
+    parser.add_argument(
+        "--gcs-address", required=True, type=str, help="The address (ip:port) of GCS."
+    )
+    parser.add_argument(
+        "--cluster-id-hex", required=True, type=str, help="The cluster id in hex."
+    )
+    parser.add_argument(
+        "--runtime-env-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of the resource directory used by runtime_env.",
+    )
+
+    parser.add_argument(
+        "--logging-level",
+        required=False,
+        type=lambda s: logging.getLevelName(s.upper()),
+        default=ray_constants.LOGGER_LEVEL,
+        choices=ray_constants.LOGGER_LEVEL_CHOICES,
+        help=ray_constants.LOGGER_LEVEL_HELP,
+    )
+    parser.add_argument(
+        "--logging-format",
+        required=False,
+        type=str,
+        default=ray_constants.LOGGER_FORMAT,
+        help=ray_constants.LOGGER_FORMAT_HELP,
+    )
+    parser.add_argument(
+        "--logging-filename",
+        required=False,
+        type=str,
+        default=runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME,
+        help="Specify the name of log file, "
+        'log to stdout if set empty, default is "{}".'.format(
+            runtime_env_consts.RUNTIME_ENV_AGENT_LOG_FILENAME
+        ),
+    )
+    parser.add_argument(
+        "--logging-rotate-bytes",
+        required=True,
+        type=int,
+        help="Specify the max bytes for rotating log file",
+    )
+    parser.add_argument(
+        "--logging-rotate-backup-count",
+        required=True,
+        type=int,
+        help="Specify the backup count of rotated log file",
+    )
+    parser.add_argument(
+        "--log-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of log directory.",
+    )
+    parser.add_argument(
+        "--temp-dir",
+        required=True,
+        type=str,
+        default=None,
+        help="Specify the path of the temporary directory use by Ray process.",
+    )
+    parser.add_argument(
+        "--stdout-filepath",
+        required=False,
+        type=str,
+        default="",
+        help="The filepath to dump runtime env agent stdout.",
+    )
+    parser.add_argument(
+        "--stderr-filepath",
+        required=False,
+        type=str,
+        default="",
+        help="The filepath to dump runtime env agent stderr.",
+    )
+
+    args = parser.parse_args()
+
+    # Disable log rotation for windows platform.
+    logging_rotation_bytes = args.logging_rotate_bytes if sys.platform != "win32" else 0
+    logging_rotation_backup_count = (
+        args.logging_rotate_backup_count if sys.platform != "win32" else 1
+    )
+
+    logging_params = dict(
+        logging_level=args.logging_level,
+        logging_format=args.logging_format,
+        log_dir=args.log_dir,
+        filename=args.logging_filename,
+        max_bytes=logging_rotation_bytes,
+        backup_count=logging_rotation_backup_count,
+    )
+
+    # Setup stdout/stderr redirect files if redirection enabled.
+    logging_utils.redirect_stdout_stderr_if_needed(
+        args.stdout_filepath,
+        args.stderr_filepath,
+        logging_rotation_bytes,
+        logging_rotation_backup_count,
+    )
+
+    gcs_client = GcsClient(address=args.gcs_address, cluster_id=args.cluster_id_hex)
+    agent = RuntimeEnvAgent(
+        runtime_env_dir=args.runtime_env_dir,
+        logging_params=logging_params,
+        gcs_client=gcs_client,
+        temp_dir=args.temp_dir,
+        address=args.node_ip_address,
+        runtime_env_agent_port=args.runtime_env_agent_port,
+    )
+
+    ray._raylet.setproctitle(ray_constants.AGENT_PROCESS_TYPE_RUNTIME_ENV_AGENT)
+
+    # POST /get_or_create_runtime_env
+    # body is serialzied protobuf GetOrCreateRuntimeEnvRequest
+    # reply is serialzied protobuf GetOrCreateRuntimeEnvReply
+    async def get_or_create_runtime_env(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.GetOrCreateRuntimeEnvRequest()
+        request.ParseFromString(data)
+        reply = await agent.GetOrCreateRuntimeEnv(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    # POST /delete_runtime_env_if_possible
+    # body is serialzied protobuf DeleteRuntimeEnvIfPossibleRequest
+    # reply is serialzied protobuf DeleteRuntimeEnvIfPossibleReply
+    async def delete_runtime_env_if_possible(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.DeleteRuntimeEnvIfPossibleRequest()
+        request.ParseFromString(data)
+        reply = await agent.DeleteRuntimeEnvIfPossible(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    # POST /get_runtime_envs_info
+    # body is serialzied protobuf GetRuntimeEnvsInfoRequest
+    # reply is serialzied protobuf GetRuntimeEnvsInfoReply
+    async def get_runtime_envs_info(request: web.Request) -> web.Response:
+        data = await request.read()
+        request = runtime_env_agent_pb2.GetRuntimeEnvsInfoRequest()
+        request.ParseFromString(data)
+        reply = await agent.GetRuntimeEnvsInfo(request)
+        return web.Response(
+            body=reply.SerializeToString(), content_type="application/octet-stream"
+        )
+
+    app = web.Application(middlewares=[get_token_auth_middleware(aiohttp)])
+
+    app.router.add_post("/get_or_create_runtime_env", get_or_create_runtime_env)
+    app.router.add_post(
+        "/delete_runtime_env_if_possible", delete_runtime_env_if_possible
+    )
+    app.router.add_post("/get_runtime_envs_info", get_runtime_envs_info)
+
+    loop = get_or_create_event_loop()
+    check_raylet_task = None
+    if sys.platform not in ["win32", "cygwin"]:
+
+        def parent_dead_callback(msg):
+            agent._logger.info(
+                "Raylet is dead! Exiting Runtime Env Agent. "
+                f"addr: {args.node_ip_address}, "
+                f"port: {args.runtime_env_agent_port}\n"
+                f"{msg}"
+            )
+
+        # No need to await this task.
+        check_raylet_task = create_check_raylet_task(
+            args.log_dir, gcs_client, parent_dead_callback, loop
+        )
+
+    port = args.runtime_env_agent_port or 0
+    infos = socket.getaddrinfo(args.node_ip_address, port, type=socket.SOCK_STREAM)
+    family, socktype, proto, _, sockaddr = infos[0]
+    sock = socket.socket(family, socktype, proto)
+    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    sock.bind(sockaddr)
+
+    bound_port = sock.getsockname()[1]
+    persist_port(
+        args.session_dir,
+        args.node_id,
+        RUNTIME_ENV_AGENT_PORT_NAME,
+        bound_port,
+    )
+
+    try:
+        web.run_app(app, sock=sock, loop=loop)
+    except SystemExit as e:
+        agent._logger.info(f"SystemExit! {e}")
+        # We have to poke the task exception, or there's an error message
+        # "task exception was never retrieved".
+        if check_raylet_task is not None:
+            check_raylet_task.exception()
+        sys.exit(e.code)