fix: ensure logs are disabled by default

This commit is contained in:
bennykok 2024-01-30 15:17:25 +08:00
parent 81bde40aeb
commit 2193dd287d
2 changed files with 49 additions and 34 deletions

View File

@ -26,6 +26,7 @@ import threading
api = None
api_task = None
prompt_metadata = {}
cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'
def post_prompt(json_data):
prompt_server = server.PromptServer.instance
@ -157,7 +158,9 @@ async def websocket_handler(request):
try:
# Send initial state to the new client
await send("status", { 'sid': sid }, sid)
await send_first_time_log(sid)
if cd_enable_log:
await send_first_time_log(sid)
async for msg in ws:
if msg.type == aiohttp.WSMsgType.ERROR:
@ -480,4 +483,5 @@ def run_in_new_thread(coroutine):
t.start()
asyncio.run_coroutine_threadsafe(coroutine, new_loop)
run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket))
if cd_enable_log:
run_in_new_thread(watch_file_changes(log_file_path, send_logs_to_websocket))

View File

@ -7,45 +7,56 @@ import threading
import logging
from logging.handlers import RotatingFileHandler
handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)
# Running with export CD_ENABLE_LOG=true; python main.py
original_stdout = sys.stdout
original_stderr = sys.stderr
# Check for 'cd-enable-log' flag in input arguments
# cd_enable_log = '--cd-enable-log' in sys.argv
cd_enable_log = os.environ.get('CD_ENABLE_LOG', 'false').lower() == 'true'
class StreamToLogger():
def __init__(self, log_level):
self.log_level = log_level
def setup():
handler = RotatingFileHandler('comfy-deploy.log', maxBytes=500000, backupCount=5)
def write(self, buf):
if (self.log_level == logging.INFO):
original_stdout.write(buf)
original_stdout.flush()
elif (self.log_level == logging.ERROR):
original_stderr.write(buf)
original_stderr.flush()
original_stdout = sys.stdout
original_stderr = sys.stderr
for line in buf.rstrip().splitlines():
handler.handle(
logging.LogRecord(
name="comfy-deploy",
level=self.log_level,
pathname="prestartup_script.py",
lineno=1,
msg=line.rstrip(),
args=None,
exc_info=None
class StreamToLogger():
def __init__(self, log_level):
self.log_level = log_level
def write(self, buf):
if (self.log_level == logging.INFO):
original_stdout.write(buf)
original_stdout.flush()
elif (self.log_level == logging.ERROR):
original_stderr.write(buf)
original_stderr.flush()
for line in buf.rstrip().splitlines():
handler.handle(
logging.LogRecord(
name="comfy-deploy",
level=self.log_level,
pathname="prestartup_script.py",
lineno=1,
msg=line.rstrip(),
args=None,
exc_info=None
)
)
)
def flush(self):
if (self.log_level == logging.INFO):
original_stdout.flush()
elif (self.log_level == logging.ERROR):
original_stderr.flush()
def flush(self):
if (self.log_level == logging.INFO):
original_stdout.flush()
elif (self.log_level == logging.ERROR):
original_stderr.flush()
# Redirect stdout and stderr to the logger
sys.stdout = StreamToLogger(logging.INFO)
sys.stderr = StreamToLogger(logging.ERROR)
# Redirect stdout and stderr to the logger
sys.stdout = StreamToLogger(logging.INFO)
sys.stderr = StreamToLogger(logging.ERROR)
if cd_enable_log:
print("** Comfy Deploy logging enabled")
setup()
try:
# Get the absolute path of the script's directory