Refactoring

Jonathan Makunga · Jonathan Makunga · commit 8227a4ceaf10 · 2024-03-26T23:49:19.000-07:00
diff --git a/src/sagemaker/serve/builder/jumpstart_builder.py b/src/sagemaker/serve/builder/jumpstart_builder.py
@@ -268,7 +268,7 @@ def _tune_for_js(self, max_tuning_duration: int = 1800):
             )
             return self.pysdk_model
 
-        if not sharded_supported(self.model, self.pysdk_model):
+        if not sharded_supported(self.model, self.js_model_config):
             logger.warning("Sharded is not supported for this model. Returning original model.")
             return self.pysdk_model
 
diff --git a/tests/unit/sagemaker/serve/builder/test_js_builder.py b/tests/unit/sagemaker/serve/builder/test_js_builder.py
@@ -36,7 +36,7 @@
 mock_tgi_most_performant_model_serving_properties = {
     "SAGEMAKER_PROGRAM": "inference.py",
     "SAGEMAKER_MODEL_SERVER_WORKERS": "1",
-    "SM_NUM_GPUS": "2",
+    "SM_NUM_GPUS": "4",
 }
 mock_tgi_model_serving_properties = {
     "SAGEMAKER_PROGRAM": "inference.py",

Original file line number	Diff line number	Diff line change
`@@ -268,7 +268,7 @@ def _tune_for_js(self, max_tuning_duration: int = 1800):`
`268`	`268`	`)`
`269`	`269`	`return self.pysdk_model`
`270`	`270`
`271`		`- if not sharded_supported(self.model, self.pysdk_model):`
	`271`	`+ if not sharded_supported(self.model, self.js_model_config):`
`272`	`272`	`logger.warning("Sharded is not supported for this model. Returning original model.")`
`273`	`273`	`return self.pysdk_model`
`274`	`274`
Original file line number	Diff line number	Diff line change
`@@ -36,7 +36,7 @@`
`36`	`36`	`mock_tgi_most_performant_model_serving_properties = {`
`37`	`37`	`"SAGEMAKER_PROGRAM": "inference.py",`
`38`	`38`	`"SAGEMAKER_MODEL_SERVER_WORKERS": "1",`
`39`		`- "SM_NUM_GPUS": "2",`
	`39`	`+ "SM_NUM_GPUS": "4",`
`40`	`40`	`}`
`41`	`41`	`mock_tgi_model_serving_properties = {`
`42`	`42`	`"SAGEMAKER_PROGRAM": "inference.py",`