# mypy: allow-untyped-defs """A diagnostic context based on SARIF.""" from __future__ import annotations import contextlib import dataclasses import gzip import logging from typing import Callable, Generic, Literal, TYPE_CHECKING, TypeVar from typing_extensions import Self from torch.onnx._internal.diagnostics import infra from torch.onnx._internal.diagnostics.infra import formatter, sarif, utils from torch.onnx._internal.diagnostics.infra.sarif import version as sarif_version if TYPE_CHECKING: from collections.abc import Generator, Mapping # This is a workaround for mypy not supporting Self from typing_extensions. _Diagnostic = TypeVar("_Diagnostic", bound="Diagnostic") diagnostic_logger: logging.Logger = logging.getLogger(__name__) @dataclasses.dataclass class Diagnostic: rule: infra.Rule level: infra.Level message: str | None = None locations: list[infra.Location] = dataclasses.field(default_factory=list) stacks: list[infra.Stack] = dataclasses.field(default_factory=list) graphs: list[infra.Graph] = dataclasses.field(default_factory=list) thread_flow_locations: list[infra.ThreadFlowLocation] = dataclasses.field( default_factory=list ) additional_messages: list[str] = dataclasses.field(default_factory=list) tags: list[infra.Tag] = dataclasses.field(default_factory=list) source_exception: Exception | None = None """The exception that caused this diagnostic to be created.""" logger: logging.Logger = dataclasses.field(init=False, default=diagnostic_logger) """The logger for this diagnostic. Defaults to 'diagnostic_logger' which has the same log level setting with `DiagnosticOptions.verbosity_level`.""" _current_log_section_depth: int = 0 def __post_init__(self) -> None: pass def sarif(self) -> sarif.Result: """Returns the SARIF Result representation of this diagnostic.""" message = self.message or self.rule.message_default_template if self.additional_messages: additional_message = "\n".join(self.additional_messages) message_markdown = ( f"{message}\n\n## Additional Message:\n\n{additional_message}" ) else: message_markdown = message kind: Literal["informational", "fail"] = ( "informational" if self.level == infra.Level.NONE else "fail" ) sarif_result = sarif.Result( message=sarif.Message(text=message, markdown=message_markdown), level=self.level.name.lower(), # type: ignore[arg-type] rule_id=self.rule.id, kind=kind, ) sarif_result.locations = [location.sarif() for location in self.locations] sarif_result.stacks = [stack.sarif() for stack in self.stacks] sarif_result.graphs = [graph.sarif() for graph in self.graphs] sarif_result.code_flows = [ sarif.CodeFlow( thread_flows=[ sarif.ThreadFlow( locations=[loc.sarif() for loc in self.thread_flow_locations] ) ] ) ] sarif_result.properties = sarif.PropertyBag( tags=[tag.value for tag in self.tags] ) return sarif_result def with_location(self: Self, location: infra.Location) -> Self: """Adds a location to the diagnostic.""" self.locations.append(location) return self def with_thread_flow_location( self: Self, location: infra.ThreadFlowLocation ) -> Self: """Adds a thread flow location to the diagnostic.""" self.thread_flow_locations.append(location) return self def with_stack(self: Self, stack: infra.Stack) -> Self: """Adds a stack to the diagnostic.""" self.stacks.append(stack) return self def with_graph(self: Self, graph: infra.Graph) -> Self: """Adds a graph to the diagnostic.""" self.graphs.append(graph) return self @contextlib.contextmanager def log_section( self, level: int, message: str, *args, **kwargs ) -> Generator[None, None, None]: """ Context manager for a section of log messages, denoted by a title message and increased indentation. Same api as `logging.Logger.log`. This context manager logs the given title at the specified log level, increases the current section depth for subsequent log messages, and ensures that the section depth is decreased again when exiting the context. Args: level: The log level. message: The title message to log. *args: The arguments to the message. Use `LazyString` to defer the expensive evaluation of the arguments until the message is actually logged. **kwargs: The keyword arguments for `logging.Logger.log`. Yields: None: This context manager does not yield any value. Example: >>> with DiagnosticContext("DummyContext", "1.0"): ... rule = infra.Rule("RuleID", "DummyRule", "Rule message") ... diagnostic = Diagnostic(rule, infra.Level.WARNING) ... with diagnostic.log_section(logging.INFO, "My Section"): ... diagnostic.log(logging.INFO, "My Message") ... with diagnostic.log_section(logging.INFO, "My Subsection"): ... diagnostic.log(logging.INFO, "My Submessage") ... diagnostic.additional_messages ['## My Section', 'My Message', '### My Subsection', 'My Submessage'] """ if self.logger.isEnabledFor(level): indented_format_message = ( f"##{'#' * self._current_log_section_depth} {message}" ) self.log( level, indented_format_message, *args, **kwargs, ) self._current_log_section_depth += 1 try: yield finally: self._current_log_section_depth -= 1 def log(self, level: int, message: str, *args, **kwargs) -> None: """Logs a message within the diagnostic. Same api as `logging.Logger.log`. If logger is not enabled for the given level, the message will not be logged. Otherwise, the message will be logged and also added to the diagnostic's additional_messages. The default setting for `DiagnosticOptions.verbosity_level` is `logging.INFO`. Based on this default, the log level recommendations are as follows. If you've set a different default verbosity level in your application, please adjust accordingly: - logging.ERROR: Log any events leading to application failure. - logging.WARNING: Log events that might result in application issues or failures, although not guaranteed. - logging.INFO: Log general useful information, ensuring minimal performance overhead. - logging.DEBUG: Log detailed debug information, which might affect performance when logged. Args: level: The log level. message: The message to log. *args: The arguments to the message. Use `LazyString` to defer the expensive evaluation of the arguments until the message is actually logged. **kwargs: The keyword arguments for `logging.Logger.log`. """ if self.logger.isEnabledFor(level): formatted_message = message % args self.logger.log(level, formatted_message, **kwargs) self.additional_messages.append(formatted_message) def debug(self, message: str, *args, **kwargs) -> None: """Logs a debug message within the diagnostic. Same api as logging.Logger.debug. Checkout `log` for more details. """ self.log(logging.DEBUG, message, *args, **kwargs) def info(self, message: str, *args, **kwargs) -> None: """Logs an info message within the diagnostic. Same api as logging.Logger.info. Checkout `log` for more details. """ self.log(logging.INFO, message, *args, **kwargs) def warning(self, message: str, *args, **kwargs) -> None: """Logs a warning message within the diagnostic. Same api as logging.Logger.warning. Checkout `log` for more details. """ self.log(logging.WARNING, message, *args, **kwargs) def error(self, message: str, *args, **kwargs) -> None: """Logs an error message within the diagnostic. Same api as logging.Logger.error. Checkout `log` for more details. """ self.log(logging.ERROR, message, *args, **kwargs) def log_source_exception(self, level: int, exception: Exception) -> None: """Logs a source exception within the diagnostic. Invokes `log_section` and `log` to log the exception in markdown section format. """ self.source_exception = exception with self.log_section(level, "Exception log"): self.log(level, "%s", formatter.lazy_format_exception(exception)) def record_python_call_stack(self, frames_to_skip: int) -> infra.Stack: """Records the current Python call stack.""" frames_to_skip += 1 # Skip this function. stack = utils.python_call_stack(frames_to_skip=frames_to_skip) self.with_stack(stack) if len(stack.frames) > 0: self.with_location(stack.frames[0].location) return stack def record_python_call( self, fn: Callable, state: Mapping[str, str], message: str | None = None, frames_to_skip: int = 0, ) -> infra.ThreadFlowLocation: """Records a python call as one thread flow step.""" frames_to_skip += 1 # Skip this function. stack = utils.python_call_stack(frames_to_skip=frames_to_skip, frames_to_log=5) location = utils.function_location(fn) location.message = message # Add function location to the top of the stack. stack.frames.insert(0, infra.StackFrame(location=location)) thread_flow_location = infra.ThreadFlowLocation( location=location, state=state, index=len(self.thread_flow_locations), stack=stack, ) self.with_thread_flow_location(thread_flow_location) return thread_flow_location class RuntimeErrorWithDiagnostic(RuntimeError): """Runtime error with enclosed diagnostic information.""" def __init__(self, diagnostic: Diagnostic): super().__init__(diagnostic.message) self.diagnostic = diagnostic @dataclasses.dataclass class DiagnosticContext(Generic[_Diagnostic]): name: str version: str options: infra.DiagnosticOptions = dataclasses.field( default_factory=infra.DiagnosticOptions ) diagnostics: list[_Diagnostic] = dataclasses.field(init=False, default_factory=list) # TODO(bowbao): Implement this. # _invocation: infra.Invocation = dataclasses.field(init=False) _inflight_diagnostics: list[_Diagnostic] = dataclasses.field( init=False, default_factory=list ) _previous_log_level: int = dataclasses.field(init=False, default=logging.WARNING) logger: logging.Logger = dataclasses.field(init=False, default=diagnostic_logger) _bound_diagnostic_type: type = dataclasses.field(init=False, default=Diagnostic) def __enter__(self): self._previous_log_level = self.logger.level self.logger.setLevel(self.options.verbosity_level) return self def __exit__(self, exc_type, exc_val, exc_tb): self.logger.setLevel(self._previous_log_level) return None def sarif(self) -> sarif.Run: """Returns the SARIF Run object.""" unique_rules = {diagnostic.rule for diagnostic in self.diagnostics} return sarif.Run( sarif.Tool( driver=sarif.ToolComponent( name=self.name, version=self.version, rules=[rule.sarif() for rule in unique_rules], ) ), results=[diagnostic.sarif() for diagnostic in self.diagnostics], ) def sarif_log(self) -> sarif.SarifLog: # type: ignore[name-defined] """Returns the SARIF Log object.""" return sarif.SarifLog( version=sarif_version.SARIF_VERSION, schema_uri=sarif_version.SARIF_SCHEMA_LINK, runs=[self.sarif()], ) def to_json(self) -> str: return formatter.sarif_to_json(self.sarif_log()) def dump(self, file_path: str, compress: bool = False) -> None: """Dumps the SARIF log to a file.""" if compress: with gzip.open(file_path, "wt") as f: f.write(self.to_json()) else: with open(file_path, "w") as f: f.write(self.to_json()) def log(self, diagnostic: _Diagnostic) -> None: """Logs a diagnostic. This method should be used only after all the necessary information for the diagnostic has been collected. Args: diagnostic: The diagnostic to add. """ if not isinstance(diagnostic, self._bound_diagnostic_type): raise TypeError( f"Expected diagnostic of type {self._bound_diagnostic_type}, got {type(diagnostic)}" ) if self.options.warnings_as_errors and diagnostic.level == infra.Level.WARNING: # type: ignore[attr-defined] diagnostic.level = infra.Level.ERROR # type: ignore[attr-defined] self.diagnostics.append(diagnostic) # type: ignore[arg-type] def log_and_raise_if_error(self, diagnostic: _Diagnostic) -> None: """Logs a diagnostic and raises an exception if it is an error. Use this method for logging non inflight diagnostics where diagnostic level is not known or lower than ERROR. If it is always expected raise, use `log` and explicit `raise` instead. Otherwise there is no way to convey the message that it always raises to Python intellisense and type checking tools. This method should be used only after all the necessary information for the diagnostic has been collected. Args: diagnostic: The diagnostic to add. """ self.log(diagnostic) if diagnostic.level == infra.Level.ERROR: if diagnostic.source_exception is not None: raise diagnostic.source_exception raise RuntimeErrorWithDiagnostic(diagnostic) @contextlib.contextmanager def add_inflight_diagnostic( self, diagnostic: _Diagnostic ) -> Generator[_Diagnostic, None, None]: """Adds a diagnostic to the context. Use this method to add diagnostics that are not created by the context. Args: diagnostic: The diagnostic to add. """ self._inflight_diagnostics.append(diagnostic) try: yield diagnostic finally: self._inflight_diagnostics.pop() def push_inflight_diagnostic(self, diagnostic: _Diagnostic) -> None: """Pushes a diagnostic to the inflight diagnostics stack. Args: diagnostic: The diagnostic to push. Raises: ValueError: If the rule is not supported by the tool. """ self._inflight_diagnostics.append(diagnostic) def pop_inflight_diagnostic(self) -> _Diagnostic: """Pops the last diagnostic from the inflight diagnostics stack. Returns: The popped diagnostic. """ return self._inflight_diagnostics.pop() def inflight_diagnostic(self, rule: infra.Rule | None = None) -> _Diagnostic: if rule is None: # TODO(bowbao): Create builtin-rules and create diagnostic using that. if len(self._inflight_diagnostics) <= 0: raise AssertionError("No inflight diagnostics") return self._inflight_diagnostics[-1] else: for diagnostic in reversed(self._inflight_diagnostics): if diagnostic.rule == rule: return diagnostic raise AssertionError(f"No inflight diagnostic for rule {rule.name}")