Skip to content

fix: Enable S3 Upload on local container -> sagemaker endpoint overwrite for TGI, TEI, MMS #4814

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 2, 2024
Merged

fix: Enable S3 Upload on local container -> sagemaker endpoint overwrite for TGI, TEI, MMS #4814

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 16 additions & 4 deletions src/sagemaker/serve/builder/tei_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ def __init__(self):
self.role_arn = None

@abstractmethod
def _prepare_for_mode(self):
def _prepare_for_mode(self, *args, **kwargs):
"""Placeholder docstring"""

@abstractmethod
Expand Down Expand Up @@ -164,15 +164,24 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
del kwargs["role"]

if not _is_optimized(self.pysdk_model):
self._prepare_for_mode()
env_vars = {}
if str(Mode.LOCAL_CONTAINER) in self.modes:
# upload model artifacts to S3 if LOCAL_CONTAINER -> SAGEMAKER_ENDPOINT
self.pysdk_model.model_data, env_vars = self._prepare_for_mode(
model_path=self.model_path, should_upload_artifacts=True
)
else:
_, env_vars = self._prepare_for_mode()

self.env_vars.update(env_vars)
self.pysdk_model.env.update(self.env_vars)

# if the weights have been cached via local container mode -> set to offline
if str(Mode.LOCAL_CONTAINER) in self.modes:
self.pysdk_model.env.update({"TRANSFORMERS_OFFLINE": "1"})
self.pysdk_model.env.update({"HF_HUB_OFFLINE": "1"})
else:
# if has not been built for local container we must use cache
# that hosting has write access to.
self.pysdk_model.env["TRANSFORMERS_CACHE"] = "/tmp"
self.pysdk_model.env["HF_HOME"] = "/tmp"
self.pysdk_model.env["HUGGINGFACE_HUB_CACHE"] = "/tmp"

Expand All @@ -191,6 +200,9 @@ def _tei_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa

predictor = self._original_deploy(*args, **kwargs)

if "HF_HUB_OFFLINE" in self.pysdk_model.env:
self.pysdk_model.env.update({"HF_HUB_OFFLINE": "0"})

predictor.serializer = serializer
predictor.deserializer = deserializer
return predictor
Expand Down
20 changes: 15 additions & 5 deletions src/sagemaker/serve/builder/tgi_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def __init__(self):
self.role_arn = None

@abstractmethod
def _prepare_for_mode(self):
def _prepare_for_mode(self, *args, **kwargs):
"""Placeholder docstring"""

@abstractmethod
Expand Down Expand Up @@ -203,15 +203,24 @@ def _tgi_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa
del kwargs["role"]

if not _is_optimized(self.pysdk_model):
self._prepare_for_mode()
env_vars = {}
if str(Mode.LOCAL_CONTAINER) in self.modes:
# upload model artifacts to S3 if LOCAL_CONTAINER -> SAGEMAKER_ENDPOINT
self.pysdk_model.model_data, env_vars = self._prepare_for_mode(
model_path=self.model_path, should_upload_artifacts=True
)
else:
_, env_vars = self._prepare_for_mode()

self.env_vars.update(env_vars)
self.pysdk_model.env.update(self.env_vars)

# if the weights have been cached via local container mode -> set to offline
if str(Mode.LOCAL_CONTAINER) in self.modes:
self.pysdk_model.env.update({"TRANSFORMERS_OFFLINE": "1"})
self.pysdk_model.env.update({"HF_HUB_OFFLINE": "1"})
else:
# if has not been built for local container we must use cache
# that hosting has write access to.
self.pysdk_model.env["TRANSFORMERS_CACHE"] = "/tmp"
self.pysdk_model.env["HF_HOME"] = "/tmp"
self.pysdk_model.env["HUGGINGFACE_HUB_CACHE"] = "/tmp"

Expand Down Expand Up @@ -242,7 +251,8 @@ def _tgi_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[PredictorBa

predictor = self._original_deploy(*args, **kwargs)

self.pysdk_model.env.update({"TRANSFORMERS_OFFLINE": "0"})
if "HF_HUB_OFFLINE" in self.pysdk_model.env:
self.pysdk_model.env.update({"HF_HUB_OFFLINE": "0"})

predictor.serializer = serializer
predictor.deserializer = deserializer
Expand Down
28 changes: 23 additions & 5 deletions src/sagemaker/serve/builder/transformers_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,7 @@ def __init__(self):
self.shared_libs = None

@abstractmethod
def _prepare_for_mode(self):
def _prepare_for_mode(self, *args, **kwargs):
"""Abstract method"""

def _create_transformers_model(self) -> Type[Model]:
Expand Down Expand Up @@ -206,8 +206,6 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
else:
raise ValueError("Mode %s is not supported!" % overwrite_mode)

self._set_instance()

serializer = self.schema_builder.input_serializer
deserializer = self.schema_builder._output_deserializer
if self.mode == Mode.LOCAL_CONTAINER:
Expand All @@ -227,14 +225,32 @@ def _transformers_model_builder_deploy_wrapper(self, *args, **kwargs) -> Type[Pr
)
return predictor

self._set_instance(kwargs)

if "mode" in kwargs:
del kwargs["mode"]
if "role" in kwargs:
self.pysdk_model.role = kwargs.get("role")
del kwargs["role"]

if not _is_optimized(self.pysdk_model):
self._prepare_for_mode()
env_vars = {}
if str(Mode.LOCAL_CONTAINER) in self.modes:
# upload model artifacts to S3 if LOCAL_CONTAINER -> SAGEMAKER_ENDPOINT
self.pysdk_model.model_data, env_vars = self._prepare_for_mode(
model_path=self.model_path, should_upload_artifacts=True
)
else:
_, env_vars = self._prepare_for_mode()

self.env_vars.update(env_vars)
self.pysdk_model.env.update(self.env_vars)

if (
"SAGEMAKER_SERVE_SECRET_KEY" in self.pysdk_model.env
and not self.pysdk_model.env["SAGEMAKER_SERVE_SECRET_KEY"]
):
del self.pysdk_model.env["SAGEMAKER_SERVE_SECRET_KEY"]

if "endpoint_logging" not in kwargs:
kwargs["endpoint_logging"] = True
Expand Down Expand Up @@ -279,9 +295,11 @@ def _build_transformers_env(self):

return self.pysdk_model

def _set_instance(self, **kwargs):
def _set_instance(self, kwargs):
"""Set the instance : Given the detected notebook type or provided instance type"""
if self.mode == Mode.SAGEMAKER_ENDPOINT:
if "instance_type" in kwargs:
return
if self.nb_instance_type and "instance_type" not in kwargs:
kwargs.update({"instance_type": self.nb_instance_type})
logger.info("Setting instance type to %s", self.nb_instance_type)
Expand Down
1 change: 0 additions & 1 deletion src/sagemaker/serve/model_server/tei/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
MODE_DIR_BINDING = "/opt/ml/model/"
_SHM_SIZE = "2G"
_DEFAULT_ENV_VARS = {
"TRANSFORMERS_CACHE": "/opt/ml/model/",
"HF_HOME": "/opt/ml/model/",
"HUGGINGFACE_HUB_CACHE": "/opt/ml/model/",
}
Expand Down
1 change: 0 additions & 1 deletion src/sagemaker/serve/model_server/tgi/server.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
MODE_DIR_BINDING = "/opt/ml/model/"
_SHM_SIZE = "2G"
_DEFAULT_ENV_VARS = {
"TRANSFORMERS_CACHE": "/opt/ml/model/",
"HF_HOME": "/opt/ml/model/",
"HUGGINGFACE_HUB_CACHE": "/opt/ml/model/",
}
Expand Down
19 changes: 19 additions & 0 deletions tests/unit/sagemaker/serve/builder/test_djl_builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
"HF_MODEL_ID": "TheBloke/Llama-2-7b-chat-fp16",
"TENSOR_PARALLEL_DEGREE": "1",
"OPTION_DTYPE": "bf16",
"MODEL_LOADING_TIMEOUT": "1800",
}

mock_schema_builder = MagicMock()
Expand All @@ -63,8 +64,13 @@ class TestDjlBuilder(unittest.TestCase):
)
@patch("sagemaker.serve.builder.djl_builder._get_ram_usage_mb", return_value=1024)
@patch("sagemaker.serve.builder.djl_builder._get_nb_instance", return_value="ml.g5.24xlarge")
@patch(
"sagemaker.serve.builder.djl_builder._get_default_djl_configurations",
return_value=(mock_default_configs, 128),
)
def test_build_deploy_for_djl_local_container(
self,
mock_default_djl_config,
mock_get_nb_instance,
mock_get_ram_usage_mb,
mock_is_jumpstart_model,
Expand Down Expand Up @@ -125,8 +131,13 @@ def test_build_deploy_for_djl_local_container(
"sagemaker.serve.builder.djl_builder._concurrent_benchmark",
side_effect=[(0.03, 16), (0.10, 4), (0.15, 2)],
)
@patch(
"sagemaker.serve.builder.djl_builder._get_default_djl_configurations",
return_value=(mock_default_configs, 128),
)
def test_tune_for_djl_local_container(
self,
mock_default_djl_config,
mock_concurrent_benchmarks,
mock_serial_benchmarks,
mock_admissible_tensor_parallel_degrees,
Expand Down Expand Up @@ -165,8 +176,10 @@ def test_tune_for_djl_local_container(
"sagemaker.serve.builder.djl_builder._get_admissible_tensor_parallel_degrees",
return_value=[4],
)
@patch("sagemaker.serve.model_server.djl_serving.utils._get_available_gpus", return_value=None)
def test_tune_for_djl_local_container_deep_ping_ex(
self,
mock_get_available_gpus,
mock_get_admissible_tensor_parallel_degrees,
mock_serial_benchmarks,
mock_get_nb_instance,
Expand Down Expand Up @@ -204,8 +217,10 @@ def test_tune_for_djl_local_container_deep_ping_ex(
"sagemaker.serve.builder.djl_builder._get_admissible_tensor_parallel_degrees",
return_value=[4],
)
@patch("sagemaker.serve.model_server.djl_serving.utils._get_available_gpus", return_value=None)
def test_tune_for_djl_local_container_load_ex(
self,
mock_get_available_gpus,
mock_get_admissible_tensor_parallel_degrees,
mock_serial_benchmarks,
mock_get_nb_instance,
Expand Down Expand Up @@ -245,8 +260,10 @@ def test_tune_for_djl_local_container_load_ex(
"sagemaker.serve.builder.djl_builder._get_admissible_tensor_parallel_degrees",
return_value=[4],
)
@patch("sagemaker.serve.model_server.djl_serving.utils._get_available_gpus", return_value=None)
def test_tune_for_djl_local_container_oom_ex(
self,
mock_get_available_gpus,
mock_get_admissible_tensor_parallel_degrees,
mock_serial_benchmarks,
mock_get_nb_instance,
Expand Down Expand Up @@ -283,8 +300,10 @@ def test_tune_for_djl_local_container_oom_ex(
"sagemaker.serve.builder.djl_builder._get_admissible_tensor_parallel_degrees",
return_value=[4],
)
@patch("sagemaker.serve.model_server.djl_serving.utils._get_available_gpus", return_value=None)
def test_tune_for_djl_local_container_invoke_ex(
self,
mock_get_available_gpus,
mock_get_admissible_tensor_parallel_degrees,
mock_serial_benchmarks,
mock_get_nb_instance,
Expand Down
Loading