Refactoring

Author: Jonathan Makunga

src/sagemaker/serve/builder/jumpstart_builder.py

@@ -783,11 +783,10 @@ def _optimize_for_jumpstart(
                     "AcceptEula": True
                 }
 
-        if quantization_config or is_compilation:
-            optimization_env_vars = _update_environment_variables(
-                optimization_env_vars, override_env
-            )
+        optimization_env_vars = _update_environment_variables(optimization_env_vars, override_env)
+        if optimization_env_vars:
             self.pysdk_model.env.update(optimization_env_vars)
+        if quantization_config or is_compilation:
             return create_optimization_job_args
         return None
 

src/sagemaker/serve/builder/model_builder.py

@@ -1258,10 +1258,8 @@ def _model_builder_optimize_wrapper(
             )
 
         if input_args:
-            print(input_args)
             self.sagemaker_session.sagemaker_client.create_optimization_job(**input_args)
             job_status = self.sagemaker_session.wait_for_optimization_job(job_name)
-            print(job_status)
             return _generate_optimized_model(self.pysdk_model, job_status)
 
         self.pysdk_model.remove_tag_with_key(Tag.OPTIMIZATION_JOB_NAME)

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

@@ -2653,22 +2653,20 @@ def test_optimize_for_hf_with_custom_s3_path(
             model_metadata={
                 "CUSTOM_MODEL_PATH": "s3://bucket/path/",
             },
+            role_arn="role-arn",
+            instance_type="ml.g5.2xlarge",
         )
 
         model_builder.pysdk_model = mock_pysdk_model
 
         out_put = model_builder._optimize_for_hf(
             job_name="job_name-123",
-            instance_type="ml.g5.2xlarge",
-            role_arn="role-arn",
             quantization_config={
                 "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
             },
             output_path="s3://bucket/code/",
         )
 
-        print(out_put)
-
         self.assertEqual(model_builder.role_arn, "role-arn")
         self.assertEqual(model_builder.instance_type, "ml.g5.2xlarge")
         self.assertEqual(model_builder.pysdk_model.env["OPTION_QUANTIZE"], "awq")
@@ -2715,14 +2713,14 @@ def test_optimize_for_hf_without_custom_s3_path(
         model_builder = ModelBuilder(
             model="meta-llama/Meta-Llama-3-8B-Instruct",
             env_vars={"HUGGING_FACE_HUB_TOKEN": "token"},
+            role_arn="role-arn",
+            instance_type="ml.g5.2xlarge",
         )
 
         model_builder.pysdk_model = mock_pysdk_model
 
         out_put = model_builder._optimize_for_hf(
             job_name="job_name-123",
-            instance_type="ml.g5.2xlarge",
-            role_arn="role-arn",
             quantization_config={
                 "OverrideEnvironment": {"OPTION_QUANTIZE": "awq"},
             },