1from __future__ import annotations 2 3import re 4from dataclasses import dataclass, field 5from datetime import datetime, timedelta, UTC 6from math import floor 7from typing import TYPE_CHECKING, Optional 8 9from lava.utils.console_format import CONSOLE_LOG 10 11if TYPE_CHECKING: 12 from lava.utils.log_section import LogSectionType 13 14 15# TODO: Add section final status to assist with monitoring 16@dataclass 17class GitlabSection: 18 id: str 19 header: str 20 type: LogSectionType 21 start_collapsed: bool = False 22 suppress_end: bool = False 23 suppress_start: bool = False 24 timestamp_relative_to: Optional[datetime] = None 25 escape: str = "\x1b[0K" 26 colour: str = f"{CONSOLE_LOG['FG_CYAN']}" 27 __start_time: Optional[datetime] = field(default=None, init=False) 28 __end_time: Optional[datetime] = field(default=None, init=False) 29 30 @classmethod 31 def section_id_filter(cls, value) -> str: 32 return str(re.sub(r"[^\w_-]+", "-", value)) 33 34 def __post_init__(self): 35 self.id = self.section_id_filter(self.id) 36 37 @property 38 def has_started(self) -> bool: 39 return self.__start_time is not None 40 41 @property 42 def has_finished(self) -> bool: 43 return self.__end_time is not None 44 45 @property 46 def start_time(self) -> datetime: 47 return self.__start_time 48 49 @property 50 def end_time(self) -> Optional[datetime]: 51 return self.__end_time 52 53 def get_timestamp(self, time: datetime) -> str: 54 unix_ts = datetime.timestamp(time) 55 return str(int(unix_ts)) 56 57 def section(self, marker: str, header: str, time: datetime) -> str: 58 preamble = f"{self.escape}section_{marker}" 59 collapse = marker == "start" and self.start_collapsed 60 collapsed = "[collapsed=true]" if collapse else "" 61 section_id = f"{self.id}{collapsed}" 62 63 timestamp = self.get_timestamp(time) 64 before_header = ":".join([preamble, timestamp, section_id]) 65 if self.timestamp_relative_to: 66 delta = self.start_time - self.timestamp_relative_to 67 reltime = f"[{floor(delta.seconds / 60):02}:{(delta.seconds % 60):02}] " 68 else: 69 reltime = "" 70 colored_header = f"{self.colour}{reltime}{header}\x1b[0m" if header else "" 71 header_wrapper = "\r" + f"{self.escape}{colored_header}" 72 73 return f"{before_header}{header_wrapper}" 74 75 def __str__(self) -> str: 76 status = "NS" if not self.has_started else "F" if self.has_finished else "IP" 77 delta = self.delta_time() 78 elapsed_time = "N/A" if delta is None else str(delta) 79 return ( 80 f"GitlabSection({self.id}, {self.header}, {self.type}, " 81 f"SC={self.start_collapsed}, S={status}, ST={self.start_time}, " 82 f"ET={self.end_time}, ET={elapsed_time})" 83 ) 84 85 def __enter__(self): 86 print(self.start()) 87 return self 88 89 def __exit__(self, exc_type, exc_val, exc_tb): 90 print(self.end()) 91 92 def start(self) -> str: 93 assert not self.has_finished, "Starting an already finished section" 94 self.__start_time = datetime.now(tz=UTC) 95 return self.print_start_section() 96 97 def print_start_section(self) -> str: 98 if self.suppress_start: 99 return "" 100 return self.section(marker="start", header=self.header, time=self.__start_time) 101 102 def end(self) -> str: 103 assert self.has_started, "Ending an uninitialized section" 104 self.__end_time = datetime.now(tz=UTC) 105 assert ( 106 self.__end_time >= self.__start_time 107 ), "Section execution time will be negative" 108 return self.print_end_section() 109 110 def print_end_section(self) -> str: 111 if self.suppress_end: 112 return "" 113 return self.section(marker="end", header="", time=self.__end_time) 114 115 def delta_time(self) -> Optional[timedelta]: 116 if self.__start_time and self.__end_time: 117 return self.__end_time - self.__start_time 118 119 if self.has_started: 120 return datetime.now(tz=UTC) - self.__start_time 121 122 return None 123