Skip to content

Commit

Permalink
Minor update to description
Browse files Browse the repository at this point in the history
  • Loading branch information
samruds committed Mar 15, 2024
1 parent 82dad36 commit a15065b
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 8 deletions.
12 changes: 6 additions & 6 deletions src/sagemaker/serve/builder/model_builder.py
Expand Up @@ -73,11 +73,11 @@
MEMORY_BUFFER_MULTIPLIER = 1.2 # 20% buffer
VERSION_DETECTION_ERROR = (
"Please install accelerate and transformers for HuggingFace (HF) model "
"size calculations pip install 'sagemaker[huggingface]'"
"size calculations e.g. pip install 'sagemaker[huggingface]'"
)


# pylint: disable=attribute-defined-outside-init
# pylint: disable=attribute-defined-outside-init, disable=E1101
@dataclass
class ModelBuilder(Triton, DJL, JumpStart, TGI, Transformers):
"""Class that builds a deployable model.
Expand Down Expand Up @@ -730,14 +730,14 @@ def _total_inference_model_size_mib(self):
to add up to an additional 20% to the given model size as found by EleutherAI.
"""
try:
import accelerate.commands.estimate.estimate_command_parser
import accelerate.commands.estimate.gather_data
import accelerate.commands.estimate.estimate_command_parser as estimate_parser
import accelerate.commands.estimate.gather_data as estimate_gather

Check warning on line 734 in src/sagemaker/serve/builder/model_builder.py

View check run for this annotation

Codecov / codecov/patch

src/sagemaker/serve/builder/model_builder.py#L732-L734

Added lines #L732 - L734 were not covered by tests

dtypes = self.env_vars.get("dtypes", "float32")
parser = accelerate.commands.estimate.estimate_command_parser.estimate_command_parser()
parser = estimate_parser()
args = parser.parse_args([self.model, "--dtypes", dtypes])

Check warning on line 738 in src/sagemaker/serve/builder/model_builder.py

View check run for this annotation

Codecov / codecov/patch

src/sagemaker/serve/builder/model_builder.py#L736-L738

Added lines #L736 - L738 were not covered by tests

output = accelerate.commands.estimate.gather_data.gather_data(
output = estimate_gather(

Check warning on line 740 in src/sagemaker/serve/builder/model_builder.py

View check run for this annotation

Codecov / codecov/patch

src/sagemaker/serve/builder/model_builder.py#L740

Added line #L740 was not covered by tests
args
) # "dtype", "Largest Layer", "Total Size Bytes", "Training using Adam"
except ImportError as e:
Expand Down
4 changes: 2 additions & 2 deletions tests/unit/sagemaker/serve/builder/test_model_builder.py
Expand Up @@ -1343,8 +1343,8 @@ def test_build_for_transformers_happy_case_with_valid_gpu_fallback(
self.assertEqual(model_builder._can_fit_on_single_gpu(), True)

@patch("sagemaker.serve.builder.model_builder.ModelBuilder._build_for_transformers", Mock())
@patch("sagemaker.serve.builder.model_builder.estimate_command_parser")
@patch("sagemaker.serve.builder.model_builder.gather_data")
@patch("sagemaker.serve.builder.model_builder.estimate_parser")
@patch("sagemaker.serve.builder.model_builder.estimate_gather")
@patch("sagemaker.image_uris.retrieve")
@patch("sagemaker.djl_inference.model.urllib")
@patch("sagemaker.djl_inference.model.json")
Expand Down

0 comments on commit a15065b

Please sign in to comment.