lemonade-sdk 9.1.1__py3-none-any.whl

lemonade/profilers/profiler.py

@@ -0,0 +1,58 @@
+import abc
+
+
+class Profiler(abc.ABC):
+
+    unique_name: str
+
+    def __init__(self, parser_arg_value=None):
+        self.parser_arg_value = parser_arg_value
+        # Statistics that will be displayed to the CLI user
+        self.status_stats = []
+
+    @staticmethod
+    @abc.abstractmethod
+    def add_arguments_to_parser(parser):
+        """
+        Adds the argument parsing for this tool to the parser.
+        Uses f"--{self.unique_name}" as the argument.
+        """
+
+    @abc.abstractmethod
+    def start(self, build_dir):
+        """
+        This method is called prior to the tool sequence starting.
+        This informs the profiler to start gathering data.
+        The build directory can be used to store profiling data.
+        """
+
+    def tool_starting(self, tool_name):
+        """
+        This method is called to inform the profiler of the name of the tool that is about to start.
+        """
+
+    def tool_stopping(self):
+        """
+        This method is called to inform the profiler that the tool has finished.
+        """
+
+    def stop(self):
+        """
+        This method is called when the tool sequence has finished.
+        This informs the profiler to stop gathering data.
+        """
+
+    @abc.abstractmethod
+    def generate_results(self, state, timestamp, start_times):
+        """
+        This method is called so that the profiler can create its output files.
+        The state is passed so that build info can be gathered and stats can be written.
+        The timestamp can be used for filename in current working directory.
+        The start times parameter is a dict with the keys being the tools names and
+        the values being the time the tool started.  There is an initial "warmup" key
+        that has a start time before the first tool and a "cool down" key that contains the
+        time when the last tool ended.
+        """
+
+
+# Copyright (c) 2025 AMD

lemonade/sequence.py

@@ -0,0 +1,363 @@
+import sys
+import time
+import os
+import platform
+import copy
+from datetime import datetime
+from typing import List, Dict, Optional
+import pytz
+import psutil
+import lemonade.common.printing as printing
+import lemonade.common.exceptions as exp
+import lemonade.common.build as build
+from lemonade.common.system_info import get_system_info_dict
+import lemonade.common.filesystem as fs
+import lemonade.common.status as status
+from lemonade.tools.tool import Tool
+from lemonade.profilers.profiler import Profiler
+from lemonade.state import State
+
+
+def _rewind_stdout(lines: int = 1):
+    """
+    Helper function for the command line monitor. Moves the cursor up a
+    certain number of lines in the terminal, corresponding to the
+    status line for a Tool, so that we can update the status of
+    that Tool.
+    """
+    rewind_stdout_one_line = "\033[1A"
+    rewind_multiple_lines = rewind_stdout_one_line * lines
+    print(rewind_multiple_lines, end="")
+    sys.stdout.flush()
+
+
+class Sequence:
+    """
+    Helper class to launch and manage build tools.
+    """
+
+    def __init__(
+        self,
+        tools: Dict[Tool, List[str]],
+        profilers: List[Profiler] = None,
+    ):
+
+        self.tools = tools
+        self.profilers = [] if profilers is None else profilers
+
+        # Make sure all the tool names are unique
+        self.tool_names = [tool.__class__.unique_name for tool in self.tools.keys()]
+
+        if len(self.tool_names) != len(set(self.tool_names)):
+            msg = f"""
+            All tools in a Sequence must have unique unique_names, however Sequence
+            received duplicates in the list of names: {self.tool_names}
+            """
+            raise ValueError(msg)
+
+        # Save the process (used to get memory usage)
+        self.process = psutil.Process()
+
+    def show_monitor(self, state: State, verbosity: bool):
+        """
+        Displays the monitor on the terminal. The purpose of the monitor
+        is to show the status of each tool (success, failure, not started yet,
+        or in-progress).
+        """
+
+        if verbosity:
+            print()
+
+            printing.logn(
+                f'Building "{state.build_name}"',
+                c=printing.Colors.BOLD,
+            )
+
+            for tool in self.tools:
+                tool.status_line(successful=None, verbosity=True)
+
+            _rewind_stdout(len(self.tools))
+
+    def _advance_cursor(self, current_tool_name: str):
+        # Advance the cursor below the monitor so
+        # we can print a message
+        tool_depth_in_sequence = len(self.tool_names) - self.tool_names.index(
+            current_tool_name
+        )
+        stdout_lines_to_advance = tool_depth_in_sequence - 2
+        cursor_down = "\n" * stdout_lines_to_advance
+
+        print(cursor_down)
+
+    def _get_mem_usage_str(self) -> str:
+        """
+        Returns a string with memory usage for the current process
+        (non-swapped physical memory).  In Windows OS, the peak memory used in the
+        process is also included.
+
+        Example: '1.100 GB (1.638 GB peak)'
+        """
+        mem_info = self.process.memory_info()
+        mem_info_str = f"{mem_info.rss / 1024 ** 3:,.3f} GB"
+        if platform.system() == "Windows":
+            mem_info_str += f" ({mem_info.peak_wset / 1024 ** 3:,.3f} GB peak)"
+        return mem_info_str
+
+    def launch(
+        self,
+        state: State,
+        lean_cache: bool = False,
+        monitor: Optional[bool] = None,
+        stats_to_save: Optional[Dict] = None,
+    ) -> State:
+        """
+        Executes the sequence of tools.
+        """
+
+        current_time = datetime.now()
+        timestamp = current_time.strftime("%Y-%m-%d-%H%M%S")
+        start_times = {"warmup": time.time()}
+
+        # Allow monitor to be globally disabled by an environment variable
+        if monitor is None:
+            if os.environ.get("LEMONADE_BUILD_MONITOR") == "False":
+                monitor_setting = False
+            else:
+                monitor_setting = True
+        else:
+            monitor_setting = monitor
+
+        # Create a build directory in the cache
+        fs.make_build_dir(state.cache_dir, state.build_name)
+
+        # Start profilers
+        build_dir = build.output_dir(state.cache_dir, state.build_name)
+        for profiler in self.profilers:
+            profiler.start(build_dir)
+
+        self.show_monitor(state, monitor_setting)
+
+        if state.build_status == build.FunctionStatus.SUCCESSFUL:
+            msg = """
+            build_model() is running a build on a model that already built successfully, which
+            should not happen because the build should have loaded from cache or rebuilt from scratch.
+            If you are using custom tools and Sequences then you have some debugging to do. Otherwise,
+            please file an issue at https://github.com/lemonade-sdk/lemonade/issues
+            """
+            raise exp.Error(msg)
+
+        # Keep a copy of any stats we loaded from disk, in case we need to
+        # restore them later
+        saved_stats = copy.deepcopy(fs.Stats(state.cache_dir, state.build_name).stats)
+
+        # Save build name to stats so it shows up on reports
+        state.save_stat(fs.Keys.BUILD_NAME, state.build_name)
+
+        # Indicate that the build is running. If the build fails for any reason,
+        # we will try to catch the exception and note it in the stats.
+        # If a concluded build still has a status of "running", this means
+        # there was an uncaught exception.
+        state.save_stat(fs.Keys.BUILD_STATUS, build.FunctionStatus.INCOMPLETE)
+
+        # Save a timestamp so that we know the order of builds within a cache
+        pacific_tz = pytz.timezone("America/Los_Angeles")
+        state.save_stat(
+            fs.Keys.TIMESTAMP,
+            datetime.now(pacific_tz),
+        )
+
+        # Save the system information used for this build
+        system_info = get_system_info_dict()
+        state.save_stat(
+            fs.Keys.SYSTEM_INFO,
+            system_info,
+        )
+
+        # Collect telemetry for the build
+        state.save_stat(
+            fs.Keys.SELECTED_SEQUENCE_OF_TOOLS,
+            self.tool_names,
+        )
+
+        # At the beginning of a sequence no tool has started
+        for tool in self.tools:
+            state.save_stat(tool.status_key, build.FunctionStatus.NOT_STARTED)
+            state.save_stat(tool.duration_key, "-")
+            state.save_stat(tool.memory_key, "-")
+
+        # Save any additional stats passed in via arguments
+        if stats_to_save:
+            for stat_key, stat_value in stats_to_save.items():
+                state.save_stat(stat_key, stat_value)
+
+        # Save initial memory as a build statistic
+        state.save_stat(f"{fs.Keys.TOOL_MEMORY}:__init__", self._get_mem_usage_str())
+
+        # Run the build
+        saved_exception = None
+        for tool, argv in self.tools.items():
+
+            start_time = time.time()
+            start_times[tool.unique_name] = start_time
+
+            # Inform profiler of name of tool about to start
+            for profiler in self.profilers:
+                profiler.tool_starting(tool.unique_name)
+
+            try:
+
+                # Set status as incomplete, since tool just started
+                state.save_stat(tool.status_key, build.FunctionStatus.INCOMPLETE)
+
+                # Collect telemetry about the tool
+                state.current_build_tool = tool.unique_name
+
+                # Run the tool
+                state = tool.parse_and_run(state, argv, monitor_setting)
+
+                # Save the state so that it can be assessed for a cache hit
+                state.save()
+
+            except exp.SkipBuild as e:
+                # SkipBuild is a special exception, which means that a build
+                # was loaded from disk, then we realized we want to skip it.
+                # In order to preserve the original stats and state of the build,
+                # we need to restore the stats file to what it was at the beginning
+                # of this function call. We also need to avoid calling state.save().
+
+                # Restore the prior stats
+                fs.save_yaml(
+                    saved_stats, fs.Stats(state.cache_dir, state.build_name).file
+                )
+
+                # Advance the cursor below the monitor so
+                # we can print a message
+                self._advance_cursor(tool.unique_name)
+                printing.log_warning(str(e))
+                return
+
+            # Broad exception is desirable as we want to capture
+            # all exceptions (including those we can't anticipate)
+            except Exception as e:  # pylint: disable=broad-except
+
+                if os.environ.get("LEMONADE_DEBUG", "").lower() == "true":
+                    # It may be useful to raise the exception here, since
+                    # if any of the subsequent lines of code raise another
+                    # exception it will be very hard to root cause e.
+                    raise e
+
+                # Update tool and build status
+                state.save_stat(tool.status_key, build.FunctionStatus.ERROR)
+                state.save_stat(fs.Keys.BUILD_STATUS, build.FunctionStatus.ERROR)
+
+                # Save the log file for the failed tool to stats for easy reference
+                stats = fs.Stats(state.cache_dir, state.build_name)
+                stats.save_eval_error_log(tool.logfile_path)
+
+                # Advance the cursor below the monitor so
+                # we can print a message
+                self._advance_cursor(tool.unique_name)
+
+                if vars(state).get("invocation_info"):
+                    state.invocation_info.status_message = f"Error: {e}"
+                    state.invocation_info.status_message_color = printing.Colors.WARNING
+                else:
+                    printing.log_error(e)
+
+                # We will raise this exception after we capture as many statistics
+                # about the build as possible
+                saved_exception = e
+
+                # Don't run any more tools
+                break
+
+            else:
+                # Update tool Status
+                state.save_stat(tool.status_key, build.FunctionStatus.SUCCESSFUL)
+                state.current_build_tool = None
+
+            finally:
+                # Store tool duration
+                execution_time = time.time() - start_time
+                state.save_stat(tool.duration_key, execution_time)
+
+                # Store current memory and peak working memory
+                state.save_stat(tool.memory_key, self._get_mem_usage_str())
+
+                # Inform profilers that tool has finished
+                for profiler in self.profilers:
+                    profiler.tool_stopping()
+
+        start_times["cool down"] = time.time()
+
+        # Tell the profilers to stop gathering data
+        for profiler in self.profilers:
+            profiler.stop()
+
+        if not saved_exception:
+            state.build_status = build.FunctionStatus.SUCCESSFUL
+            state.save_stat(fs.Keys.BUILD_STATUS, build.FunctionStatus.SUCCESSFUL)
+            if vars(state).get("invocation_info"):
+                state.invocation_info.status_message = (
+                    f"Successful build! {state.invocation_info.extra_status}"
+                )
+                state.invocation_info.status_message_color = printing.Colors.OKGREEN
+
+        # Generate profiler output
+        for profiler in self.profilers:
+            profiler.generate_results(state, timestamp, start_times)
+
+        if vars(state).get("models_found") and vars(state).get("invocation_info"):
+
+            # Present status statistics from the tools
+            for tool in self.tools:
+                state.invocation_info.stats_keys += tool.status_stats
+
+            # Present status statistics from the profilers
+            for profiler in self.profilers:
+                state.invocation_info.stats_keys += profiler.status_stats
+
+            print()
+
+            status.recursive_print(
+                models_found=state.models_found,
+                build_name=state.build_name,
+                cache_dir=state.cache_dir,
+                parent_model_hash=None,
+                parent_invocation_hash=None,
+                script_names_visited=[],
+            )
+
+        if lean_cache:
+            printing.log_info("Removing build artifacts...")
+            fs.clean_output_dir(state.cache_dir, state.build_name)
+
+        state.save()
+
+        if saved_exception:
+            raise saved_exception
+
+        printing.log_success(
+            f"\n    Saved to **{build.output_dir(state.cache_dir, state.build_name)}**"
+        )
+
+        return state
+
+    def status_line(self, verbosity):
+        """
+        Print a status line in the monitor for every tool in the sequence
+        """
+        for tool in self.tools:
+            tool.status_line(successful=None, verbosity=verbosity)
+
+    @property
+    def info(self) -> Dict[str, Dict]:
+        """
+        Return a dictionary of tool_name:argv for the sequence
+        """
+
+        return {tool.__class__.unique_name: argv for tool, argv in self.tools.items()}
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/state.py

@@ -0,0 +1,159 @@
+import os
+import sys
+from typing import Dict, Optional, Any
+import yaml
+import lemonade.common.build as build
+import lemonade.common.filesystem as fs
+from lemonade.version import __version__ as lemonade_version
+
+
+def _is_nice_to_write(value):
+    """
+    Checks whether a value is nice to write to YAML.
+    Returns True if the value is a string, int, float, bool, list, dict, or tuple.
+    Returns False otherwise.
+    """
+    if isinstance(value, (str, int, float, bool)):
+        return True
+    elif isinstance(value, list) or isinstance(value, tuple):
+        # Check if all elements in the list are nice to write
+        return all(_is_nice_to_write(item) for item in value)
+    elif isinstance(value, dict):
+        # Check if all values in the dictionary are nice to write
+        return all(_is_nice_to_write(item) for item in value.values())
+    return False
+
+
+def _sanitize_for_yaml(input_dict: Dict) -> Dict:
+    """
+    Creates a new dictionary containing only nice-to-write values
+    from the original dictionary.
+    """
+    result = {}
+    for key, value in input_dict.items():
+        if _is_nice_to_write(value):
+            result[key] = value
+    return result
+
+
+class State:
+    """
+    The State class is meant to carry build state, starting with the user's
+    initial arguments, through each build Tool in the Sequence, and finally
+    to the disk, where it is used to assess cache hits.
+
+    State is initialized with the key members that are shared by every build,
+    and reasonable default values are assigned as appropriate.
+
+    Tool developers can also add any members they wish. To get or set an
+    attribute, reference it as an attribute:
+        1. get: `my_variable = state.attribute_name`
+        2. set: `state.attribute_name = my_variable`
+
+    Build State can be saved and loaded from disk in the form of a state.yaml file
+    via State.save() and load_state(), respectively. Note that while State can
+    contain members of any type, only YAML-safe members (str, int, bool, float,
+    list, dict, tuple) will be saved and loaded.
+    """
+
+    def __init__(
+        self,
+        cache_dir: str,
+        build_name: Optional[str] = None,
+        sequence_info: Dict[str, Dict] = None,
+        **kwargs,
+    ):
+
+        # The default model name is the name of the python file that calls build_model()
+        if build_name is None:
+            build_name = os.path.basename(sys.argv[0])
+
+        # Support "~" in the cache_dir argument
+        parsed_cache_dir = os.path.expanduser(cache_dir)
+
+        # Save settings as State members
+        self.cache_dir = parsed_cache_dir
+        self.build_name = build_name
+        self.sequence_info = sequence_info
+        self.lemonade_version = lemonade_version
+        self.build_status = build.FunctionStatus.NOT_STARTED
+        self.downcast_applied = False
+        self.uid = build.unique_id()
+        self.results = None
+
+        # Store any additional kwargs as members
+        for key, value in kwargs.items():
+            self.__dict__[key] = value
+
+    def __setattr__(self, name: str, value: Any) -> None:
+        """
+        Tool developers can add a new member to State by simply
+        assigning it as an attribute, i.e., `state.new_member = value`.
+        """
+        return super().__setattr__(name, value)
+
+    def save_stat(self, key: str, value):
+        """
+        Save statistics to an yaml file in the build directory
+        """
+
+        stats = fs.Stats(self.cache_dir, self.build_name)
+        stats.save_stat(key, value)
+
+    def save_sub_stat(self, parent_key: str, key: str, value):
+        """
+        Save statistics to an yaml file in the build directory
+        """
+
+        stats = fs.Stats(self.cache_dir, self.build_name)
+        stats.save_sub_stat(parent_key, key, value)
+
+    def save(self):
+        """
+        Save all YAML-friendly members to disk as a state.yaml file.
+
+        Note that `model` and `inputs` will typically not be saved since
+        they are typically in non-YAML-friendly types such as `torch.nn.Module`
+        and `torch.tensor`.
+        """
+
+        state_to_save = _sanitize_for_yaml(vars(self))
+
+        # Create a build directory in the cache
+        fs.make_build_dir(self.cache_dir, self.build_name)
+
+        with open(
+            build.state_file(self.cache_dir, self.build_name),
+            "w",
+            encoding="utf8",
+        ) as outfile:
+            yaml.dump(state_to_save, outfile)
+
+
+def load_state(
+    cache_dir=None,
+    build_name=None,
+    state_path=None,
+) -> State:
+    """
+    Read a state.yaml file corresponding to a specific build in a specific
+    cache, and use its contents to initialize a State instance.
+    """
+
+    if state_path is not None:
+        file_path = state_path
+    elif build_name is not None and cache_dir is not None:
+        file_path = build.state_file(cache_dir, build_name)
+    else:
+        raise ValueError(
+            "This function requires either build_name and cache_dir to be set, "
+            "or state_path to be set, not both or neither"
+        )
+
+    state_dict = build.load_yaml(file_path)
+
+    return State(**state_dict)
+
+
+# This file was originally licensed under Apache 2.0. It has been modified.
+# Modifications Copyright (c) 2025 AMD

lemonade/tools/__init__.py

@@ -0,0 +1 @@
+from .tool import Tool, FirstTool, NiceHelpFormatter