Add local optoin for output json file (#1426)

* Add local optoin for output json file

Summary:
We are planning to run the benchmarks in local machine for now since there are some issues when trying to run in OSS CI,
so we have to fill in some fields that is populated by the OSS CI before

Test Plan:
```
python torchao/_models/llama/generate.py --checkpoint_path "${CHECKPOINT_PATH}/${MODEL_REPO}/model.pth" --compile --compile_prefill --output_json_path benchmark-results.json --output_json_local

python torchao/_models/llama/generate.py --checkpoint_path "${CHECKPOINT_PATH}/${MODEL_REPO}/model.pth" --compile --compile_prefill --quantization autoquant --output_json_path benchmark-results.json --output_js\on_local

python torchao/_models/sam/eval_combo.py --coco_root_dir datasets/coco2017 --coco_slice_name val2017 --sam_checkpoint_base_path checkpoints --sam_model_type vit_h --point_sampling_cache_dir tmp/sam_coco_mask_ce\nter_cache --mask_debug_out_dir tmp/sam_eval_masks_out --batch_size 32 --num_workers 8 --use_compile max-autotune --use_half bfloat16 --device cuda --output_json_path benchmark-results.json --output_json_local

python torchao/_models/sam/eval_combo.py --coco_root_dir datasets/coco2017 --coco_slice_name val2017 --sam_checkpoint_base_path checkpoints --sam_model_type vit_h --point_sampling_cache_dir tmp/sam_coco_mask_ce\nter_cache --mask_debug_out_dir tmp/sam_eval_masks_out --batch_size 32 --num_workers 8 --use_compile max-autotune --use_half bfloat16 --device cuda --compression autoquant --output_json_path benchmark-results.j\son --output_json_local

cd examples/sam2_amg_server
python server.py ${CHECKPOINT_PATH}/sam2 large --port 4000 --host localhost --fast --benchmark --dry --output_json_path ../../benchmark-results.json --output_json_local
python server.py ${CHECKPOINT_PATH}/sam2 large --port 4000 --host localhost --fast --use_autoquant --benchmark --dry --output_json_path ../../benchmark-results.json --output_json_local
cd ../..
```
Reviewers:

Subscribers:

Tasks:

Tags:

* fix

Author: jerryzh168

examples/sam2_amg_server/server.py

@@ -30,7 +30,8 @@
 import contextlib
 from torchao._models.utils import (
     get_arch_name,
-    write_json_result,
+    write_json_result_ossci,
+    write_json_result_local,
 )
 
 from torch._inductor import config as inductorconfig
@@ -564,7 +565,8 @@ def main(checkpoint_path,
          batch_size=1,
          load_fast="",
          save_fast="",
-         output_json_path=None):
+         output_json_path=None,
+         output_json_local=False):
     if verbose:
         logging.basicConfig(level=logging.INFO,
                             format='%(asctime)s - %(levelname)s - %(message)s',
@@ -661,6 +663,7 @@ def main(checkpoint_path,
             memory_result = [name, dtype, device, arch, "memory(MiB)", max_memory_allocated_bytes, None]
             memory_percent_result = [name, dtype, device, arch, "memory(%)", max_memory_allocated_percentage, None]
             performance_result = [name, dtype, device, arch, "time_s(avg)", avg_time_per_run, None]
+            write_json_result = write_json_result_local if output_json_local else write_json_result_ossci
             write_json_result(output_json_path, headers, memory_result)
             write_json_result(output_json_path, headers, memory_percent_result)
             write_json_result(output_json_path, headers, performance_result)

torchao/_models/llama/generate.py

@@ -20,7 +20,8 @@
 from torchao.utils import get_model_size_in_bytes, TORCH_VERSION_AT_LEAST_2_5
 from torchao._models.utils import (
     get_arch_name,
-    write_json_result,
+    write_json_result_ossci,
+    write_json_result_local,
 )
 
 torch.sparse.SparseSemiStructuredTensor._FORCE_CUTLASS = False
@@ -134,7 +135,7 @@ def decode_n_tokens(
             next_token, next_prob = next_token.clone(), next_prob.clone()
             input_pos += 1
             # in some instances not having this causes weird issues with the stored tokens when you run the next decode_one_token step
-            new_tokens.append(next_token.clone()) 
+            new_tokens.append(next_token.clone())
             callback(new_tokens[-1])
             new_probs.append(next_prob)
             cur_token = next_token
@@ -278,6 +279,7 @@ def main(
     precision=torch.bfloat16,
     write_result: Optional[Path] = None,
     output_json_path: Optional[Path] = None,
+    output_json_local: bool = False,
 ) -> None:
     """Generates text samples based on a pre-trained Transformer model and tokenizer."""
 
@@ -941,6 +943,7 @@ def callback(x):
         dtype = quantization or str(precision)
         memory_result = [name, dtype, device, arch, "mem/s", bandwidth, None]
         performance_result = [name, dtype, device, arch, "tok/s", tokpersec, None]
+        write_json_result = write_json_result_local if output_json_local else write_json_result_ossci
         write_json_result(output_json_path, headers, memory_result)
         write_json_result(output_json_path, headers, performance_result)
 
@@ -1042,6 +1045,11 @@ def callback(x):
         default=None,
         help="Path where to write the json result for dashboard",
     )
+    parser.add_argument(
+        "--output_json_local",
+        action="store_true",
+        help="Whether to output json result for local machine or for CI machine, local option will fill in some dummy fields",
+    )
 
     args = parser.parse_args()
     print(args)
@@ -1069,4 +1077,5 @@ def callback(x):
         args.precision,
         args.write_result,
         args.output_json_path,
+        args.output_json_local,
     )

torchao/_models/sam/eval_combo.py

@@ -23,7 +23,8 @@
 from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 from torchao._models.utils import (
     get_arch_name,
-    write_json_result,
+    write_json_result_ossci,
+    write_json_result_local,
 )
 
 torch._dynamo.config.cache_size_limit = 50000
@@ -253,6 +254,7 @@ def run(
     memory_path=None,
     device="cuda",
     output_json_path=None,
+    output_json_local=False,
 ):
     from torch._inductor import config as inductorconfig
     inductorconfig.triton.unique_kernel_names = True
@@ -468,6 +470,7 @@ def mlp_only(mod, name):
         dtype = compress or str(use_half) or "torch.float32"
         memory_result = [name, dtype, device, arch, "memory(MiB)", max_memory_allocated_bytes, None]
         performance_result = [name, dtype, device, arch, "img_s(avg)", img_s, None]
+        write_json_result = write_json_result_local if output_json_local else write_json_result_ossci
         write_json_result(output_json_path, headers, memory_result)
         write_json_result(output_json_path, headers, performance_result)
 

torchao/_models/utils.py

@@ -2,6 +2,9 @@
 import torch
 import platform
 import os
+import datetime
+import hashlib
+import time
 
 def get_arch_name() -> str:
     if torch.cuda.is_available():
@@ -11,14 +14,65 @@ def get_arch_name() -> str:
         return platform.machine()
 
 
-def write_json_result(output_json_path, headers, row):
+def write_json_result_ossci(output_json_path, headers, row):
     """
     Write the result into JSON format, so that it can be uploaded to the benchmark database
     to be displayed on OSS dashboard. The JSON format is defined at
     https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database
+
+    OSS CI version, that will leave many fields to be filled in by CI
+    """
+    mapping_headers = {headers[i]: v for i, v in enumerate(row)}
+    record = {
+        "benchmark": {
+            "name": "TorchAO benchmark",
+            "mode": "inference",
+            "dtype": mapping_headers["dtype"],
+            "extra_info": {
+                "device": mapping_headers["device"],
+                "arch": mapping_headers["arch"],
+            },
+        },
+        "model": {
+            "name": mapping_headers["name"],
+            "type": "model",
+            "origins": ["pytorch"],
+        },
+        "metric": {
+            "name": mapping_headers["metric"],
+            "benchmark_values": [mapping_headers["actual"]],
+            "target_value": mapping_headers["target"],
+        },
+    }
+
+    with open(f"{os.path.splitext(output_json_path)[0]}.json", "a") as f:
+        print(json.dumps(record), file=f)
+
+
+def write_json_result_local(output_json_path, headers, row):
+    """
+    Write the result into JSON format, so that it can be uploaded to the benchmark database
+    to be displayed on OSS dashboard. The JSON format is defined at
+    https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database
+
+    Local version (filling in dummy values for fields that should be populated by CI)
     """
     mapping_headers = {headers[i]: v for i, v in enumerate(row)}
+    today = datetime.date.today()
+    sha_hash = hashlib.sha256(str(today).encode("utf-8")).hexdigest()
+    first_second = datetime.datetime.combine(today, datetime.time.min)
+    workflow_id = int(first_second.timestamp())
+    job_id = workflow_id + 1
     record = {
+        "timestamp": int(time.time()),
+        "schema_version": "v3",
+        "name": "devvm local benchmark",
+        "repo": "pytorch/ao",
+        "head_branch": "main",
+        "head_sha": sha_hash,
+        "workflow_id": workflow_id,
+        "run_attempt": 1,
+        "job_id": job_id,
         "benchmark": {
             "name": "TorchAO benchmark",
             "mode": "inference",