Setup dependencies

Author: samruds

doc/requirements.txt

@@ -4,3 +4,4 @@ docutils==0.15.2
 packaging==20.9
 jinja2==3.1.3
 schema==0.7.5
+accelerate>=0.24.1,<=0.27.0

requirements/extras/huggingface_requirements.txt

@@ -0,0 +1 @@
+accelerate>=0.24.1,<=0.27.0

setup.py

@@ -79,6 +79,7 @@ def read_requirements(filename):
     "feature-processor": read_requirements(
         "requirements/extras/feature-processor_requirements.txt"
     ),
+    "huggingface": read_requirements("requirements/extras/huggingface_requirements.txt"),
 }
 # Meta dependency groups
 extras["all"] = [item for group in extras.values() for item in group]

src/sagemaker/serve/builder/model_builder.py

@@ -40,7 +40,7 @@
 from sagemaker.serve.save_retrive.version_1_0_0.metadata.metadata import Metadata
 from sagemaker.serve.spec.inference_spec import InferenceSpec
 from sagemaker.serve.utils import task
-from sagemaker.serve.utils.estimate_parser import estimate_command_parser, gather_data
+from accelerate.commands.estimate import estimate_command_parser, gather_data
 from sagemaker.serve.utils.exceptions import TaskNotFoundException
 from sagemaker.serve.utils.predictors import _get_local_mode_predictor
 from sagemaker.serve.utils.hardware_detector import _get_gpu_info, _get_gpu_info_fallback