add torch output directory option in performance scripts (#15032)

* add torch output directory

Signed-off-by: Brian Coutinho <[email protected]>

* Apply isort and black reformatting

Signed-off-by: briancoutinho <[email protected]>

---------

Signed-off-by: Brian Coutinho <[email protected]>
Signed-off-by: briancoutinho <[email protected]>
Co-authored-by: Brian Coutinho <[email protected]>
Co-authored-by: briancoutinho <[email protected]>

Author: briancoutinho

scripts/performance/argument_parser.py

@@ -174,6 +174,14 @@ def parse_cli_args():
         help="Enable torch profiler. Disabled by default",
         action="store_true",
     )
+    parser.add_argument(
+        "-tpo",
+        "--torch_profile_out_path",
+        type=str,
+        help="Path to the output file of torch profiling",
+        required=False,
+        default=None,
+    )
     parser.add_argument(
         "-tb",
         "--tensorboard",

scripts/performance/helpers.py

@@ -637,7 +637,9 @@ def build_torch_profiler_plugin(args):
         return PyTorchProfilerPlugin(
             start_step=start_iter,
             end_step=end_iter,
-            output_path=args.log_dir,  # a subdir torch_profiles will be created here.
+            output_path=os.environ.get(
+                'TORCH_PROFILES_DIR', args.torch_profile_out_path
+            ),  # a subdir torch_profiles will be created here.
             profiler_kwargs={
                 "with_stack": os.environ.get('TORCH_PROFILER_WITH_STACK', '0') == '1',
             },