diff --git a/gigl/common/logger.py b/gigl/common/logger.py index 3dd98bb81..67184b7e2 100644 --- a/gigl/common/logger.py +++ b/gigl/common/logger.py @@ -6,6 +6,8 @@ from google.cloud import logging as google_cloud_logging +from gigl.env.constants import GIGL_DEBUG, is_env_flag_enabled + _BASE_LOG_FILE_PATH = "/tmp/research/gbml/logs" @@ -19,6 +21,8 @@ class Logger(logging.LoggerAdapter): extra (Optional[dict[str, Any]]): Extra information to be added to the log message. """ + _DID_ALERT_FOR_LOG_LEVEL: bool = False + def __init__( self, logger: Optional[logging.Logger] = None, @@ -26,14 +30,29 @@ def __init__( log_to_file: bool = False, extra: Optional[dict[str, Any]] = None, ): + gigl_debug = is_env_flag_enabled(GIGL_DEBUG) + if gigl_debug: + log_level = logging.DEBUG + else: + log_level = logging.INFO + if logger is None: logger = logging.getLogger(name) - self._setup_logger(logger, name, log_to_file) + self._setup_logger(logger, name, log_to_file, log_level) super().__init__(logger, extra or {}) + if not Logger._DID_ALERT_FOR_LOG_LEVEL: + Logger._DID_ALERT_FOR_LOG_LEVEL = True + level_name = logging.getLevelName(log_level) + self.info(f"{GIGL_DEBUG}={gigl_debug}, using log level {level_name}") + def _setup_logger( - self, logger: logging.Logger, name: Optional[str], log_to_file: bool + self, + logger: logging.Logger, + name: Optional[str], + log_to_file: bool, + log_level: int, ) -> None: handler: logging.Handler if not logger.handlers: @@ -42,7 +61,7 @@ def _setup_logger( ): # Google Cloud Logging client = google_cloud_logging.Client() - client.setup_logging(log_level=logging.INFO) + client.setup_logging(log_level=log_level) else: # Logging locally. Set up logging to console or file if log_to_file: @@ -61,7 +80,7 @@ def _setup_logger( ) handler.setFormatter(formatter) logger.addHandler(handler) - logger.setLevel(logging.INFO) + logger.setLevel(log_level) def process(self, msg: str, kwargs: MutableMapping[str, Any]) -> Any: if "extra" in kwargs: diff --git a/gigl/distributed/base_dist_loader.py b/gigl/distributed/base_dist_loader.py index 919d939a0..79e5c8b98 100644 --- a/gigl/distributed/base_dist_loader.py +++ b/gigl/distributed/base_dist_loader.py @@ -8,7 +8,6 @@ - Graph Store mode: barrier loop + async RPC dispatch + channel creation """ -import os import sys import time from collections import Counter @@ -62,7 +61,7 @@ patch_fanout_for_sampling, strip_non_ppr_edge_types, ) -from gigl.env.constants import GIGL_ENABLE_PERF_MONITORING +from gigl.env.constants import GIGL_ENABLE_PERF_MONITORING, is_env_flag_enabled from gigl.src.common.utils.metrics_service_provider import get_metrics_service_instance from gigl.types.graph import DEFAULT_HOMOGENEOUS_NODE_TYPE from gigl.utils.share_memory import share_memory @@ -452,10 +451,7 @@ def create_colocated_channel( RuntimeError: If GIGL_ENABLE_PERF_MONITORING is set but user did not previously call `gigl.src.common.utils.metrics_service_provider.initialize_metrics()`. """ - enable_channel_monitoring = os.environ.get( - GIGL_ENABLE_PERF_MONITORING, "" - ).strip().lower() in ("1", "True") - if enable_channel_monitoring: + if is_env_flag_enabled(GIGL_ENABLE_PERF_MONITORING): logger.info(f"{GIGL_ENABLE_PERF_MONITORING} set, using MonitoredShmChannel") try: get_metrics_service_instance() diff --git a/gigl/env/constants.py b/gigl/env/constants.py index 9b9633b15..425bce1b6 100644 --- a/gigl/env/constants.py +++ b/gigl/env/constants.py @@ -24,6 +24,7 @@ GIGL_COMPONENT_ENV_KEY: Final[str] = "GIGL_COMPONENT" GIGL_BIGQUERY_QUOTA_PROJECT_ENV_KEY: Final[str] = "GIGL_BIGQUERY_QUOTA_PROJECT" GIGL_ENABLE_PERF_MONITORING: Final[str] = "GIGL_ENABLE_PERF_MONITORING" +GIGL_DEBUG: Final[str] = "GIGL_DEBUG" def read_resource_config_uri_from_env() -> Optional[str]: @@ -33,3 +34,8 @@ def read_resource_config_uri_from_env() -> Optional[str]: The URI string if the variable is set, else ``None``. """ return os.environ.get(GIGL_RESOURCE_CONFIG_URI_ENV_KEY) + + +def is_env_flag_enabled(env_var_name: str) -> bool: + """Parses an environment variable and returns True iff the variable is set to '1' or 'True' (case-insensitive).""" + return os.environ.get(env_var_name, "").strip().lower() in ("1", "true") diff --git a/tests/unit/env/constants.py b/tests/unit/env/constants.py new file mode 100644 index 000000000..d138b184a --- /dev/null +++ b/tests/unit/env/constants.py @@ -0,0 +1,32 @@ +import os +from unittest.mock import patch + +from gigl.env.constants import is_env_flag_enabled +from tests.test_assets.test_case import TestCase + + +class EnvFlagEnabledTest(TestCase): + def test_env_flag_disabled(self): + env_var = "TEST_ENV_VAR" + with patch.dict(os.environ, {}): + self.assertFalse(is_env_flag_enabled(env_var)) + + def test_env_flag_enabled_1(self): + env_var = "TEST_ENV_VAR" + with patch.dict(os.environ, {env_var: "1"}): + self.assertTrue(is_env_flag_enabled(env_var)) + + def test_env_flag_enabled_True(self): + env_var = "TEST_ENV_VAR" + with patch.dict(os.environ, {env_var: "True"}): + self.assertTrue(is_env_flag_enabled(env_var)) + + def test_env_flag_enabled_true(self): + env_var = "TEST_ENV_VAR" + with patch.dict(os.environ, {env_var: "true"}): + self.assertTrue(is_env_flag_enabled(env_var)) + + def test_env_flag_enabled_true_with_whitespace(self): + env_var = "TEST_ENV_VAR" + with patch.dict(os.environ, {env_var: " true "}): + self.assertTrue(is_env_flag_enabled(env_var))