Add progress tracking lib

PiperOrigin-RevId: 871707321

Author: chunnienc

litert_torch/__init__.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+
 from litert_torch._config import config
 from litert_torch._convert.interface import convert
 from litert_torch._convert.interface import experimental_add_compilation_backend

litert_torch/_config.py

@@ -93,5 +93,14 @@ def lazy_constant_getter_chunk_size(self) -> int:
   def lazy_constant_getter_chunk_size(self, value: int) -> None:
     os.environ["LAZY_CONSTANT_GETTER_CHUNK_SIZE"] = str(value)
 
+  @property
+  def show_progress(self) -> bool:
+    """True if the progress should be shown."""
+    return _get_bool_env_var("LITERT_TORCH_SHOW_PROGRESS", default=True)
+
+  @show_progress.setter
+  def show_progress(self, value: bool):
+    os.environ["LITERT_TORCH_SHOW_PROGRESS"] = "y" if value else "n"
+
 
 config = _Config()

litert_torch/_convert/core.py

@@ -20,6 +20,7 @@
 
 from litert_torch import fx_infra
 from litert_torch import model
+from litert_torch import progress
 from litert_torch._convert import fx_passes
 from litert_torch._convert import litert_converter
 from litert_torch._convert import signature
@@ -70,6 +71,7 @@ def _warn_training_modules(signatures: list[signature.Signature]):
     logging.warning(message)
 
 
+@progress.task("LiteRT-Torch Convert")
 def convert_signatures(
     signatures: list[signature.Signature],
     *,
@@ -127,18 +129,20 @@ def export(**kwargs):
     )
     return exported_program
 
-  exported_programs = [
-      export(
+  exported_programs = []
+  for sig in signatures:
+    with progress.task(f"Torch Export: {sig.name}"):
+      exported_program = export(
           mod=sig.module,
           args=sig.args,
           kwargs=sig.kwargs,
           dynamic_shapes=sig.dynamic_shapes,
       )
-      for sig in signatures
-  ]
+    exported_programs.append(exported_program)
 
   # Apply default fx passes
-  exported_programs = list(map(_run_convert_passes, exported_programs))
+  with progress.task("Run FX Passes"):
+    exported_programs = list(map(_run_convert_passes, exported_programs))
 
   exporter = litert_converter.exported_programs_to_flatbuffer(
       exported_programs,

litert_torch/_convert/litert_converter.py

@@ -18,6 +18,7 @@
 
 from litert_torch import backend
 from litert_torch import model as model_lib
+from litert_torch import progress
 from litert_torch._convert import signature
 from litert_torch.backend import inline_consts as inline_consts_lib
 from litert_torch.quantize import quant_config as qcfg
@@ -75,16 +76,20 @@ def to_file(self, path: str):
         f.write(self.content)
       return
 
-    try:
-      # TODO b/478909085 - Remove the try-except once converter_api_ext is
-      # stable in OSS.
-      converter_api_ext.export_flatbuffer_to_file(self._module_op, path)
-    except TypeError:
-      converter_api_ext.export_flatbuffer_to_file(self.module, path)
+    with progress.task(f"Write Model to {path}"):
+      try:
+        # TODO b/478909085 - Remove the try-except once converter_api_ext is
+        # stable in OSS.
+        converter_api_ext.export_flatbuffer_to_file(self._module_op, path)
+      except TypeError:
+        converter_api_ext.export_flatbuffer_to_file(self.module, path)
 
   def to_bytes(self) -> bytes:
     """Returns the flatbuffer bytes of the module."""
-    if self.content is None:
+    if self.content is not None:
+      return self.content
+
+    with progress.task("Write Model to Bytes"):
       try:
         # TODO b/478909085 - Remove the try-except once converter_api_ext is
         # stable in OSS.
@@ -93,8 +98,8 @@ def to_bytes(self) -> bytes:
         )
       except TypeError:
         self.content = converter_api_ext.export_flatbuffer_to_bytes(self.module)
-      self.module = None
 
+    self.module = None
     return self.content
 
 
@@ -114,19 +119,19 @@ def exported_programs_to_flatbuffer(
     )
 
   ir_context = backend.export_utils.create_ir_context()
-
   cross_program_inline_consts_ctx = inline_consts_lib.InlineConstsContext(
       enable_lazy_constants=lightweight_conversion,
   )
 
   lowered_programs = []
   for exported_program, sig in zip(exported_programs, signatures):
     # Convert ExportedProgram to Mlir Module.
-    lowered = backend.export.exported_program_to_mlir(
-        exported_program,
-        ir_context=ir_context,
-        lowering_context_plugins=[cross_program_inline_consts_ctx],
-    )
+    with progress.task(f"Lower to MLIR: {sig.name}"):
+      lowered = backend.export.exported_program_to_mlir(
+          exported_program,
+          ir_context=ir_context,
+          lowering_context_plugins=[cross_program_inline_consts_ctx],
+      )
 
     # Set signature.
     sig_name = sig.name
@@ -141,9 +146,10 @@ def exported_programs_to_flatbuffer(
     lowered_programs.append(lowered)
 
   # Merge all lowered modules into one module.
-  merged_module = converter_api_ext.merge_modules(
-      [lowered.module for lowered in lowered_programs]
-  )
+  with progress.task("Merge MLIR Modules"):
+    merged_module = converter_api_ext.merge_modules(
+        [lowered.module for lowered in lowered_programs]
+    )
 
   # Prepare ai-edge-quantizer recipe.
   translated_recipe = None
@@ -170,7 +176,7 @@ def exported_programs_to_flatbuffer(
   config.canonicalizing_inf_as_min_max_float = False
 
   # Run LiteRT converter passes.
-  with ir_context:
+  with ir_context, progress.task("Run LiteRT Converter Passes"):
     pass_manager = passmanager.PassManager()
     converter_api_ext.run_convert_to_tfl_passes(
         merged_module, pass_manager, config
@@ -181,9 +187,10 @@ def exported_programs_to_flatbuffer(
 
   # Quantize the model if needed.
   if translated_recipe:
-    model_bytes = translate_recipe.quantize_model(
-        exporter.to_bytes(), translated_recipe
-    )
-    exporter = LazyModelExporter(content=model_bytes)
+    with progress.task("Quantize Model"):
+      model_bytes = translate_recipe.quantize_model(
+          exporter.to_bytes(), translated_recipe
+      )
+      exporter = LazyModelExporter(content=model_bytes)
 
   return exporter

litert_torch/backend/export.py

@@ -21,6 +21,7 @@
 from typing import Any, Callable, Optional
 
 from litert_torch import fx_infra
+from litert_torch import progress
 from litert_torch.backend import _torch_future
 from litert_torch.backend import debuginfo
 from litert_torch.backend import export_utils
@@ -375,7 +376,7 @@ def exported_program_to_mlir(
   if not ir_context:
     ir_context = export_utils.create_ir_context()
 
-  with ir_context, ir.Location.unknown():
+  with ir_context, ir.Location.unknown(), progress.task("Create MLIR Module"):
 
     module = ir.Module.create()
     lctx = LoweringContext(ir_context, module)

litert_torch/fx_infra/_safe_run_decompositions.py

@@ -13,8 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 """ExportedProgram.run_decompositions wrapper to handle unexpected export behavior."""
+
 import operator
 from typing import Any, Callable
+from litert_torch import progress
 import torch
 
 
@@ -69,6 +71,7 @@ def annotate_force_decomp(decomp: Callable[..., Any]):
   return decomp
 
 
+@progress.task("ExportedProgram Run Decompositions")
 def safe_run_decompositions(exported_program, decomp_table=None, can_skip=True):
   """Wrapper for ExportedProgram.run_decompositions to handle unexpected export behavior."""
 

litert_torch/progress.py

@@ -0,0 +1,101 @@
+# Copyright 2024 The LiteRT Torch 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.
+# ==============================================================================
+
+"""Progress tracking and logging library."""
+
+import contextlib
+import dataclasses
+import time
+from typing import Generator
+from litert_torch import _config
+from rich import console
+
+__all__ = ["task", "console"]
+
+config = _config.config
+console = console.Console(color_system="auto")
+
+
+@dataclasses.dataclass
+class Task:
+  name: str
+  start_time: float = 0.0
+
+
+_task_stack: list[Task] = []
+
+
+def _fmt_elapsed_time(elapsed_time: float) -> str:
+  minutes, seconds = divmod(elapsed_time, 60)
+  return f"{int(minutes):02d}:{int(seconds):02d}"
+
+
+def _style_elapsed_time(elapsed: float) -> str:
+  formatted = _fmt_elapsed_time(elapsed)
+  if elapsed < 5:
+    return f"[dim default]{formatted}[/dim default]"
+  elif elapsed < 60:
+    return f"[yellow]{formatted}[/yellow]"
+  elif elapsed < 60 * 10:
+    return f"[bold yellow]{formatted}[/bold yellow]"
+  else:
+    return f"[bold red]{formatted}[/bold red]"
+
+
+def _task_stack_repr() -> str:
+  return " > ".join(
+      f"[bold default]{task.name}[/bold default]" for task in _task_stack
+  )
+
+
+def _stack_elapsed_time() -> str:
+  elapsed = (
+      0 if not _task_stack else time.perf_counter() - _task_stack[0].start_time
+  )
+  return f"[dim default]({_fmt_elapsed_time(elapsed)})[/dim default]"
+
+
+@contextlib.contextmanager
+def task(name: str) -> Generator[None, None, None]:
+  """Context manager for tracking one task."""
+  if not config.show_progress:
+    yield
+    return
+
+  current_task = Task(name, time.perf_counter())
+  _task_stack.append(current_task)
+
+  stack_view = _task_stack_repr()
+  console.print(
+      f"{_stack_elapsed_time()} [bold cyan][START][/bold cyan] {stack_view}"
+  )
+
+  try:
+    yield
+  except Exception:
+    console.print(
+        f"{_stack_elapsed_time()} [bold red][ FAIL][/bold red] {stack_view}"
+    )
+    raise
+  else:
+    elapsed_time = time.perf_counter() - current_task.start_time
+    console.print(
+        f"{_stack_elapsed_time()} [bold green][ DONE][/bold green]"
+        f" [dim]{stack_view}[/dim]"
+        f" [dim](+[/dim]{_style_elapsed_time(elapsed_time)}[dim])[/dim]"
+    )
+  finally:
+    if _task_stack:
+      _task_stack.pop()