chore: update docstrings, add tests

Author: evakravi

src/sagemaker/jumpstart/model.py

@@ -412,6 +412,9 @@ def attach(
         model_version: Optional[str] = None,
         sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
     ) -> "JumpStartModel":
+        """Attaches a JumpStartModel object to an existing SageMaker Endpoint.
+
+        The model id, version (and inference component name) can be inferred from the tags."""
 
         inferred_model_id, inferred_model_version, inferred_inference_component_name = (
             get_model_id_version_from_endpoint(

src/sagemaker/model.py

@@ -408,6 +408,7 @@ def attach(
         inference_component_name: Optional[str] = None,
         sagemaker_session=None,
     ) -> "Model":
+        """Attaches a Model object to an existing SageMaker Endpoint."""
         raise NotImplementedError
 
     @runnable_by_pipeline

tests/integ/sagemaker/jumpstart/model/test_jumpstart_model.py

@@ -219,6 +219,11 @@ def test_jumpstart_gated_model_inference_component_enabled(setup):
 
     assert response is not None
 
+    model = JumpStartModel.attach(predictor.endpoint_name, sagemaker_session=get_sm_session())
+    assert model.model_id == model_id
+    assert model.endpoint_name == predictor.endpoint_name
+    assert model.inference_component_name == predictor.inference_component_name
+
 
 @mock.patch("sagemaker.jumpstart.cache.JUMPSTART_LOGGER.warning")
 def test_instatiating_model(mock_warning_logger, setup):

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -1292,6 +1292,28 @@ def test_model_artifact_variant_model(
             enable_network_isolation=True,
         )
 
+    @mock.patch("sagemaker.jumpstart.model.get_model_id_version_from_endpoint")
+    @mock.patch("sagemaker.jumpstart.model.JumpStartModel.__init__")
+    def test_attach(
+        self,
+        mock_js_model_init,
+        mock_get_model_id_version_from_endpoint,
+    ):
+        mock_js_model_init.return_value = None
+        mock_get_model_id_version_from_endpoint.return_value = "model-id", "model-version", None
+        val = JumpStartModel.attach("some-endpoint")
+        mock_get_model_id_version_from_endpoint.assert_called_once_with(
+            endpoint_name="some-endpoint",
+            inference_component_name=None,
+            sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+        )
+        mock_js_model_init.assert_called_once_with(
+            model_id="model-id",
+            model_version="model-version",
+            sagemaker_session=DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+        )
+        assert isinstance(val, JumpStartModel)
+
     @mock.patch("sagemaker.jumpstart.model.validate_model_id_and_get_type")
     @mock.patch("sagemaker.jumpstart.factory.model.Session")
     @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")