Revert "feat: Add Optional task to Model"

This reverts commit fd3e86b.

Author: Xiong Zeng

src/sagemaker/model.py

@@ -156,7 +156,6 @@ def __init__(
         dependencies: Optional[List[str]] = None,
         git_config: Optional[Dict[str, str]] = None,
         resources: Optional[ResourceRequirements] = None,
-        task: Optional[Union[str, PipelineVariable]] = None,
     ):
         """Initialize an SageMaker ``Model``.
 
@@ -320,9 +319,7 @@ def __init__(
                 for a model to be deployed to an endpoint. Only
                 EndpointType.INFERENCE_COMPONENT_BASED supports this feature.
                 (Default: None).
-            task (str or PipelineVariable): Task values used to override the HuggingFace task
-                Examples are: "audio-classification", "depth-estimation",
-                "feature-extraction" etc. (default: None).
+
         """
         self.model_data = model_data
         self.image_uri = image_uri
@@ -399,7 +396,6 @@ def __init__(
         self.content_types = None
         self.response_types = None
         self.accept_eula = None
-        self.task = task
 
     @runnable_by_pipeline
     def register(

src/sagemaker/serve/builder/model_builder.py

@@ -205,7 +205,6 @@ class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
             "help": (
                 'Model object with "predict" method to perform inference '
                 "or HuggingFace/JumpStart Model ID"
-                "or HuggingFace Task to override"
             )
         },
     )

tests/integ/sagemaker/serve/test_serve_transformers.py

@@ -13,7 +13,6 @@
 from __future__ import absolute_import
 
 import pytest
-from sagemaker.model import Model
 from sagemaker.serve.builder.schema_builder import SchemaBuilder
 from sagemaker.serve.builder.model_builder import ModelBuilder, Mode
 
@@ -29,11 +28,6 @@
 
 logger = logging.getLogger(__name__)
 
-MODEL_DATA = "s3://bucket/model.tar.gz"
-MODEL_IMAGE = "mi"
-ROLE = "some-role"
-HF_TASK = "fill-mask"
-
 sample_input = {
     "inputs": "The man worked as a [MASK].",
 }
@@ -86,17 +80,6 @@ def model_builder_model_schema_builder():
     )
 
 
-@pytest.fixture
-def model_builder_model_with_task_builder():
-    model = Model(
-        MODEL_IMAGE, MODEL_DATA, task=HF_TASK, name="bert-base-uncased", role=ROLE
-    )
-    return ModelBuilder(
-        model_path=HF_DIR,
-        model=model,
-    )
-
-
 @pytest.fixture
 def model_builder(request):
     return request.getfixturevalue(request.param)
@@ -139,42 +122,3 @@ def test_pytorch_transformers_sagemaker_endpoint(
                 assert (
                     False
                 ), f"{caught_ex} was thrown when running pytorch transformers sagemaker endpoint test"
-
-
-@pytest.mark.skipif(
-    PYTHON_VERSION_IS_NOT_310,
-    reason="Testing Optional task",
-)
-@pytest.mark.parametrize("model_builder", ["model_builder_model_with_task_builder"], indirect=True)
-def test_happy_path_with_task_sagemaker_endpoint(
-        sagemaker_session, model_builder, gpu_instance_type, input
-):
-    logger.info("Running in SAGEMAKER_ENDPOINT mode...")
-    caught_ex = None
-
-    iam_client = sagemaker_session.boto_session.client("iam")
-    role_arn = iam_client.get_role(RoleName="SageMakerRole")["Role"]["Arn"]
-
-    model = model_builder.build(
-        mode=Mode.SAGEMAKER_ENDPOINT, role_arn=role_arn, sagemaker_session=sagemaker_session
-    )
-
-    with timeout(minutes=SERVE_SAGEMAKER_ENDPOINT_TIMEOUT):
-        try:
-            logger.info("Deploying and predicting in SAGEMAKER_ENDPOINT mode...")
-            predictor = model.deploy(instance_type=gpu_instance_type, initial_instance_count=1)
-            logger.info("Endpoint successfully deployed.")
-            predictor.predict(input)
-        except Exception as e:
-            caught_ex = e
-        finally:
-            cleanup_model_resources(
-                sagemaker_session=model_builder.sagemaker_session,
-                model_name=model.name,
-                endpoint_name=model.endpoint_name,
-            )
-            if caught_ex:
-                logger.exception(caught_ex)
-                assert (
-                    False
-                ), f"{caught_ex} was thrown when running pytorch transformers sagemaker endpoint test"

tests/unit/sagemaker/model/test_deploy.py

@@ -85,7 +85,6 @@
     },
     limits={},
 )
-HF_TASK = "audio-classification"
 
 
 @pytest.fixture
@@ -1028,57 +1027,3 @@ def test_deploy_with_name_and_resources(sagemaker_session):
         async_inference_config_dict=None,
         live_logging=False,
     )
-
-
-@patch("sagemaker.model.Model._create_sagemaker_model", Mock())
-@patch("sagemaker.production_variant", return_value=BASE_PRODUCTION_VARIANT)
-def test_deploy_with_name_and_task(sagemaker_session):
-    sagemaker_session.sagemaker_config = {}
-
-    model = Model(
-        MODEL_IMAGE, MODEL_DATA, task=HF_TASK, name=MODEL_NAME, role=ROLE, sagemaker_session=sagemaker_session
-    )
-
-    endpoint_name = "testing-task-input"
-    predictor = model.deploy(
-        endpoint_name=endpoint_name,
-        instance_type=INSTANCE_TYPE,
-        initial_instance_count=INSTANCE_COUNT,
-    )
-
-    sagemaker_session.create_model.assert_called_with(
-        name=MODEL_IMAGE,
-        role=ROLE,
-        task=HF_TASK
-    )
-
-    assert isinstance(predictor, sagemaker.predictor.Predictor)
-    assert predictor.endpoint_name == endpoint_name
-    assert predictor.sagemaker_session == sagemaker_session
-
-
-@patch("sagemaker.model.Model._create_sagemaker_model", Mock())
-@patch("sagemaker.production_variant", return_value=BASE_PRODUCTION_VARIANT)
-def test_deploy_with_name_and_without_task(sagemaker_session):
-    sagemaker_session.sagemaker_config = {}
-
-    model = Model(
-        MODEL_IMAGE, MODEL_DATA, name=MODEL_NAME, role=ROLE, sagemaker_session=sagemaker_session
-    )
-
-    endpoint_name = "testing-without-task-input"
-    predictor = model.deploy(
-        endpoint_name=endpoint_name,
-        instance_type=INSTANCE_TYPE,
-        initial_instance_count=INSTANCE_COUNT,
-    )
-
-    sagemaker_session.create_model.assert_called_with(
-        name=MODEL_IMAGE,
-        role=ROLE,
-        task=None,
-    )
-
-    assert isinstance(predictor, sagemaker.predictor.Predictor)
-    assert predictor.endpoint_name == endpoint_name
-    assert predictor.sagemaker_session == sagemaker_session