Lint, and add _s suffix to metrics

benchmarks/experiment_runner.py

@@ -254,7 +254,9 @@ def dump_profile_info(self, prof, model_name):
   def collect_profile_to_metrics(self, prof, metrics):
     assert prof is not None, 'Expecting profiler to be defined!'
     if not self._args.profile_cuda_cpu_collect:
-      logger.warning('Profiling enabled, but collection of CPU/CUDA profiling info disabled.')
+      logger.warning(
+          'Profiling enabled, but collection of CPU/CUDA profiling info disabled.'
+      )
       return
 
     kernel_dump = prof.profiler.total_average()
@@ -276,10 +278,12 @@ def collect_profile_to_metrics(self, prof, metrics):
 
     total_cpu_time /= 1000000
     total_cuda_time /= 1000000
-    metrics["total_cpu_time"] = total_cpu_time
-    metrics["total_cuda_time"] = total_cuda_time
-    metrics["per_iter_cpu_time"] = total_cpu_time / self._args.iterations_per_run
-    metrics["per_iter_cuda_time"] = total_cuda_time / self._args.iterations_per_run
+    metrics["total_cpu_time_s"] = total_cpu_time
+    metrics["total_cuda_time_s"] = total_cuda_time
+    metrics[
+        "per_iter_cpu_time_s"] = total_cpu_time / self._args.iterations_per_run
+    metrics[
+        "per_iter_cuda_time_s"] = total_cuda_time / self._args.iterations_per_run
 
   def timed_run(self, benchmark_experiment, benchmark_model):
     reset_rng_state(benchmark_experiment)
@@ -323,7 +327,6 @@ def loop(prof=None):
     else:
       output = loop()
 
-
     t_end = time.perf_counter()
     if enable_prof:
       self.dump_profile_info(prof, benchmark_model.model_name)

benchmarks/result_analyzer.py

@@ -88,18 +88,22 @@ def get_calculated_metrics(self, d, dataline):
       d["xla_median_trace_per_iter_time"] = -1
       d["xla_compile_time"] = -1
 
-    if "total_cpu_time" in dataline["metrics"]:
-      total_cpu_time = np.asarray(dataline["metrics"]["total_cpu_time"], dtype="float")
-      d["median_total_cpu_time"] = np.median(total_cpu_time)
-    if "per_iter_cpu_time" in dataline["metrics"]:
-      per_iter_cpu_time = np.asarray(dataline["metrics"]["per_iter_cpu_time"], dtype="float")
-      d["median_per_iter_cpu_time"] = np.median(per_iter_cpu_time)
-    if "total_cuda_time" in dataline["metrics"]:
-      total_cuda_time = np.asarray(dataline["metrics"]["total_cuda_time"], dtype="float")
-      d["median_total_cuda_time"] = np.median(total_cuda_time)
-    if "per_iter_cuda_time" in dataline["metrics"]:
-      per_iter_cuda_time = np.asarray(dataline["metrics"]["per_iter_cuda_time"], dtype="float")
-      d["median_per_iter_cuda_time"] = np.median(per_iter_cuda_time)
+    if "total_cpu_time_s" in dataline["metrics"]:
+      total_cpu_time = np.asarray(
+          dataline["metrics"]["total_cpu_time_s"], dtype="float")
+      d["median_total_cpu_time_s"] = np.median(total_cpu_time)
+    if "per_iter_cpu_time_s" in dataline["metrics"]:
+      per_iter_cpu_time = np.asarray(
+          dataline["metrics"]["per_iter_cpu_time_s"], dtype="float")
+      d["median_per_iter_cpu_time_s"] = np.median(per_iter_cpu_time)
+    if "total_cuda_time_s" in dataline["metrics"]:
+      total_cuda_time = np.asarray(
+          dataline["metrics"]["total_cuda_time_s"], dtype="float")
+      d["median_total_cuda_time_s"] = np.median(total_cuda_time)
+    if "per_iter_cuda_time_s" in dataline["metrics"]:
+      per_iter_cuda_time = np.asarray(
+          dataline["metrics"]["per_iter_cuda_time_s"], dtype="float")
+      d["median_per_iter_cuda_time_s"] = np.median(per_iter_cuda_time)
 
     if dataline["experiment"]["dynamo"]:
       d["dynamo_compile_time"] = np.max(total_time) - np.median(total_time)