Skip JS model mapping with env vars or image URI provided (#1599)

xiongz945 · benieric · nargokul · pintaoz-aws · commit 2438a3fc81a3 · 2024-12-04T01:44:04.000-08:00
* Base model trainer (#1521) * Base model trainer * flake8 * add testing notebook * add param validation & set defaults * Implement simple train method * feature: support script mode with local train.sh (#1523) * feature: support script mode with local train.sh * Stop tracking train.sh and add it to .gitignore * update message * make dir if not exist * fix docs * fix: docstyle * Address comments * fix hyperparams * Revert pydantic custom error * pylint * Image Spec refactoring and updates (#1525) * Image Spec refactoring and updates * Unit tests and update function for Image Spec * Fix hugging face test * Fix Tests * Add unit tests for ModelTrainer (#1527) * Add unit tests for ModelTrainer * Flake8 * format * Add example notebook (#1528) * Add testing notebook * format * use smaller data * remove large dataset * update * pylint * flake8 * ignore docstyle in directories with test * format * format * Add enviornment variable bootstrapping script (#1530) * Add enviornment variables scripts * format * fix comment * add docstrings * fix comment * feature: add utility function to capture local snapshot (#1524) * local snapshot * Update pip list command * Remove function calls * Address comments * Address comments * Support intelligent parameters (#1540) * Support intelligent parameters * fix codestyle * Revert Image Spec (#1541) * Cleanup ModelTrainer (#1542) * General image builder (#1546) * General image builder * General image builder * Fix codestyle * Fix codestyle * Move location * Add warnings * Add integ tests * Fix integ test * Fix integ test * Fix region error * Add region * Latest Container Image (#1545) * Latest Container Image * Test Fixes * Parameterized tests and some logic updates * Test fixes * Move to Image URI * Fixes for unit test * Fixes for unit test * Fix codestyle error checks * Cleanup ModelTrainer code (#1552) * feat: add pre-processing and post-processing logic to inference_spec (#1560) * add pre-processing and post-processing logic to inference_spec * fix format * make accept_type and content_type optional * remove accept_type and content_type from pre/post processing * correct typo * Add Distributed Training Support Model Trainer (#1536) * Add path to set Additional Settings in ModelTrainer (#1555) * feature: support HuggingFace models with JumpStart configs * Update bucket name for the model mapping * Mask Sensitive Env Logs in Container (#1568) * Fix unit test * Fix bug in script mode setup ModelTrainer (#1575) * Save mapping as attribute * Fix style issues * Fix style issues * Fix: bypass jumpstart mapping when not in endpoint mode * Skip JS model mapping with env vars or image URI provided * Revert "Merge branch 'aws:master' into dev-morpheus" This reverts commit 26a0b0bb37e0343b3287f5c5c484df22726fc858, reversing changes made to d19d4e178442be4b6e1d07d55498dd76dfac50f0. * Merge branch 'aws:master' into dev-morpheus This reverts commit 076442bd83e5ca977bf5b6ce1b716474d2794feb. * Rebase on master-morpheus * Fix unit test description * Fix TEI integ test * Fix style issue * Fix style issues * Fix schema builder integ tests * Fix TEI integ test * Fix code style issue --------- Co-authored-by: Erick Benitez-Ramos <141277478+benieric@users.noreply.github.com> Co-authored-by: Gokul Anantha Narayanan <166456257+nargokul@users.noreply.github.com> Co-authored-by: pintaoz-aws <167920275+pintaoz-aws@users.noreply.github.com> Co-authored-by: Pravali Uppugunduri <46845440+pravali96@users.noreply.github.com> Co-authored-by: Xiong Zeng <xionzeng@amazon.com> Co-authored-by: Gary Wang <38331932+gwang111@users.noreply.github.com>
diff --git a/src/sagemaker/serve/builder/model_builder.py b/src/sagemaker/serve/builder/model_builder.py
@@ -872,7 +872,6 @@ def build(  # pylint: disable=R0911
         Returns:
             Type[Model]: A deployable ``Model`` object.
         """
-        from sagemaker.modules.train.model_trainer import ModelTrainer
 
         self.modes = dict()
 
@@ -1728,10 +1727,24 @@ def _use_jumpstart_equivalent(self):
 
         Replace it with the equivalent if there's one
         """
+        # Do not use the equivalent JS model if image_uri or env_vars is provided
+        if self.image_uri or self.env_vars:
+            return False
         if not hasattr(self, "_has_jumpstart_equivalent"):
             self._jumpstart_mapping = self._retrieve_hugging_face_model_mapping()
             self._has_jumpstart_equivalent = self.model in self._jumpstart_mapping
         if self._has_jumpstart_equivalent:
+            # Use schema builder from HF model metadata
+            if not self.schema_builder:
+                model_task = None
+                if self.model_metadata:
+                    model_task = self.model_metadata.get("HF_TASK")
+                hf_model_md = get_huggingface_model_metadata(self.model)
+                if not model_task:
+                    model_task = hf_model_md.get("pipeline_tag")
+                if model_task:
+                    self._hf_schema_builder_init(model_task)
+
             huggingface_model_id = self.model
             jumpstart_model_id = self._jumpstart_mapping[huggingface_model_id]["jumpstart-model-id"]
             self.model = jumpstart_model_id
@@ -1743,7 +1756,7 @@ def _use_jumpstart_equivalent(self):
                 "artifact S3 URI and compare them."
             )
             logger.warning(  # pylint: disable=logging-fstring-interpolation
-                "Please note that for this model we are using the JumpStart's"
+                "Please note that for this model we are using the JumpStart's "
                 f'local copy "{jumpstart_model_id}" '
                 f'of the HuggingFace model "{huggingface_model_id}" you chose. '
                 "We strive to keep our local copy synced with the HF model hub closely. "
diff --git a/tests/integ/sagemaker/serve/test_schema_builder.py b/tests/integ/sagemaker/serve/test_schema_builder.py
@@ -111,13 +111,13 @@ def test_model_builder_happy_path_with_task_provided_local_schema_mode(
             if container_startup_timeout:
                 predictor = model.deploy(
                     role=role_arn,
-                    instance_count=1,
+                    initial_instance_count=1,
                     instance_type=instance_type_provided,
                     container_startup_health_check_timeout=container_startup_timeout,
                 )
             else:
                 predictor = model.deploy(
-                    role=role_arn, instance_count=1, instance_type=instance_type_provided
+                    role=role_arn, initial_instance_count=1, instance_type=instance_type_provided
                 )
 
             predicted_outputs = predictor.predict(inputs)
@@ -181,7 +181,7 @@ def test_model_builder_happy_path_with_task_provided_remote_schema_mode(
 
             logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
             predictor = model.deploy(
-                role=role_arn, instance_count=1, instance_type=instance_type_provided
+                role=role_arn, initial_instance_count=1, instance_type=instance_type_provided
             )
 
             predicted_outputs = predictor.predict(inputs)
diff --git a/tests/integ/sagemaker/serve/test_serve_tei.py b/tests/integ/sagemaker/serve/test_serve_tei.py
@@ -44,6 +44,10 @@ def model_builder_model_schema_builder():
         model_path=HF_DIR,
         model="BAAI/bge-m3",
         schema_builder=SchemaBuilder(sample_input, loaded_response),
+        env_vars={
+            # Add this to bypass JumpStart model mapping
+            "HF_MODEL_ID": "BAAI/bge-m3"
+        },
     )
 
 

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,10 @@ def model_builder_model_schema_builder():`
`44`	`44`	`model_path=HF_DIR,`
`45`	`45`	`model="BAAI/bge-m3",`
`46`	`46`	`schema_builder=SchemaBuilder(sample_input, loaded_response),`
	`47`	`+ env_vars={`
	`48`	`+ # Add this to bypass JumpStart model mapping`
	`49`	`+ "HF_MODEL_ID": "BAAI/bge-m3"`
	`50`	`+ },`
`47`	`51`	`)`
`48`	`52`
`49`	`53`