Pretty Print JS DJL Benchmark results

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -25,16 +25,26 @@
 from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees
 from sagemaker.serve.model_server.tgi.prepare import prepare_tgi_js_resources, _create_dir_structure
 from sagemaker.serve.mode.function_pointers import Mode
-from sagemaker.serve.utils.exceptions import LocalDeepPingException, LocalModelOutOfMemoryException, \
-    LocalModelInvocationException, LocalModelLoadException, SkipTuningComboException
+from sagemaker.serve.utils.exceptions import (
+    LocalDeepPingException,
+    LocalModelOutOfMemoryException,
+    LocalModelInvocationException,
+    LocalModelLoadException,
+    SkipTuningComboException
+)
 from sagemaker.serve.utils.predictors import (
     DjlLocalModePredictor,
     TgiLocalModePredictor,
 )
 from sagemaker.serve.utils.local_hardware import _get_nb_instance, _get_ram_usage_mb
 from sagemaker.serve.utils.telemetry_logger import _capture_telemetry
-from sagemaker.serve.utils.tuning import _serial_benchmark, _concurrent_benchmark, _more_performant, \
-    _pretty_print_results_tgi, _pretty_print_results_tgi_js
+from sagemaker.serve.utils.tuning import (
+    _serial_benchmark,
+    _concurrent_benchmark,
+    _more_performant,
+    _pretty_print_results_djl_js,
+    _pretty_print_results_tgi_js
+)
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.base_predictor import PredictorBase
 from sagemaker.jumpstart.model import JumpStartModel
@@ -375,7 +385,7 @@ def tune_for_djl_jumpstart(self, max_tuning_duration: int = 1800):
                 "OPTION_TENSOR_PARALLEL_DEGREE": str(best_tuned_combination[1])
             })
 
-            _pretty_print_results_tgi(benchmark_results)
+            _pretty_print_results_djl_js(benchmark_results)
             logger.info(
                 "Model Configuration: %s was most performant with avg latency: %s, "
                 "p90 latency: %s, average tokens per second: %s, throughput/s: %s, "

src/sagemaker/serve/utils/tuning.py

@@ -135,6 +135,43 @@ def _pretty_print_results_tgi_js(results: dict):
     )
 
 
+def _pretty_print_results_djl_js(results: dict):
+    """Placeholder docstring"""
+    avg_latencies = []
+    option_tp_degree = []
+    p90s = []
+    avg_tokens_per_seconds = []
+    throughput_per_seconds = []
+    standard_deviations = []
+    ordered = collections.OrderedDict(sorted(results.items()))
+
+    for key, value in ordered.items():
+        avg_latencies.append(key)
+        option_tp_degree.append(value[0]["OPTION_TENSOR_PARALLEL_DEGREE"])
+        p90s.append(value[1])
+        avg_tokens_per_seconds.append(value[2])
+        throughput_per_seconds.append(value[3])
+        standard_deviations.append(value[4])
+
+    df = pd.DataFrame(
+        {
+            "AverageLatency (Serial)": avg_latencies,
+            "P90_Latency (Serial)": p90s,
+            "AverageTokensPerSecond (Serial)": avg_tokens_per_seconds,
+            "ThroughputPerSecond (Concurrent)": throughput_per_seconds,
+            "StandardDeviationResponse (Concurrent)": standard_deviations,
+            "OPTION_TENSOR_PARALLEL_DEGREE": option_tp_degree,
+        }
+    )
+    logger.info(
+        "\n================================================================== Benchmark "
+        "Results ==================================================================\n%s"
+        "\n============================================================================"
+        "===========================================================================\n",
+        df.to_string(),
+    )
+
+
 def _tokens_per_second(generated_text: str, max_token_length: int, latency: float) -> int:
     """Placeholder docstring"""
     est_tokens = (_tokens_from_chars(generated_text) + _tokens_from_words(generated_text)) / 2