Fix repr

Author: samruds

src/sagemaker/serve/builder/model_builder.py

@@ -20,6 +20,7 @@
 
 from pathlib import Path
 
+from accelerate.commands.estimate import estimate_command_parser, gather_data
 from sagemaker import Session
 from sagemaker.model import Model
 from sagemaker.base_predictor import PredictorBase
@@ -730,14 +731,11 @@ def _total_inference_model_size_mib(self):
         to add up to an additional 20% to the given model size as found by EleutherAI.
         """
         try:
-            import accelerate.commands.estimate.estimate_command_parser as estimate_parser
-            import accelerate.commands.estimate.gather_data as estimate_gather
-
             dtypes = self.env_vars.get("dtypes", "float32")
-            parser = estimate_parser()
+            parser = estimate_command_parser()
             args = parser.parse_args([self.model, "--dtypes", dtypes])
 
-            output = estimate_gather(
+            output = gather_data(
                 args
             )  # "dtype", "Largest Layer", "Total Size Bytes", "Training using Adam"
         except ImportError as e:

src/sagemaker/serve/builder/schema_builder.py

@@ -208,13 +208,20 @@ def _get_inverse(self, obj):
 
     def __repr__(self):
         """Placeholder docstring"""
-        return (
-            f"SchemaBuilder(\n"
-            f"input_serializer={self.input_serializer}\n"
-            f"output_serializer={self.output_serializer}\n"
-            f"input_deserializer={self.input_deserializer._deserializer}\n"
-            f"output_deserializer={self.output_deserializer._deserializer})"
-        )
+        if hasattr(self, "input_serializer") and hasattr(self, "output_serializer"):
+            return (
+                f"SchemaBuilder(\n"
+                f"input_serializer={self.input_serializer}\n"
+                f"output_serializer={self.output_serializer}\n"
+                f"input_deserializer={self.input_deserializer._deserializer}\n"
+                f"output_deserializer={self.output_deserializer._deserializer})"
+            )
+        elif hasattr(self, "custom_input_translator") and hasattr(self, "custom_output_translator"):
+            return (
+                f"SchemaBuilder(\n"
+                f"custom_input_translator={self.custom_input_translator}\n"
+                f"custom_output_translator={self.custom_output_translator}\n"
+            )
 
     def generate_marshalling_map(self) -> dict:
         """Generate marshalling map for the schema builder"""

tests/integ/sagemaker/serve/test_serve_pt_happy.py

@@ -181,7 +181,6 @@ def model_builder(request):
 #                 ), f"{caught_ex} was thrown when running pytorch squeezenet local container test"
 
 
-@pytest.mark.skip(reason="Failing test. Fix is pending.")
 @pytest.mark.skipif(
     PYTHON_VERSION_IS_NOT_310,  # or NOT_RUNNING_ON_INF_EXP_DEV_PIPELINE,
     reason="The goal of these test are to test the serving components of our feature",

tests/unit/sagemaker/serve/builder/test_model_builder.py

@@ -1343,9 +1343,8 @@ def test_build_for_transformers_happy_case_with_valid_gpu_fallback(
         self.assertEqual(model_builder._can_fit_on_single_gpu(), True)
 
     @patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers", Mock())
-    @patch("sagemaker.serve.builder.model_builder.accelerate.commands.estimate"
-           ".estimate_command_parser")
-    @patch("sagemaker.serve.builder.model_builder.accelerate.commands.estimate.gather_data")
+    @patch("sagemaker.serve.builder.model_builder.estimate_command_parser")
+    @patch("sagemaker.serve.builder.model_builder.gather_data")
     @patch("sagemaker.image_uris.retrieve")
     @patch("sagemaker.djl_inference.model.urllib")
     @patch("sagemaker.djl_inference.model.json")