fixing OSS ci for tritonbench (#2439)

Summary:
Pull Request resolved: #2439

two problems found so far in CI runs
1. the ci file was not added to the target deps so the ci file couldnt be found even if the flag was used
2. `--op` was required so we fixed that here where `--op` is only required if `--ci` is not used

Reviewed By: danzimm, xuzhao9

Differential Revision: D61809602

fbshipit-source-id: e84844f0b6a697892222a008f45a8c20db8ef6bd

Author: adamomainz

userbenchmark/triton/ci.py

@@ -4,7 +4,7 @@
 from typing import List, Any, Dict
 
 from . import BM_NAME
-from .run import parse_args, _run
+from .run import get_parser, _run
 from userbenchmark.utils import get_default_output_json_path, get_output_json
 from torchbenchmark.util.triton_op import BenchmarkOperatorResult
 
@@ -22,7 +22,7 @@ def run_ci():
     ci_result = []
     for test_opts in CI_TESTS:
         logging.info(f"Running the test opts: {test_opts}")
-        test_args, test_extra_args = parse_args(test_opts)
+        test_args, test_extra_args = get_parser(test_opts).parse_known_args(test_opts)
         metrics = _run(test_args, test_extra_args)
         ci_result.append(metrics)
     result = ci_result_to_userbenchmark_json(ci_result)

userbenchmark/triton/run.py

@@ -25,12 +25,12 @@
 
 TRITON_BENCH_CSV_DUMP_PATH = tempfile.gettempdir() + "/tritonbench/"
 
-def get_parser():
+def get_parser(args = None):
     parser = argparse.ArgumentParser(allow_abbrev=False)
     parser.add_argument(
         "--op",
         type=str,
-        required=True,
+        required=False,
         help="Operator to benchmark."
     )
     parser.add_argument(
@@ -140,6 +140,12 @@ def get_parser():
     )
     if not hasattr(torch_version, "git_version"):
         parser.add_argument("--log-scuba", action="store_true", help="Log to scuba.")
+    
+    args, extra_args = parser.parse_known_args(args)
+    if  args.op and  args.ci:
+        parser.error("cannot specify operator when in CI mode")
+    elif not args.op and not args.ci:
+        parser.error("must specify operator when not in CI mode")
     return parser
 
 def _run(args: argparse.Namespace, extra_args: List[str]) -> BenchmarkOperatorResult:
@@ -187,5 +193,6 @@ def run(args: List[str] = []):
         from .ci import run_ci
         run_ci()
         return
+    
     with gpu_lockdown(args.gpu_lockdown):
         _run(args, extra_args)