diff --git a/examples/airflow-torch-training/callables.py b/examples/airflow-torch-training/callables.py index c9d8c9a4e..891147309 100644 --- a/examples/airflow-torch-training/callables.py +++ b/examples/airflow-torch-training/callables.py @@ -4,7 +4,7 @@ from torch_example_for_airflow import DownloadData, SimpleTrainer -logger = get_logger(name=__name__) +logger = get_logger() def bring_up_cluster_callable(**kwargs): diff --git a/runhouse/logger.py b/runhouse/logger.py index 694e29c3f..e26a1b9b9 100644 --- a/runhouse/logger.py +++ b/runhouse/logger.py @@ -1,15 +1,59 @@ import logging import os +import sys -def get_logger(name: str = __name__): - logger = logging.getLogger(name) +def get_logger(reinitialize: bool = False) -> logging.Logger: + """ + Creates and returns a logger with the specified name. - level = os.getenv("RH_LOG_LEVEL") - if level: - try: - logger.setLevel(getattr(logging, level.upper())) - except AttributeError as e: - raise e + Ensures a universal logger configuration across the codebase with the format: + "levelname - asctime - filename:lineno - message" + + Args: + name (str): Name of the logger. Defaults to None, which gets the root logger. + + Returns: + logging.Logger: Configured logger instance. + """ + # Create or retrieve the logger + logger = logging.getLogger("_rh_universal_logger") + + level = os.getenv("RH_LOG_LEVEL") or "INFO" + try: + level = getattr(logging, level.upper()) + except AttributeError as e: + raise e + + if reinitialize: + for handler in logger.handlers: + logger.removeHandler(handler) + + # Check if the logger has already been configured to prevent duplicate handlers + if not logger.handlers: + + # Set the level for the new logger + logger.setLevel(level) + + # Create a console handler that outputs to stdout + console_handler = logging.StreamHandler(stream=sys.stdout) + + # The handler should capture all log levels, so we set the level to debug + console_handler.setLevel(logging.DEBUG) + + # Define the log format + formatter = logging.Formatter( + "%(levelname)s | %(asctime)s | %(filename)s:%(lineno)d | %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + # Apply the formatter to the handler + console_handler.setFormatter(formatter) + + # Add the handler to the logger + logger.addHandler(console_handler) + + # Prevent log messages from propagating to the root logger + logger.propagate = False return logger diff --git a/runhouse/main.py b/runhouse/main.py index e0c0d4f08..45282b844 100644 --- a/runhouse/main.py +++ b/runhouse/main.py @@ -46,7 +46,7 @@ # For printing with typer console = Console() -logger = get_logger(name=__name__) +logger = get_logger() @app.command() diff --git a/runhouse/resources/envs/conda_env.py b/runhouse/resources/envs/conda_env.py index c8e5b2c90..b3261845e 100644 --- a/runhouse/resources/envs/conda_env.py +++ b/runhouse/resources/envs/conda_env.py @@ -14,7 +14,7 @@ from .env import Env -logger = get_logger(name=__name__) +logger = get_logger() class CondaEnv(Env): diff --git a/runhouse/resources/envs/env.py b/runhouse/resources/envs/env.py index 3971c0056..50c9f7edc 100644 --- a/runhouse/resources/envs/env.py +++ b/runhouse/resources/envs/env.py @@ -13,7 +13,7 @@ from runhouse.resources.resource import Resource from runhouse.utils import run_with_logs -logger = get_logger(name=__name__) +logger = get_logger() class Env(Resource): diff --git a/runhouse/resources/folders/folder.py b/runhouse/resources/folders/folder.py index 2736fa7ce..3a4ae744e 100644 --- a/runhouse/resources/folders/folder.py +++ b/runhouse/resources/folders/folder.py @@ -14,7 +14,7 @@ from runhouse.rns.utils.api import generate_uuid, relative_file_path from runhouse.utils import locate_working_dir -logger = get_logger(name=__name__) +logger = get_logger() class Folder(Resource): diff --git a/runhouse/resources/folders/folder_factory.py b/runhouse/resources/folders/folder_factory.py index b977f7d24..5da0dd987 100644 --- a/runhouse/resources/folders/folder_factory.py +++ b/runhouse/resources/folders/folder_factory.py @@ -6,7 +6,7 @@ from runhouse.resources.folders.folder import Folder from runhouse.resources.hardware.utils import _get_cluster_from -logger = get_logger(name=__name__) +logger = get_logger() def folder( diff --git a/runhouse/resources/folders/gcs_folder.py b/runhouse/resources/folders/gcs_folder.py index 0a1ebcdc8..c5a5f5bf8 100644 --- a/runhouse/resources/folders/gcs_folder.py +++ b/runhouse/resources/folders/gcs_folder.py @@ -8,7 +8,7 @@ from .folder import Folder -logger = get_logger(name=__name__) +logger = get_logger() class GCSFolder(Folder): diff --git a/runhouse/resources/folders/s3_folder.py b/runhouse/resources/folders/s3_folder.py index 33f5a2f8d..590f133b4 100644 --- a/runhouse/resources/folders/s3_folder.py +++ b/runhouse/resources/folders/s3_folder.py @@ -13,7 +13,7 @@ POLL_INTERVAL = 1 TIMEOUT_SECONDS = 3600 -logger = get_logger(name=__name__) +logger = get_logger() class S3Folder(Folder): diff --git a/runhouse/resources/functions/aws_lambda.py b/runhouse/resources/functions/aws_lambda.py index c74de599e..fe6bfe472 100644 --- a/runhouse/resources/functions/aws_lambda.py +++ b/runhouse/resources/functions/aws_lambda.py @@ -25,7 +25,7 @@ CRED_PATH = f"{Path.home()}/.aws/credentials" LOG_GROUP_PREFIX = "/aws/lambda/" -logger = get_logger(name=__name__) +logger = get_logger() class LambdaFunction(Function): diff --git a/runhouse/resources/functions/function.py b/runhouse/resources/functions/function.py index e3f0f4d43..4dc550029 100644 --- a/runhouse/resources/functions/function.py +++ b/runhouse/resources/functions/function.py @@ -11,7 +11,7 @@ from runhouse.resources.resource import Resource -logger = get_logger(name=__name__) +logger = get_logger() class Function(Module): diff --git a/runhouse/resources/functions/function_factory.py b/runhouse/resources/functions/function_factory.py index ad015cfda..08b2c06f1 100644 --- a/runhouse/resources/functions/function_factory.py +++ b/runhouse/resources/functions/function_factory.py @@ -8,7 +8,7 @@ from runhouse.resources.functions.function import Function from runhouse.resources.packages import git_package -logger = get_logger(name=__name__) +logger = get_logger() def function( diff --git a/runhouse/resources/hardware/cluster.py b/runhouse/resources/hardware/cluster.py index aa5045663..85ab63c22 100644 --- a/runhouse/resources/hardware/cluster.py +++ b/runhouse/resources/hardware/cluster.py @@ -55,7 +55,7 @@ from runhouse.servers.http import HTTPClient -logger = get_logger(name=__name__) +logger = get_logger() class Cluster(Resource): diff --git a/runhouse/resources/hardware/cluster_factory.py b/runhouse/resources/hardware/cluster_factory.py index a44454aba..044cc2513 100644 --- a/runhouse/resources/hardware/cluster_factory.py +++ b/runhouse/resources/hardware/cluster_factory.py @@ -15,7 +15,7 @@ from .on_demand_cluster import OnDemandCluster from .sagemaker.sagemaker_cluster import SageMakerCluster -logger = get_logger(name=__name__) +logger = get_logger() # Cluster factory method def cluster( diff --git a/runhouse/resources/hardware/on_demand_cluster.py b/runhouse/resources/hardware/on_demand_cluster.py index e4c3181e2..e1ba65cd0 100644 --- a/runhouse/resources/hardware/on_demand_cluster.py +++ b/runhouse/resources/hardware/on_demand_cluster.py @@ -31,7 +31,7 @@ from .cluster import Cluster -logger = get_logger(name=__name__) +logger = get_logger() class OnDemandCluster(Cluster): diff --git a/runhouse/resources/hardware/ray_utils.py b/runhouse/resources/hardware/ray_utils.py index 23451bc6a..7cf26d241 100644 --- a/runhouse/resources/hardware/ray_utils.py +++ b/runhouse/resources/hardware/ray_utils.py @@ -6,7 +6,7 @@ from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() def check_for_existing_ray_instance(): diff --git a/runhouse/resources/hardware/sagemaker/sagemaker_cluster.py b/runhouse/resources/hardware/sagemaker/sagemaker_cluster.py index e398260d1..308db5c3d 100644 --- a/runhouse/resources/hardware/sagemaker/sagemaker_cluster.py +++ b/runhouse/resources/hardware/sagemaker/sagemaker_cluster.py @@ -44,7 +44,7 @@ from runhouse.utils import generate_default_name -logger = get_logger(name=__name__) +logger = get_logger() #################################################################################################### # Caching mechanisms for SSHTunnelForwarder #################################################################################################### diff --git a/runhouse/resources/hardware/sky/command_runner.py b/runhouse/resources/hardware/sky/command_runner.py index 8bd2f7a8d..32b33f6b3 100644 --- a/runhouse/resources/hardware/sky/command_runner.py +++ b/runhouse/resources/hardware/sky/command_runner.py @@ -22,7 +22,7 @@ ##### RH modification ##### from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() ##### RH modification ##### diff --git a/runhouse/resources/hardware/sky/subprocess_utils.py b/runhouse/resources/hardware/sky/subprocess_utils.py index 4a823b4af..c11475563 100644 --- a/runhouse/resources/hardware/sky/subprocess_utils.py +++ b/runhouse/resources/hardware/sky/subprocess_utils.py @@ -5,7 +5,7 @@ from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() class CommandError(Exception): """Raised when a command fails. diff --git a/runhouse/resources/hardware/sky_ssh_runner.py b/runhouse/resources/hardware/sky_ssh_runner.py index f3c42e2fb..205389065 100644 --- a/runhouse/resources/hardware/sky_ssh_runner.py +++ b/runhouse/resources/hardware/sky_ssh_runner.py @@ -25,7 +25,7 @@ ) -logger = get_logger(name=__name__) +logger = get_logger() # Get rid of the constant "Found credentials in shared credentials file: ~/.aws/credentials" message try: diff --git a/runhouse/resources/module.py b/runhouse/resources/module.py index b8b39efff..b4da7c5d0 100644 --- a/runhouse/resources/module.py +++ b/runhouse/resources/module.py @@ -52,7 +52,7 @@ "_dumb_signature_cache", ] -logger = get_logger(name=__name__) +logger = get_logger() class Module(Resource): diff --git a/runhouse/resources/packages/package.py b/runhouse/resources/packages/package.py index 7dd3adc92..24d003b11 100644 --- a/runhouse/resources/packages/package.py +++ b/runhouse/resources/packages/package.py @@ -24,7 +24,7 @@ INSTALL_METHODS = {"local", "reqs", "pip", "conda", "rh"} -logger = get_logger(name=__name__) +logger = get_logger() class CodeSyncError(Exception): diff --git a/runhouse/resources/provenance.py b/runhouse/resources/provenance.py index 570dff16f..8c9907220 100644 --- a/runhouse/resources/provenance.py +++ b/runhouse/resources/provenance.py @@ -20,7 +20,7 @@ from runhouse.rns.utils.api import log_timestamp, resolve_absolute_path from runhouse.utils import StreamTee -logger = get_logger(name=__name__) +logger = get_logger() class RunStatus(str, Enum): diff --git a/runhouse/resources/resource.py b/runhouse/resources/resource.py index eb9456e7b..aef4ca8d5 100644 --- a/runhouse/resources/resource.py +++ b/runhouse/resources/resource.py @@ -19,7 +19,7 @@ ResourceVisibility, ) -logger = get_logger(name=__name__) +logger = get_logger() class Resource: diff --git a/runhouse/resources/secrets/provider_secrets/ssh_secret.py b/runhouse/resources/secrets/provider_secrets/ssh_secret.py index 895e05acf..40836f0bd 100644 --- a/runhouse/resources/secrets/provider_secrets/ssh_secret.py +++ b/runhouse/resources/secrets/provider_secrets/ssh_secret.py @@ -10,7 +10,7 @@ from runhouse.resources.hardware.cluster import Cluster from runhouse.resources.secrets.provider_secrets.provider_secret import ProviderSecret -logger = get_logger(name=__name__) +logger = get_logger() class SSHSecret(ProviderSecret): diff --git a/runhouse/resources/secrets/secret.py b/runhouse/resources/secrets/secret.py index b88b10779..8b5862834 100644 --- a/runhouse/resources/secrets/secret.py +++ b/runhouse/resources/secrets/secret.py @@ -14,7 +14,7 @@ from runhouse.rns.utils.api import load_resp_content, read_resp_data from runhouse.utils import generate_default_name -logger = get_logger(name=__name__) +logger = get_logger() class Secret(Resource): diff --git a/runhouse/resources/secrets/utils.py b/runhouse/resources/secrets/utils.py index d8dca10e4..d2c667e4f 100644 --- a/runhouse/resources/secrets/utils.py +++ b/runhouse/resources/secrets/utils.py @@ -11,7 +11,7 @@ USER_ENDPOINT = "user/secret" -logger = get_logger(name=__name__) +logger = get_logger() def load_config(name: str, endpoint: str = USER_ENDPOINT): diff --git a/runhouse/rns/defaults.py b/runhouse/rns/defaults.py index f22b77872..c59a02dc6 100644 --- a/runhouse/rns/defaults.py +++ b/runhouse/rns/defaults.py @@ -16,7 +16,7 @@ req_ctx = contextvars.ContextVar("rh_ctx", default={}) -logger = get_logger(name=__name__) +logger = get_logger() class Defaults: diff --git a/runhouse/rns/login.py b/runhouse/rns/login.py index 018a97338..9ab485126 100644 --- a/runhouse/rns/login.py +++ b/runhouse/rns/login.py @@ -7,7 +7,7 @@ from runhouse.globals import configs, rns_client from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() def is_interactive(): diff --git a/runhouse/rns/rns_client.py b/runhouse/rns/rns_client.py index 3e077a948..5897756e0 100644 --- a/runhouse/rns/rns_client.py +++ b/runhouse/rns/rns_client.py @@ -23,7 +23,7 @@ from runhouse.utils import locate_working_dir -logger = get_logger(name=__name__) +logger = get_logger() # This is a copy of the Pydantic model that we use to validate in Den diff --git a/runhouse/rns/top_level_rns_fns.py b/runhouse/rns/top_level_rns_fns.py index 8d92891ed..51c198519 100644 --- a/runhouse/rns/top_level_rns_fns.py +++ b/runhouse/rns/top_level_rns_fns.py @@ -8,7 +8,7 @@ from runhouse.logger import get_logger from runhouse.servers.obj_store import ClusterServletSetupOption -logger = get_logger(name=__name__) +logger = get_logger() logging.getLogger("numexpr").setLevel(logging.WARNING) diff --git a/runhouse/servers/autostop_helper.py b/runhouse/servers/autostop_helper.py index c4f06954c..f3613e6e9 100644 --- a/runhouse/servers/autostop_helper.py +++ b/runhouse/servers/autostop_helper.py @@ -3,7 +3,7 @@ from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() class AutostopHelper: diff --git a/runhouse/servers/caddy/config.py b/runhouse/servers/caddy/config.py index 3862f1612..0a20b93c0 100644 --- a/runhouse/servers/caddy/config.py +++ b/runhouse/servers/caddy/config.py @@ -5,7 +5,7 @@ from runhouse.constants import DEFAULT_SERVER_PORT from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() SYSTEMCTL_ERROR = "systemctl: command not found" diff --git a/runhouse/servers/cluster_servlet.py b/runhouse/servers/cluster_servlet.py index e9594d909..32030ba48 100644 --- a/runhouse/servers/cluster_servlet.py +++ b/runhouse/servers/cluster_servlet.py @@ -27,7 +27,7 @@ from runhouse.servers.http.auth import AuthCache from runhouse.utils import ColoredFormatter, sync_function -logger = get_logger(name=__name__) +logger = get_logger() class ClusterServletError(Exception): diff --git a/runhouse/servers/env_servlet.py b/runhouse/servers/env_servlet.py index 0ededb46e..ff418343f 100644 --- a/runhouse/servers/env_servlet.py +++ b/runhouse/servers/env_servlet.py @@ -16,7 +16,7 @@ from runhouse.servers.obj_store import ClusterServletSetupOption from runhouse.utils import arun_in_thread, get_node_ip -logger = get_logger(name=__name__) +logger = get_logger() def error_handling_decorator(func): diff --git a/runhouse/servers/http/auth.py b/runhouse/servers/http/auth.py index b61f34fef..1454d2929 100644 --- a/runhouse/servers/http/auth.py +++ b/runhouse/servers/http/auth.py @@ -5,7 +5,7 @@ from runhouse.rns.utils.api import load_resp_content, ResourceAccess from runhouse.servers.http.http_utils import username_from_token -logger = get_logger(name=__name__) +logger = get_logger() class AuthCache: diff --git a/runhouse/servers/http/certs.py b/runhouse/servers/http/certs.py index 971526f81..60fc78a7a 100644 --- a/runhouse/servers/http/certs.py +++ b/runhouse/servers/http/certs.py @@ -12,7 +12,7 @@ from runhouse.rns.utils.api import resolve_absolute_path -logger = get_logger(name=__name__) +logger = get_logger() class TLSCertConfig: diff --git a/runhouse/servers/http/http_client.py b/runhouse/servers/http/http_client.py index 7d4e09261..f951a0dd8 100644 --- a/runhouse/servers/http/http_client.py +++ b/runhouse/servers/http/http_client.py @@ -41,7 +41,7 @@ session = requests.Session() session.timeout = None -logger = get_logger(name=__name__) +logger = get_logger() def retry_with_exponential_backoff(func): diff --git a/runhouse/servers/http/http_server.py b/runhouse/servers/http/http_server.py index 100405c59..79b7eba60 100644 --- a/runhouse/servers/http/http_server.py +++ b/runhouse/servers/http/http_server.py @@ -68,7 +68,7 @@ app = FastAPI(docs_url=None, redoc_url=None) -logger = get_logger(name=__name__) +logger = get_logger() def validate_cluster_access(func): diff --git a/runhouse/servers/http/http_utils.py b/runhouse/servers/http/http_utils.py index b9dca7be8..1e18c3d93 100644 --- a/runhouse/servers/http/http_utils.py +++ b/runhouse/servers/http/http_utils.py @@ -18,7 +18,7 @@ from runhouse.servers.obj_store import RunhouseStopIteration from runhouse.utils import ClusterLogsFormatter -logger = get_logger(name=__name__) +logger = get_logger() class RequestContext(BaseModel): diff --git a/runhouse/servers/obj_store.py b/runhouse/servers/obj_store.py index b87f8c9c3..9b2c67e26 100644 --- a/runhouse/servers/obj_store.py +++ b/runhouse/servers/obj_store.py @@ -23,7 +23,7 @@ sync_function, ) -logger = get_logger(name=__name__) +logger = get_logger() class RaySetupOption(str, Enum): diff --git a/runhouse/utils.py b/runhouse/utils.py index e698d52fe..cc1b5c16d 100644 --- a/runhouse/utils.py +++ b/runhouse/utils.py @@ -16,7 +16,6 @@ ) import inspect import json -import logging import os import re import subprocess @@ -33,7 +32,7 @@ from runhouse.constants import LOGS_DIR from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() #################################################################################################### # Python package utilities #################################################################################################### @@ -379,7 +378,6 @@ class LogToFolder: def __init__(self, name: str): self.name = name self.directory = self._base_local_folder_path(name) - self.root_logger = logging.getLogger("") # We do exist_ok=True here because generator runs are separate calls to the same directory. os.makedirs(self.directory, exist_ok=True) @@ -388,14 +386,12 @@ def __enter__(self): sys.stdout = StreamTee(sys.stdout, [Path(self._stdout_path).open(mode="a")]) sys.stderr = StreamTee(sys.stderr, [Path(self._stderr_path).open(mode="a")]) - # Add the stdout and stderr handlers to the root logger - self._stdout_handler = logging.StreamHandler(sys.stdout) - self.root_logger.addHandler(self._stdout_handler) + # Reinitialize the universal logger that we're using, so that it points to the new fake sys.stdout + get_logger(reinitialize=True) return self def __exit__(self, exc_type, exc_val, exc_tb): - self.root_logger.removeHandler(self._stdout_handler) # Flush stdout and stderr # sys.stdout.flush() @@ -407,6 +403,9 @@ def __exit__(self, exc_type, exc_val, exc_tb): if hasattr(sys.stderr, "instream"): sys.stderr = sys.stderr.instream + # Reinitialize this again to be the original sys.stdout + get_logger(reinitialize=True) + # return False to propagate any exception that occurred inside the with block return False diff --git a/tests/test_obj_store.py b/tests/test_obj_store.py index 842422fcf..4d518665a 100644 --- a/tests/test_obj_store.py +++ b/tests/test_obj_store.py @@ -14,7 +14,7 @@ TEMP_FILE = "my_file.txt" TEMP_FOLDER = "~/runhouse-tests" -logger = get_logger(name=__name__) +logger = get_logger() UNIT = {"cluster": []} LOCAL = { diff --git a/tests/test_performance.py b/tests/test_performance.py index 1dc0e5faa..7106b3a9e 100644 --- a/tests/test_performance.py +++ b/tests/test_performance.py @@ -5,7 +5,7 @@ from runhouse.globals import rns_client from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() def profile(func, reps=10): diff --git a/tests/test_resources/test_modules/test_functions/test_function.py b/tests/test_resources/test_modules/test_functions/test_function.py index e79e8bafc..bb9e02d75 100644 --- a/tests/test_resources/test_modules/test_functions/test_function.py +++ b/tests/test_resources/test_modules/test_functions/test_function.py @@ -11,7 +11,7 @@ from tests.utils import friend_account -logger = get_logger(name=__name__) +logger = get_logger() def get_remote_func_name(test_folder): diff --git a/tests/test_resources/test_modules/test_module.py b/tests/test_resources/test_modules/test_module.py index d01831685..24dc0ddca 100644 --- a/tests/test_resources/test_modules/test_module.py +++ b/tests/test_resources/test_modules/test_module.py @@ -15,7 +15,7 @@ from runhouse.constants import TEST_ORG from runhouse.logger import get_logger -logger = get_logger(name=__name__) +logger = get_logger() """ Tests for runhouse.Module. Structure: - Test call_module_method rpc, with various envs diff --git a/tests/test_servers/conftest.py b/tests/test_servers/conftest.py index fb836c4fe..6e8dc3a1b 100644 --- a/tests/test_servers/conftest.py +++ b/tests/test_servers/conftest.py @@ -16,7 +16,7 @@ from tests.utils import friend_account, get_ray_servlet_and_obj_store -logger = get_logger(name=__name__) +logger = get_logger() # -------- HELPERS ----------- # def summer(a, b):