142 lines
4.5 KiB
Python
142 lines
4.5 KiB
Python
# Copyright 2023 The Kubeflow Authors
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
"""Utilitites for formatting, coloring, and controlling the output of logs."""
|
|
import builtins
|
|
import contextlib
|
|
import datetime
|
|
import logging
|
|
import sys
|
|
from typing import Any, Dict, Generator, List
|
|
|
|
from kfp import dsl
|
|
|
|
|
|
class Color:
|
|
CYAN = '\033[96m'
|
|
GREEN = '\033[92m'
|
|
RED = '\033[91m'
|
|
RESET = '\033[0m'
|
|
|
|
|
|
class MillisecondFormatter(logging.Formatter):
|
|
|
|
def formatTime(
|
|
self,
|
|
record: logging.LogRecord,
|
|
datefmt: str = None,
|
|
) -> str:
|
|
created = datetime.datetime.fromtimestamp(record.created)
|
|
s = created.strftime(datefmt)
|
|
# truncate microseconds to milliseconds
|
|
return s[:-3]
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def local_logger_context() -> Generator[None, None, None]:
|
|
"""Context manager for creating and reseting the local execution logger."""
|
|
|
|
logger = logging.getLogger()
|
|
original_level = logger.level
|
|
original_handlers = logger.handlers[:]
|
|
formatter = MillisecondFormatter(
|
|
fmt='%(asctime)s - %(levelname)s - %(message)s',
|
|
datefmt='%H:%M:%S.%f',
|
|
)
|
|
# use sys.stdout so that both inner process and outer process logs
|
|
# go to stdout
|
|
# this is needed for logs to present sequentially in a colab notebook,
|
|
# since stderr will print above stdout
|
|
handler = logging.StreamHandler(sys.stdout)
|
|
handler.setFormatter(formatter)
|
|
logger.handlers.clear()
|
|
logger.addHandler(handler)
|
|
logger.setLevel(logging.INFO)
|
|
|
|
try:
|
|
yield
|
|
finally:
|
|
logger.setLevel(original_level)
|
|
logger.handlers.clear()
|
|
for handler in original_handlers:
|
|
logger.addHandler(handler)
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def indented_print(num_spaces: int = 4) -> Generator[None, None, None]:
|
|
"""Context manager to indent all print statements in its scope by
|
|
num_prints.
|
|
|
|
Useful for visually separating a subprocess logs from the outer
|
|
process logs.
|
|
"""
|
|
original_print = builtins.print
|
|
|
|
def indented_print_function(*args, **kwargs):
|
|
original_print(' ' * num_spaces, end='')
|
|
return original_print(*args, **kwargs)
|
|
|
|
builtins.print = indented_print_function
|
|
try:
|
|
yield
|
|
finally:
|
|
builtins.print = original_print
|
|
|
|
|
|
def color_text(text: str, color: Color) -> str:
|
|
return f'{color}{text}{Color.RESET}'
|
|
|
|
|
|
def make_log_lines_for_artifact(artifact: dsl.Artifact,) -> List[str]:
|
|
"""Returns a list of log lines that represent a single artifact output."""
|
|
artifact_class_name_and_paren = f'{artifact.__class__.__name__}( '
|
|
# name
|
|
artifact_lines = [f"{artifact_class_name_and_paren}name='{artifact.name}',"]
|
|
newline_spaces = len(artifact_class_name_and_paren) * ' '
|
|
# uri
|
|
artifact_lines.append(f"{newline_spaces}uri='{artifact.uri}',")
|
|
# metadata
|
|
artifact_lines.append(f'{newline_spaces}metadata={artifact.metadata} )')
|
|
return artifact_lines
|
|
|
|
|
|
def make_log_lines_for_outputs(outputs: Dict[str, Any]) -> List[str]:
|
|
"""Returns a list of log lines to repesent the outputs of a task."""
|
|
INDENT = ' ' * 4
|
|
SEPARATOR = ': '
|
|
output_lines = []
|
|
for key, value in outputs.items():
|
|
key_chars = INDENT + key + SEPARATOR
|
|
|
|
# present artifacts
|
|
if isinstance(value, dsl.Artifact):
|
|
artifact_lines = make_log_lines_for_artifact(value)
|
|
|
|
first_artifact_line = artifact_lines[0]
|
|
output_lines.append(f'{key_chars}{first_artifact_line}')
|
|
|
|
remaining_artifact_lines = artifact_lines[1:]
|
|
# indent to align with first char in artifact
|
|
# to visually separate output keys
|
|
remaining_artifact_lines = [
|
|
len(key_chars) * ' ' + l for l in remaining_artifact_lines
|
|
]
|
|
output_lines.extend(remaining_artifact_lines)
|
|
|
|
# present params
|
|
else:
|
|
value = f"'{value}'" if isinstance(value, str) else value
|
|
output_lines.append(f'{key_chars}{value}')
|
|
|
|
return output_lines
|