From d3c1f8c8b77e5d3b4eee856c6e92b5da31a70b08 Mon Sep 17 00:00:00 2001
From: Andrew Aikawa <asai@berkeley.edu>
Date: Sun, 30 Jun 2024 10:51:00 -0700
Subject: [PATCH] [k8s] suppress connection error warnings when disconnected
 from k8s (#3674)

* suppress connection error warnings when disconnected from k8s

format

set urllib3 log level

set to ERROR level

* Update sky/provision/kubernetes/utils.py

Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com>

* format

* decorate k8s apis to suppress logs

* Add docstr

* lint

---------

Co-authored-by: Zhanghao Wu <zhanghao.wu@outlook.com>
Co-authored-by: Romil Bhardwaj <romil.bhardwaj@gmail.com>
---
 sky/adaptors/kubernetes.py | 37 ++++++++++++++++++++++++++++++++++++-
 sky/sky_logging.py         | 11 +++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/sky/adaptors/kubernetes.py b/sky/adaptors/kubernetes.py
index 7cdb3ff3059..7f52a099f56 100644
--- a/sky/adaptors/kubernetes.py
+++ b/sky/adaptors/kubernetes.py
@@ -2,9 +2,11 @@
 
 # pylint: disable=import-outside-toplevel
 
+import logging
 import os
 
 from sky.adaptors import common
+from sky.sky_logging import set_logging_level
 from sky.utils import env_options
 from sky.utils import ux_utils
 
@@ -28,6 +30,33 @@
 API_TIMEOUT = 5
 
 
+def _decorate_methods(obj, decorator):
+    for attr_name in dir(obj):
+        attr = getattr(obj, attr_name)
+        if callable(attr) and not attr_name.startswith('__'):
+            setattr(obj, attr_name, decorator(attr))
+    return obj
+
+
+def _api_logging_decorator(logger: str, level: int):
+    """Decorator to set logging level for API calls.
+
+    This is used to suppress the verbose logging from urllib3 when calls to the
+    Kubernetes API timeout.
+    """
+
+    def decorated_api(api):
+
+        def wrapped(*args, **kwargs):
+            obj = api(*args, **kwargs)
+            _decorate_methods(obj, set_logging_level(logger, level))
+            return obj
+
+        return wrapped
+
+    return decorated_api
+
+
 def _load_config():
     global _configured
     if _configured:
@@ -65,15 +94,16 @@ def _load_config():
     _configured = True
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def core_api():
     global _core_api
     if _core_api is None:
         _load_config()
         _core_api = kubernetes.client.CoreV1Api()
-
     return _core_api
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def auth_api():
     global _auth_api
     if _auth_api is None:
@@ -83,6 +113,7 @@ def auth_api():
     return _auth_api
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def networking_api():
     global _networking_api
     if _networking_api is None:
@@ -92,6 +123,7 @@ def networking_api():
     return _networking_api
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def custom_objects_api():
     global _custom_objects_api
     if _custom_objects_api is None:
@@ -101,6 +133,7 @@ def custom_objects_api():
     return _custom_objects_api
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def node_api():
     global _node_api
     if _node_api is None:
@@ -110,6 +143,7 @@ def node_api():
     return _node_api
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def apps_api():
     global _apps_api
     if _apps_api is None:
@@ -119,6 +153,7 @@ def apps_api():
     return _apps_api
 
 
+@_api_logging_decorator('urllib3', logging.ERROR)
 def api_client():
     global _api_client
     if _api_client is None:
diff --git a/sky/sky_logging.py b/sky/sky_logging.py
index dbaf1dd0479..c8a243c72cf 100644
--- a/sky/sky_logging.py
+++ b/sky/sky_logging.py
@@ -95,6 +95,17 @@ def init_logger(name: str):
     return logging.getLogger(name)
 
 
+@contextlib.contextmanager
+def set_logging_level(logger: str, level: int):
+    logger = logging.getLogger(logger)
+    original_level = logger.level
+    logger.setLevel(level)
+    try:
+        yield
+    finally:
+        logger.setLevel(original_level)
+
+
 @contextlib.contextmanager
 def silent():
     """Make all sky_logging.print() and logger.{info, warning...} silent.