diff --git a/tetra/__init__.py b/tetra/__init__.py index fd8a196d..ea97bbdb 100644 --- a/tetra/__init__.py +++ b/tetra/__init__.py @@ -1,4 +1,5 @@ -# Import all the packahges here +from .logger import setup_logging +setup_logging() from . import remote_execution_pb2, remote_execution_pb2_grpc from .core.resources.resource_manager import ResourceManager diff --git a/tetra/core/resources/resource_manager.py b/tetra/core/resources/resource_manager.py index 173b0792..b6886c93 100644 --- a/tetra/core/resources/resource_manager.py +++ b/tetra/core/resources/resource_manager.py @@ -1,3 +1,6 @@ +from tetra.logger import get_logger +log = get_logger() + import json from typing import Any, Dict from pathlib import Path @@ -36,7 +39,7 @@ async def get_or_create_resource(self, config: DeployableResource) -> str: # Check if resource already exists if resource_id in self._resources: - print(f"Resource {resource_id} already exists, reusing.") + log.debug(f"Resource {resource_id} already exists, reusing.") return self._resources[resource_id]["server_name"] # Deploy new resource based on type diff --git a/tetra/core/resources/serverless.py b/tetra/core/resources/serverless.py index cc7e3ea6..688883f1 100644 --- a/tetra/core/resources/serverless.py +++ b/tetra/core/resources/serverless.py @@ -1,3 +1,6 @@ +from tetra.logger import get_logger +log = get_logger() + from typing import Dict, List, Optional from enum import Enum from urllib.parse import urljoin @@ -82,7 +85,7 @@ async def deploy(self) -> "ServerlessResource": try: # If the resource is already deployed, return it if self.id: - print(f"Endpoint exists: {self.url}") + log.debug(f"Endpoint exists: {self.url}") return self result = runpod.create_endpoint( @@ -100,9 +103,9 @@ async def deploy(self) -> "ServerlessResource": ) if endpoint := ServerlessResource(**result): - print(f"Endpoint deployed: {endpoint.url}") + log.debug(f"Endpoint deployed: {endpoint.url}") return endpoint except Exception as e: - print(f"Endpoint failed to deploy: {e}") + log.error(f"Endpoint failed to deploy: {e}") raise diff --git a/tetra/core/utils/logger.py b/tetra/core/utils/logger.py deleted file mode 100644 index 8aa14139..00000000 --- a/tetra/core/utils/logger.py +++ /dev/null @@ -1,21 +0,0 @@ -# logger_util.py -import logging -import inspect - - -def setup_logging(level=logging.INFO, fmt=None): - if fmt is None: - fmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s" - logging.basicConfig(level=level, format=fmt) - - -def get_logger(name=None): - """ - Returns a logger. If no name is provided, it infers the caller's module name. - """ - if name is None: - # Get the caller's module name. - frame = inspect.stack()[1] - module = inspect.getmodule(frame[0]) - name = module.__name__ if module else "__main__" - return logging.getLogger(name) diff --git a/tetra/logger.py b/tetra/logger.py new file mode 100644 index 00000000..5e092a9d --- /dev/null +++ b/tetra/logger.py @@ -0,0 +1,38 @@ +import inspect +import logging +import os +import sys + + +def setup_logging(level: int = logging.INFO, stream=sys.stdout, fmt: str = None): + """ + Sets up the root logger with a stream handler and basic formatting. + Does nothing if handlers are already configured. + """ + if fmt is None: + fmt = "%(asctime)s | %(name)s | %(levelname)s | %(message)s" + + root_logger = logging.getLogger() + if not root_logger.hasHandlers(): + handler = logging.StreamHandler(stream) + handler.setFormatter(logging.Formatter(fmt)) + root_logger.setLevel(level) + root_logger.addHandler(handler) + + # Optionally allow log level override via env var + env_level = os.environ.get("LOG_LEVEL") + if env_level: + root_logger.setLevel(env_level.upper()) + + +def get_logger(name: str = None) -> logging.Logger: + """ + Returns a logger. If no name is provided, it infers the caller's module name. + """ + if name is None: + # Get the caller's module name. + frame = inspect.stack()[1] + module = inspect.getmodule(frame[0]) + name = module.__name__ if module else "__main__" + + return logging.getLogger(name)