Format

Author: samruds

src/sagemaker/serve/builder/model_builder.py

@@ -43,8 +43,11 @@
 from sagemaker.serve.utils import task
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
-from sagemaker.serve.utils.hardware_detector import _get_gpu_info, _get_gpu_info_fallback,\
-    _total_inference_model_size_mib
+from sagemaker.serve.utils.hardware_detector import (
+    _get_gpu_info,
+    _get_gpu_info_fallback,
+    _total_inference_model_size_mib,
+)
 from sagemaker.serve.detector.image_detector import (
     auto_detect_container,
     _detect_framework_and_version,
@@ -724,12 +727,15 @@ def _can_fit_on_single_gpu(self) -> Type[bool]:
         """
         try:
             single_gpu_size_mib = self._try_fetch_gpu_info()
-            if _total_inference_model_size_mib(self.model, self.env_vars.get("dtypes", "float32")) \
-                    <= single_gpu_size_mib:
+            if (
+                _total_inference_model_size_mib(self.model, self.env_vars.get("dtypes", "float32"))
+                <= single_gpu_size_mib
+            ):
                 logger.info(
                     "Total inference model size MIB %s, single GPU size for instance MIB %s",
-                    _total_inference_model_size_mib(self.model, self.env_vars.get("dtypes",
-                                                                                 "float32")),
+                    _total_inference_model_size_mib(
+                        self.model, self.env_vars.get("dtypes", "float32")
+                    ),
                     single_gpu_size_mib,
                 )
                 return True

src/sagemaker/serve/utils/hardware_detector.py

@@ -33,6 +33,7 @@
     "size calculations e.g. pip install 'sagemaker[huggingface]'"
 )
 
+
 def _get_gpu_info(instance_type: str, session: Session) -> Tuple[int, int]:
     """Get GPU info for the provided instance