[Dashboard] Add GET /log_proxy API (#13165)

dashboard/modules/log/log_agent.py

@@ -1,7 +1,6 @@
 import logging
 
-import mimetypes
-
+import ray.new_dashboard.modules.log.log_utils as log_utils
 import ray.new_dashboard.utils as dashboard_utils
 
 logger = logging.getLogger(__name__)
@@ -11,8 +10,7 @@ routes = dashboard_utils.ClassMethodRouteTable
 class LogAgent(dashboard_utils.DashboardAgentModule):
     def __init__(self, dashboard_agent):
         super().__init__(dashboard_agent)
-        mimetypes.add_type("text/plain", ".err")
-        mimetypes.add_type("text/plain", ".out")
+        log_utils.register_mimetypes()
         routes.static("/logs", self._dashboard_agent.log_dir, show_index=True)
 
     async def run(self, server):

dashboard/modules/log/log_consts.py

@@ -0,0 +1,3 @@
+MIME_TYPES = {
+    "text/plain": [".err", ".out", ".log"],
+}

dashboard/modules/log/log_head.py

@@ -1,8 +1,7 @@
 import logging
 
-import mimetypes
-
 import aiohttp.web
+import ray.new_dashboard.modules.log.log_utils as log_utils
 import ray.new_dashboard.utils as dashboard_utils
 from ray.new_dashboard.datacenter import DataSource, GlobalSignals
 
@@ -15,8 +14,9 @@ class LogHead(dashboard_utils.DashboardHeadModule):
 
     def __init__(self, dashboard_head):
         super().__init__(dashboard_head)
-        mimetypes.add_type("text/plain", ".err")
-        mimetypes.add_type("text/plain", ".out")
+        # We disable auto_decompress when forward / proxy log url.
+        self._proxy_session = aiohttp.ClientSession(auto_decompress=False)
+        log_utils.register_mimetypes()
         routes.static("/logs", self._dashboard_head.log_dir, show_index=True)
         GlobalSignals.node_info_fetched.append(
             self.insert_log_url_to_node_info)
@@ -52,6 +52,26 @@ class LogHead(dashboard_utils.DashboardHeadModule):
         return aiohttp.web.Response(
             text=self._directory_as_html(url_list), content_type="text/html")
 
+    @routes.get("/log_proxy")
+    async def get_log_from_proxy(self, req) -> aiohttp.web.StreamResponse:
+        url = req.query.get("url")
+        if not url:
+            raise Exception("url is None.")
+        body = await req.read()
+        async with self._proxy_session.request(
+                req.method, url, data=body, headers=req.headers) as r:
+            sr = aiohttp.web.StreamResponse(
+                status=r.status, reason=r.reason, headers=req.headers)
+            sr.content_length = r.content_length
+            sr.content_type = r.content_type
+            sr.charset = r.charset
+
+            writer = await sr.prepare(req)
+            async for data in r.content.iter_any():
+                await writer.write(data)
+
+            return sr
+
     @staticmethod
     def _directory_as_html(url_list) -> str:
         # returns directory's index as html

dashboard/modules/log/log_utils.py

@@ -0,0 +1,8 @@
+import mimetypes
+import ray.new_dashboard.modules.log.log_consts as log_consts
+
+
+def register_mimetypes():
+    for _type, extensions in log_consts.MIME_TYPES.items():
+        for ext in extensions:
+            mimetypes.add_type(_type, ext)

dashboard/modules/log/tests/test_log.py

@@ -107,5 +107,39 @@ def test_log(disable_aiohttp_cache, ray_start_with_dashboard):
                 raise Exception(f"Timed out while testing, {ex_stack}")
 
 
+def test_log_proxy(ray_start_with_dashboard):
+    assert (wait_until_server_available(ray_start_with_dashboard["webui_url"])
+            is True)
+    webui_url = ray_start_with_dashboard["webui_url"]
+    webui_url = format_web_url(webui_url)
+
+    timeout_seconds = 5
+    start_time = time.time()
+    last_ex = None
+    while True:
+        time.sleep(1)
+        try:
+            # Test range request.
+            response = requests.get(
+                f"{webui_url}/log_proxy?url={webui_url}/logs/dashboard.log",
+                headers={"Range": "bytes=43-51"})
+            response.raise_for_status()
+            assert response.text == "Dashboard"
+            # Test 404.
+            response = requests.get(f"{webui_url}/log_proxy?"
+                                    f"url={webui_url}/logs/not_exist_file.log")
+            assert response.status_code == 404
+            break
+        except Exception as ex:
+            last_ex = ex
+        finally:
+            if time.time() > start_time + timeout_seconds:
+                ex_stack = traceback.format_exception(
+                    type(last_ex), last_ex,
+                    last_ex.__traceback__) if last_ex else []
+                ex_stack = "".join(ex_stack)
+                raise Exception(f"Timed out while testing, {ex_stack}")
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main(["-v", __file__]))