Refactoring

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -22,7 +22,8 @@
 from sagemaker.model import Model
 from sagemaker import model_uris
 from sagemaker.serve.model_server.djl_serving.prepare import prepare_djl_js_resources
-from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees
+from sagemaker.serve.model_server.djl_serving.utils import _get_admissible_tensor_parallel_degrees, \
+    _get_admissible_dtypes
 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 (
@@ -42,8 +43,7 @@
     _serial_benchmark,
     _concurrent_benchmark,
     _more_performant,
-    _pretty_print_results_djl_js,
-    _pretty_print_results_tgi_js
+    _pretty_print_benchmark_results
 )
 from sagemaker.serve.utils.types import ModelServer
 from sagemaker.base_predictor import PredictorBase
@@ -254,169 +254,6 @@ def _build_for_tgi_jumpstart(self):
 
         self.pysdk_model.env.update(env)
 
-    def _log_delete_me(self, data: any):
-        """Placeholder docstring"""
-        logger.debug("*****************************************")
-        logger.debug(data)
-        logger.debug("*****************************************")
-
-    @_capture_telemetry("djl_jumpstart.tune")
-    def tune_for_djl_jumpstart(self, max_tuning_duration: int = 1800):
-        """pass"""
-        if self.mode != Mode.LOCAL_CONTAINER:
-            logger.warning(
-                "Tuning is only a %s capability. Returning original model.", Mode.LOCAL_CONTAINER
-            )
-            return self.pysdk_model
-
-        initial_model_configuration = copy.deepcopy(self.pysdk_model.env)
-        self._log_delete_me(f"initial_model_configuration: {initial_model_configuration}")
-
-        self._log_delete_me(f"self.js_model_config: {self.js_model_config}")
-        admissible_tensor_parallel_degrees = _get_admissible_tensor_parallel_degrees(self.js_model_config)
-        self._log_delete_me(f"admissible_tensor_parallel_degrees: {admissible_tensor_parallel_degrees}")
-
-        benchmark_results = {}
-        best_tuned_combination = None
-        timeout = datetime.now() + timedelta(seconds=max_tuning_duration)
-        for tensor_parallel_degree in admissible_tensor_parallel_degrees:
-            self._log_delete_me(f"tensor_parallel_degree: {tensor_parallel_degree}")
-            if datetime.now() > timeout:
-                logger.info("Max tuning duration reached. Tuning stopped.")
-                break
-
-            self.pysdk_model.env.update({
-                "OPTION_TENSOR_PARALLEL_DEGREE": str(tensor_parallel_degree)
-            })
-
-            try:
-                predictor = self.pysdk_model.deploy(
-                    model_data_download_timeout=max_tuning_duration
-                )
-
-                avg_latency, p90, avg_tokens_per_second = _serial_benchmark(
-                    predictor, self.schema_builder.sample_input
-                )
-                throughput_per_second, standard_deviation = _concurrent_benchmark(
-                    predictor, self.schema_builder.sample_input
-                )
-
-                tested_env = self.pysdk_model.env.copy()
-                logger.info(
-                    "Average latency: %s, throughput/s: %s for configuration: %s",
-                    avg_latency,
-                    throughput_per_second,
-                    tested_env,
-                )
-                benchmark_results[avg_latency] = [
-                    tested_env,
-                    p90,
-                    avg_tokens_per_second,
-                    throughput_per_second,
-                    standard_deviation,
-                ]
-
-                if not best_tuned_combination:
-                    best_tuned_combination = [
-                        avg_latency,
-                        tensor_parallel_degree,
-                        None,
-                        p90,
-                        avg_tokens_per_second,
-                        throughput_per_second,
-                        standard_deviation,
-                    ]
-                else:
-                    tuned_configuration = [
-                        avg_latency,
-                        tensor_parallel_degree,
-                        None,
-                        p90,
-                        avg_tokens_per_second,
-                        throughput_per_second,
-                        standard_deviation,
-                    ]
-                    if _more_performant(best_tuned_combination, tuned_configuration):
-                        best_tuned_combination = tuned_configuration
-            except LocalDeepPingException as e:
-                logger.warning(
-                    "Deployment unsuccessful with OPTION_TENSOR_PARALLEL_DEGREE: %s. "
-                    "Failed to invoke the model server: %s",
-                    tensor_parallel_degree,
-                    str(e),
-                )
-                break
-            except LocalModelOutOfMemoryException as e:
-                logger.warning(
-                    "Deployment unsuccessful with OPTION_TENSOR_PARALLEL_DEGREE: %s. "
-                    "Out of memory when loading the model: %s",
-                    tensor_parallel_degree,
-                    str(e),
-                )
-                break
-            except LocalModelInvocationException as e:
-                logger.warning(
-                    "Deployment unsuccessful with OPTION_TENSOR_PARALLEL_DEGREE: %s. "
-                    "Failed to invoke the model server: %s"
-                    "Please check that model server configurations are as expected "
-                    "(Ex. serialization, deserialization, content_type, accept).",
-                    tensor_parallel_degree,
-                    str(e),
-                )
-                break
-            except LocalModelLoadException as e:
-                logger.warning(
-                    "Deployment unsuccessful with OPTION_TENSOR_PARALLEL_DEGREE: %s. "
-                    "Failed to load the model: %s.",
-                    tensor_parallel_degree,
-                    str(e),
-                )
-                break
-            except SkipTuningComboException as e:
-                logger.warning(
-                    "Deployment with OPTION_TENSOR_PARALLEL_DEGREE: %s. "
-                    "was expected to be successful. However failed with: %s. "
-                    "Trying next combination.",
-                    tensor_parallel_degree,
-                    str(e),
-                )
-                break
-            except Exception:
-                logger.exception(
-                    "Deployment unsuccessful with OPTION_TENSOR_PARALLEL_DEGREE: %s. "
-                    "with uncovered exception",
-                    tensor_parallel_degree
-                )
-                break
-
-        if best_tuned_combination:
-            self.pysdk_model.env.update({
-                "OPTION_TENSOR_PARALLEL_DEGREE": str(best_tuned_combination[1])
-            })
-
-            _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, "
-                "standard deviation of request %s",
-                self.pysdk_model.env,
-                best_tuned_combination[0],
-                best_tuned_combination[3],
-                best_tuned_combination[4],
-                best_tuned_combination[5],
-                best_tuned_combination[6],
-            )
-        else:
-            self.pysdk_model.env.update(initial_model_configuration)
-            logger.debug(
-                "Failed to gather any tuning results. "
-                "Please inspect the stack trace emitted from live logging for more details. "
-                "Falling back to default model environment variable configurations: %s",
-                self.pysdk_model.env,
-            )
-
-        return self.pysdk_model
-
     @_capture_telemetry("tgi_jumpstart.tune")
     def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
         """Placeholder docstring"""
@@ -426,19 +263,13 @@ def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
             )
             return self.pysdk_model
 
-        initial_model_configuration = copy.deepcopy(self.pysdk_model.env)
-        self._log_delete_me(f"initial_model_configuration: {initial_model_configuration}")
-
-        self._log_delete_me(f"self.js_model_config: {self.js_model_config}")
-
+        initial_env_vars = copy.deepcopy(self.pysdk_model.env)
         admissible_tensor_parallel_degrees = _get_admissible_tensor_parallel_degrees(self.js_model_config)
-        self._log_delete_me(f"admissible_tensor_parallel_degrees: {admissible_tensor_parallel_degrees}")
 
         benchmark_results = {}
         best_tuned_combination = None
         timeout = datetime.now() + timedelta(seconds=max_tuning_duration)
         for tensor_parallel_degree in admissible_tensor_parallel_degrees:
-            self._log_delete_me(f"tensor_parallel_degree: {tensor_parallel_degree}")
             if datetime.now() > timeout:
                 logger.info("Max tuning duration reached. Tuning stopped.")
                 break
@@ -554,7 +385,9 @@ def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
                 "SM_NUM_GPUS": str(best_tuned_combination[1])
             })
 
-            _pretty_print_results_tgi_js(benchmark_results)
+            _pretty_print_benchmark_results(benchmark_results, [
+                "SM_NUM_GPUS"
+            ])
             logger.info(
                 "Model Configuration: %s was most performant with avg latency: %s, "
                 "p90 latency: %s, average tokens per second: %s, throughput/s: %s, "
@@ -567,7 +400,7 @@ def tune_for_tgi_jumpstart(self, max_tuning_duration: int = 1800):
                 best_tuned_combination[6],
             )
         else:
-            self.pysdk_model.env.update(initial_model_configuration)
+            self.pysdk_model.env.update(initial_env_vars)
             logger.debug(
                 "Failed to gather any tuning results. "
                 "Please inspect the stack trace emitted from live logging for more details. "
@@ -612,6 +445,4 @@ def _build_for_jumpstart(self):
 
         if self.model_server == ModelServer.TGI:
             self.pysdk_model.tune = self.tune_for_tgi_jumpstart
-        elif self.model_server == ModelServer.DJL_SERVING:
-            self.pysdk_model.tune = self.tune_for_djl_jumpstart
         return self.pysdk_model

src/sagemaker/serve/utils/tuning.py

@@ -98,47 +98,16 @@ def _pretty_print_results_tgi(results: dict):
     )
 
 
-def _pretty_print_results_tgi_js(results: dict):
+def _pretty_print_benchmark_results(results: dict, model_env_vars=None):
     """Placeholder docstring"""
-    avg_latencies = []
-    sm_num_gpus = []
-    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)
-        sm_num_gpus.append(value[0]["SM_NUM_GPUS"])
-        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,
-            "SM_NUM_GPUS": sm_num_gpus,
-        }
-    )
-    logger.info(
-        "\n================================================================== Benchmark "
-        "Results ==================================================================\n%s"
-        "\n============================================================================"
-        "===========================================================================\n",
-        df.to_string(),
-    )
+    if model_env_vars is None:
+        model_env_vars = []
 
+    __env_var_data = {}
+    for e in model_env_vars:
+        __env_var_data[e] = []
 
-def _pretty_print_results_djl_js(results: dict):
-    """Placeholder docstring"""
     avg_latencies = []
-    option_tp_degree = []
     p90s = []
     avg_tokens_per_seconds = []
     throughput_per_seconds = []
@@ -147,20 +116,22 @@ def _pretty_print_results_djl_js(results: dict):
 
     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])
 
+        for k in __env_var_data:
+            __env_var_data[k].append(value[0][k])
+
     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,
+            **__env_var_data
         }
     )
     logger.info(