Install snakeviz for cProfile flamegraph (pytorch#1915)

Summary:
Pull Request resolved: pytorch#1915

fix ci test: https://github.com/pytorch/executorch/actions/runs/7848669862/job/21420323950

Reviewed By: mergennachin

Differential Revision: D53620967

fbshipit-source-id: b4860566016fb05e28ec18f7aeaa751e71cec730

Author: lucylq

.ci/scripts/test.sh

@@ -60,6 +60,8 @@ run_portable_executor_runner() {
 
 test_model() {
   if [[ "${MODEL_NAME}" == "llama2" ]]; then
+    # Install requirements for export_llama
+    bash examples/models/llama2/install_requirements.sh
     # Test export_llama script: python3 -m examples.models.llama2.export_llama
     "${PYTHON_EXECUTABLE}" -m examples.models.llama2.export_llama -c examples/models/llama2/params/demo_rand_params.pth -p examples/models/llama2/params/demo_config.json
     run_portable_executor_runner

examples/models/llama2/export_llama_lib.py

@@ -25,7 +25,6 @@
 from executorch.exir.passes.quant_fusion_pass import QuantFusionPass
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
 from executorch.util.activation_memory_profiler import generate_memory_trace
-from executorch.util.python_profiler import CProfilerFlameGraph
 from torch._export import capture_pre_autograd_graph
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 from torch.ao.quantization.quantizer.composable_quantizer import ComposableQuantizer
@@ -266,8 +265,16 @@ def _get_quantization_options(args):
 
 def export_llama(modelname, args) -> str:
     if args.profile_path is not None:
-        with CProfilerFlameGraph(args.profile_path):
-            return _export_llama(modelname, args)
+        try:
+            from executorch.util.python_profiler import CProfilerFlameGraph
+
+            with CProfilerFlameGraph(args.profile_path):
+                return _export_llama(modelname, args)
+        except ImportError:
+            print(
+                "Please run `pip install snakeviz` to install required dependencies for cProfiler flamegraph."
+            )
+            return ""
     else:
         return _export_llama(modelname, args)
 

examples/models/llama2/install_requirements.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Install snakeviz for cProfile flamegraph
+pip install snakeviz