diff --git a/.lintrunner.toml b/.lintrunner.toml index cd8a8d535e..00275ff7f9 100644 --- a/.lintrunner.toml +++ b/.lintrunner.toml @@ -303,7 +303,7 @@ include_patterns = [ 'runtime/**/*.py', 'scripts/**/*.py', # 'test/**/*.py', - # 'util/**/*.py', + 'util/**/*.py', '*.py', ] exclude_patterns = [ diff --git a/.mypy.ini b/.mypy.ini index bb1d574ab5..922b912cb3 100644 --- a/.mypy.ini +++ b/.mypy.ini @@ -80,4 +80,4 @@ ignore_missing_imports = True ignore_missing_imports = True [mypy-zstd] -ignore_missing_imports = True \ No newline at end of file +ignore_missing_imports = True diff --git a/util/activation_memory_profiler.py b/util/activation_memory_profiler.py index f459dfafaf..c149a46122 100644 --- a/util/activation_memory_profiler.py +++ b/util/activation_memory_profiler.py @@ -9,7 +9,7 @@ import json import typing from dataclasses import dataclass, field -from typing import List +from typing import Any, Dict, List, Optional import executorch.exir.memory as memory import torch @@ -52,7 +52,7 @@ def create_tensor_allocation_info(graph: torch.fx.Graph) -> List[MemoryTimeline] allocations at that timestep. """ nodes = graph.nodes - memory_timeline = [None] * len(nodes) + memory_timeline: List[Optional[MemoryTimeline]] = [None for _ in range(len(nodes))] for _, node in enumerate(nodes): if node.op == "output": continue @@ -72,11 +72,11 @@ def create_tensor_allocation_info(graph: torch.fx.Graph) -> List[MemoryTimeline] stack_trace = node.meta.get("stack_trace") fqn = _get_module_hierarchy(node) for j in range(start, end + 1): - if memory_timeline[j] is None: - # pyre-ignore - memory_timeline[j] = MemoryTimeline() - # pyre-ignore - memory_timeline[j].allocations.append( + memory_timeline_j = memory_timeline[j] + if memory_timeline_j is None: + memory_timeline_j = MemoryTimeline() + assert memory_timeline_j + memory_timeline_j.allocations.append( Allocation( node.name, node.target, @@ -87,8 +87,7 @@ def create_tensor_allocation_info(graph: torch.fx.Graph) -> List[MemoryTimeline] stack_trace, ) ) - # pyre-ignore - return memory_timeline + return memory_timeline # type: ignore[return-value] def _validate_memory_planning_is_done(exported_program: ExportedProgram): @@ -129,7 +128,7 @@ def generate_memory_trace( memory_timeline = create_tensor_allocation_info(exported_program.graph) root = {} - trace_events = [] + trace_events: List[Dict[str, Any]] = [] root["traceEvents"] = trace_events tid = 0 @@ -138,7 +137,7 @@ def generate_memory_trace( if memory_timeline_event is None: continue for allocation in memory_timeline_event.allocations: - e = {} + e: Dict[str, Any] = {} e["name"] = allocation.name e["cat"] = "memory_allocation" e["ph"] = "X" diff --git a/util/python_profiler.py b/util/python_profiler.py index 632187f56a..8993beb942 100644 --- a/util/python_profiler.py +++ b/util/python_profiler.py @@ -12,18 +12,20 @@ import re from pstats import Stats -from snakeviz.stats import json_stats, table_rows -from tornado import template +from snakeviz.stats import json_stats, table_rows # type: ignore[import-not-found] +from tornado import template # type: ignore[import-not-found] module_found = True +snakeviz_templates_dir: str = "" + try: - import snakeviz + import snakeviz # type: ignore[import-not-found] + + snakeviz_dir = os.path.dirname(os.path.abspath(snakeviz.__file__)) + snakeviz_templates_dir = os.path.join(snakeviz_dir, "templates") except ImportError: module_found = False -snakeviz_dir = os.path.dirname(os.path.abspath(snakeviz.__file__)) -snakeviz_templates_dir = os.path.join(snakeviz_dir, "templates") - def _from_pstat_to_static_html(stats: Stats, html_filename: str): """