Merge master into master-assets (#4515)

* fix: make sure gpus are found in local_gpu run (#4384) * fix: make sure gpus are found in local_gpu run * fix: black formatting * fix: adjust unit test * feat: pin dll version to support python3.11 to the sdk (#4472) Co-authored-by: Ashwin Krishna <ashwikri@amazon.com> * fix: Skip No Canvas regions for test_deploy_best_candidate (#4477) * prepare release v2.211.0 * update development version to v2.211.1.dev0 * change: Enhance model builder selection logic to include model size (#4429) * change: Enhance model builder selection logic to include model size * Fix conflicts * Address PR comments * fix formatting * fix formatting of test * Fix token in tasks.json * Increase coverage for tests * fix formatting * Fix requirements * Import code instead of importing accelerate * Fix formatting * Setup dependencies * change: Upgrade smp to version 2.2 (#4479) * upgrading smp to version 2.2 * fixing linting issue * fixing syntax error with multiline if statement * upgrading smp to version 2.2 * fixing linting issue * fixing syntax error with multiline if statement * fixing formatting --------- Co-authored-by: Andrew Tian <tinandr@amazon.com> * feat: Update SM Python SDK for PT 2.2.0 SM DLC (#4481) * update pt2.2 sm training dlc pysdk * update pt2.2 sm inference dlc pysdk and region list * fix: Create custom tarfile extractall util to fix backward compatibility issue (#4476) * fix: Create custom tarfile extractall util to fix backward compatibility issue * Address review comments * fix logger.error statements * prepare release v2.212.0 * update development version to v2.212.1.dev0 * change: Update tblib constraint (#4452) * fix: make unit tests compatible with pytest-xdist (#4486) * fix: make unit tests compatible with pytest-xdist * fix failing test * feature: Add overriding logic in ModelBuilder when task is provided (#4460) * feat: Add Optional task to Model * Revert "feat: Add Optional task to Model" This reverts commit fd3e86b. * Add override logic in ModelBuilder with task provided * Adjusted formatting * Add extra unit tests for invalid inputs * Address PR comments * Add more test inputs to integration test * Add model_metadata field to ModelBuilder * Update doc * Update doc * Adjust formatting --------- Co-authored-by: Samrudhi Sharma <samruds@amazon.com> Co-authored-by: Xiong Zeng <xionzeng@amazon.com> * feature: Accept user-defined env variables for the entry-point (#4175) * fix: Move sagemaker pysdk version check after bootstrap in remote job (#4487) * change: enable github actions for PRs (#4489) * change: enable github actions for PRs * Update codebuild-ci.yml * trigger on pull_request_target * add source-version-override * fix permission * feature: Add ModelDataSource and SourceUri support for model package and while registering (#4492) Co-authored-by: Erick Benitez-Ramos <141277478+benieric@users.noreply.github.com> * feat: support JumpStart proprietary models (#4467) * feat: add proprietary manifest/specs parsing add unittests for test_cache small refactoring address comments and more unittests fix linting and fix more tests fix: pylint feat: JumpStartModel class for prop models * remove unused imports and fix docstyle * fix: remove unused args * fix: remove unused args * fix: more unused vars * fix: slow tests * fix: unittests * added more tests to cover some lines * remove estimator warn check * chore: address comments re performance * fix: address comments * complete list experience and other fixes * fix: pylint * add doc utils and fix pylint * fix: docstyle * fix: doc * fix: default payloads * fix: doc and tags and enums * fix: jumpstart doc * rename to open_weights and fix filtering * update filter name * doc update * fix: black * rename to proprietary model and fix unittests * address comments * fix: docstyle and flake8 * address more comments and fix doc * put back doc utils for future refactoring * add prop model title in doc * doc update --------- Co-authored-by: liujiaor <128006184+liujiaorr@users.noreply.github.com> * chore: emit warning when no instance specific gated training env var is available, and raise exception when accept_eula flag is not supplied (#4485) * fix: raise exception when no instance specific gated training env var available * chore: raise client exception if accept_eula flag is not set for gated models * chore: address flake8 errors * chore: emit warning when instance type is chosen with no gated training artifacts * fix: sagemaker session region not being used (#4469) * fix: sagemaker session region not being used * chore: add unit tests * fix: remove all JUMPSTART_DEFAULT_REGION_NAME default arguments * chore: use get_region_fallback throughout * chore: remove unnecessary if statement * chore: remove unnecessary if statement (2) --------- Co-authored-by: Erick Benitez-Ramos <141277478+benieric@users.noreply.github.com> * fix: add PT 2.2 support for smdistributed, pytorchddp, and torch_distributed distributions (#4480) * Add support for smdistributed, pytorchddp, torch_distributed for PT 2.2 * formatting * formatting --------- Co-authored-by: liujiaor <128006184+liujiaorr@users.noreply.github.com> * change: split coverage out from testenv in tox.ini (#4495) Co-authored-by: Ashwin Krishna <ashwikri@amazon.com> * change: add ci-health checks (#4493) * feat: tgi optimum 0.0.19, 0.0.20 releases (#4496) * feature: Add support for Streaming Inference (#4497) * feature: Add support for Streaming Inference * fix: codestyle-docs-test * fix: codestyle-docs-test * Add AutoML -> AutoMLV2 mapper (#4500) Co-authored-by: liujiaor <128006184+liujiaorr@users.noreply.github.com> * Skip of tests which are long running and causing the ResourceLimitInUse exception (#4504) * Improvement of the tuner documentation (#4506) * prepare release v2.213.0 * update development version to v2.213.1.dev0 * fix:urge customers to install latest version (#4507) * fix: list jumpstart models with invalid version strings (#4511) * fix: list jumpstart models with invalid versions * docstyle * docstyle * pylint * add more test * fix * fix: skip failing pt test (#4512) * fix: skip failing pt test * black-format --------- Co-authored-by: gv <gverkes@users.noreply.github.com> Co-authored-by: akrishna1995 <38850354+akrishna1995@users.noreply.github.com> Co-authored-by: Ashwin Krishna <ashwikri@amazon.com> Co-authored-by: Kalyani Nikure <110067132+knikure@users.noreply.github.com> Co-authored-by: ci <ci> Co-authored-by: Samrudhi Sharma <154457034+samruds@users.noreply.github.com> Co-authored-by: adtian2 <55163384+adtian2@users.noreply.github.com> Co-authored-by: Andrew Tian <tinandr@amazon.com> Co-authored-by: Sirut Buasai <73297481+sirutBuasai@users.noreply.github.com> Co-authored-by: Danny Bushkanets <d.bushkanets@gmail.com> Co-authored-by: xiongz945 <54782408+xiongz945@users.noreply.github.com> Co-authored-by: Samrudhi Sharma <samruds@amazon.com> Co-authored-by: Xiong Zeng <xionzeng@amazon.com> Co-authored-by: martinRenou <martin.renou@gmail.com> Co-authored-by: qidewenwhen <32910701+qidewenwhen@users.noreply.github.com> Co-authored-by: mrudulmn <161017394+mrudulmn@users.noreply.github.com> Co-authored-by: Haotian An <33510317+Captainia@users.noreply.github.com> Co-authored-by: liujiaor <128006184+liujiaorr@users.noreply.github.com> Co-authored-by: evakravi <69981223+evakravi@users.noreply.github.com> Co-authored-by: ruhanprasad <52712386+ruhanprasad@users.noreply.github.com> Co-authored-by: Mufaddal Rohawala <89424143+mufaddal-rohawala@users.noreply.github.com> Co-authored-by: Anton Repushko <repushko.a@gmail.com>
aws · Mar 19, 2024 · 3cdefeb · 3cdefeb
1 parent 12bcf05
commit 3cdefeb
Show file tree

Hide file tree

Showing 141 changed files with 7,160 additions and 890 deletions.
diff --git a/.github/workflows/codebuild-ci-health.yml b/.github/workflows/codebuild-ci-health.yml
@@ -0,0 +1,97 @@
+name: CI Health
+on:
+  schedule:
+    - cron: "0 */3 * * *"
+  workflow_dispatch:
+
+permissions:
+    id-token: write # This is required for requesting the JWT
+
+jobs:
+  codestyle-doc-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Codestyle & Doc Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-health-codestyle-doc-tests
+  unit-tests:
+    runs-on: ubuntu-latest
+    strategy:
+        fail-fast: false
+        matrix:
+          python-version: ["py38", "py39", "py310"]
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Unit Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-health-unit-tests
+          env-vars-for-codebuild: |
+            PY_VERSION
+        env:
+          PY_VERSION: ${{ matrix.python-version }}
+  integ-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Integ Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        id: codebuild
+        with:
+          project-name: sagemaker-python-sdk-ci-health-integ-tests
+  slow-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Slow Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-health-slow-tests
+  localmode-tests:
+      runs-on: ubuntu-latest
+      steps:
+        - name: Configure AWS Credentials
+          uses: aws-actions/configure-aws-credentials@v4
+          with:
+            role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+            aws-region: us-west-2
+            role-duration-seconds: 10800
+        - name: Run Local Mode Tests
+          uses: aws-actions/aws-codebuild-run-build@v1
+          with:
+            project-name: sagemaker-python-sdk-ci-health-localmode-tests
+  notebook-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Notebook Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-health-notebook-tests
diff --git a/.github/workflows/codebuild-ci.yml b/.github/workflows/codebuild-ci.yml
@@ -0,0 +1,48 @@
+name: PR Checks
+on:
+    pull_request_target:
+
+concurrency:
+    group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.head_ref }}
+    cancel-in-progress: true
+
+permissions:
+    id-token: write # This is required for requesting the JWT
+
+jobs:
+  codestyle-doc-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Codestyle & Doc Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-codestyle-doc-tests
+          source-version-override: 'pr/${{ github.event.pull_request.number }}'
+  unit-tests:
+    runs-on: ubuntu-latest
+    strategy:
+        fail-fast: false
+        matrix:
+          python-version: ["py38", "py39", "py310"]
+    steps:
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v4
+        with:
+          role-to-assume: ${{ secrets.CI_AWS_ROLE_ARN }}
+          aws-region: us-west-2
+          role-duration-seconds: 10800
+      - name: Run Unit Tests
+        uses: aws-actions/aws-codebuild-run-build@v1
+        with:
+          project-name: sagemaker-python-sdk-ci-unit-tests
+          source-version-override: 'pr/${{ github.event.pull_request.number }}'
+          env-vars-for-codebuild: |
+            PY_VERSION
+        env:
+          PY_VERSION: ${{ matrix.python-version }}
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,66 @@
 # Changelog
 
+## v2.213.0 (2024-03-15)
+
+### Features
+
+ * Add support for Streaming Inference
+ * tgi optimum 0.0.19, 0.0.20 releases
+ * support JumpStart proprietary models
+ * Add ModelDataSource and SourceUri support for model package and while registering
+ * Accept user-defined env variables for the entry-point
+ * Add overriding logic in ModelBuilder when task is provided
+
+### Bug Fixes and Other Changes
+
+ * Improvement of the tuner documentation
+ * Skip of tests which are long running and causing the ResourceLimitInUse exception
+ * Add AutoML -> AutoMLV2 mapper
+ * add ci-health checks
+ * split coverage out from testenv in tox.ini
+ * add PT 2.2 support for smdistributed, pytorchddp, and torch_distributed distributions
+ * sagemaker session region not being used
+ * chore: emit warning when no instance specific gated training env var is available, and raise exception when accept_eula flag is not supplied
+ * enable github actions for PRs
+ * Move sagemaker pysdk version check after bootstrap in remote job
+ * make unit tests compatible with pytest-xdist
+ * Update tblib constraint
+
+## v2.212.0 (2024-03-06)
+
+### Features
+
+ * Update SM Python SDK for PT 2.2.0 SM DLC
+
+### Bug Fixes and Other Changes
+
+ * Create custom tarfile extractall util to fix backward compatibility issue
+ * Upgrade smp to version 2.2
+ * Enhance model builder selection logic to include model size
+
+## v2.211.0 (2024-03-05)
+
+### Features
+
+ * pin dll version to support python3.11 to the sdk
+ * instance specific jumpstart host requirements
+ * Add TensorFlow 2.14 image configs
+ * Add AutoMLV2 support
+ * Support selective pipeline execution between function step and regular step
+ * Add new Triton DLC URIs
+
+### Bug Fixes and Other Changes
+
+ * Skip No Canvas regions for test_deploy_best_candidate
+ * make sure gpus are found in local_gpu run
+ * Bump Apache Airflow version to 2.8.2
+ * properly close sagemaker config file after loading config
+ * remove enable_network_isolation from the python doc
+
+### Documentation Changes
+
+ * Add doc for new feature processor APIs and classes
+
 ## v2.210.0 (2024-02-28)
 
 ### Features

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -77,8 +77,8 @@ Before sending us a pull request, please ensure that:
 1. Install coverage using `pip install .[test]`
 1. cd into the sagemaker-python-sdk folder: `cd sagemaker-python-sdk` or `cd /environment/sagemaker-python-sdk`
 1. Run the following tox command and verify that all code checks and unit tests pass: `tox tests/unit`
-
-You can also run a single test with the following command: `tox -e py310 -- -s -vv <path_to_file><file_name>::<test_function_name>`
+1. You can also run a single test with the following command: `tox -e py310 -- -s -vv <path_to_file><file_name>::<test_function_name>`
+1. You can run coverage via runcvoerage env : `tox -e runcoverage -- tests/unit` or `tox -e py310 -- tests/unit --cov=sagemaker --cov-append --cov-report xml`
   * Note that the coverage test will fail if you only run a single test, so make sure to surround the command with `export IGNORE_COVERAGE=-` and `unset IGNORE_COVERAGE`
   * Example: `export IGNORE_COVERAGE=- ; tox -e py310 -- -s -vv tests/unit/test_estimator.py::test_sagemaker_model_s3_uri_invalid ; unset IGNORE_COVERAGE`
 

diff --git a/README.rst b/README.rst
@@ -22,6 +22,10 @@ SageMaker Python SDK
    :target: https://sagemaker.readthedocs.io/en/stable/
    :alt: Documentation Status
 
+.. image:: https://github.com/benieric/sagemaker-python-sdk/actions/workflows/codebuild-ci-health.yml/badge.svg
+    :target: https://github.com/benieric/sagemaker-python-sdk/actions/workflows/codebuild-ci-health.yml
+    :alt: CI Health
+
 SageMaker Python SDK is an open source library for training and deploying machine learning models on Amazon SageMaker.
 
 With the SDK, you can train and deploy models using popular deep learning frameworks **Apache MXNet** and **TensorFlow**.
@@ -63,11 +67,10 @@ Table of Contents
 Installing the SageMaker Python SDK
 -----------------------------------
 
-The SageMaker Python SDK is built to PyPI and can be installed with pip as follows:
-
+The SageMaker Python SDK is built to PyPI and the latest version of the SageMaker Python SDK can be installed with pip as follows
 ::
 
-    pip install sagemaker
+    pip install sagemaker==<Latest version from pyPI from https://pypi.org/project/sagemaker/>
 
 You can install from source by cloning this repository and running a pip install command in the root directory of the repository:
 

diff --git a/VERSION b/VERSION
@@ -1 +1 @@
-2.210.1.dev0
+2.213.1.dev0
diff --git a/doc/doc_utils/jumpstart_doc_utils.py b/doc/doc_utils/jumpstart_doc_utils.py
@@ -74,9 +74,12 @@ class Frameworks(str, Enum):
 
 JUMPSTART_REGION = "eu-west-2"
 SDK_MANIFEST_FILE = "models_manifest.json"
+PROPRIETARY_SDK_MANIFEST_FILE = "proprietary-sdk-manifest.json"
 JUMPSTART_BUCKET_BASE_URL = "https://jumpstart-cache-prod-{}.s3.{}.amazonaws.com".format(
     JUMPSTART_REGION, JUMPSTART_REGION
 )
+PROPRIETARY_DOC_BUCKET = "https://jumpstart-cache-prod-us-west-2.s3.us-west-2.amazonaws.com"
+
 TASK_MAP = {
     Tasks.IC: ProblemTypes.IMAGE_CLASSIFICATION,
     Tasks.IC_EMBEDDING: ProblemTypes.IMAGE_EMBEDDING,
@@ -152,18 +155,26 @@ class Frameworks(str, Enum):
 }
 
 
-def get_jumpstart_sdk_manifest():
-    url = "{}/{}".format(JUMPSTART_BUCKET_BASE_URL, SDK_MANIFEST_FILE)
+def get_public_s3_json_object(url):
     with request.urlopen(url) as f:
         models_manifest = f.read().decode("utf-8")
     return json.loads(models_manifest)
 
 
-def get_jumpstart_sdk_spec(key):
-    url = "{}/{}".format(JUMPSTART_BUCKET_BASE_URL, key)
-    with request.urlopen(url) as f:
-        model_spec = f.read().decode("utf-8")
-    return json.loads(model_spec)
+def get_jumpstart_sdk_manifest():
+    return get_public_s3_json_object(f"{JUMPSTART_BUCKET_BASE_URL}/{SDK_MANIFEST_FILE}")
+
+
+def get_proprietary_sdk_manifest():
+    return get_public_s3_json_object(f"{PROPRIETARY_DOC_BUCKET}/{PROPRIETARY_SDK_MANIFEST_FILE}")
+
+
+def get_jumpstart_sdk_spec(s3_key: str):
+    return get_public_s3_json_object(f"{JUMPSTART_BUCKET_BASE_URL}/{s3_key}")
+
+
+def get_proprietary_sdk_spec(s3_key: str):
+    return get_public_s3_json_object(f"{PROPRIETARY_DOC_BUCKET}/{s3_key}")
 
 
 def get_model_task(id):
@@ -196,6 +207,45 @@ def get_model_source(url):
         return "Source"
 
 
+def create_proprietary_model_table():
+    proprietary_content_intro = []
+    proprietary_content_intro.append("\n")
+    proprietary_content_intro.append(".. list-table:: Available Proprietary Models\n")
+    proprietary_content_intro.append("   :widths: 50 20 20 20 20\n")
+    proprietary_content_intro.append("   :header-rows: 1\n")
+    proprietary_content_intro.append("   :class: datatable\n")
+    proprietary_content_intro.append("\n")
+    proprietary_content_intro.append("   * - Model ID\n")
+    proprietary_content_intro.append("     - Fine Tunable?\n")
+    proprietary_content_intro.append("     - Supported Version\n")
+    proprietary_content_intro.append("     - Min SDK Version\n")
+    proprietary_content_intro.append("     - Source\n")
+
+    sdk_manifest = get_proprietary_sdk_manifest()
+    sdk_manifest_top_versions_for_models = {}
+
+    for model in sdk_manifest:
+        if model["model_id"] not in sdk_manifest_top_versions_for_models:
+            sdk_manifest_top_versions_for_models[model["model_id"]] = model
+        else:
+            if str(sdk_manifest_top_versions_for_models[model["model_id"]]["version"]) < str(
+                model["version"]
+            ):
+                sdk_manifest_top_versions_for_models[model["model_id"]] = model
+
+    proprietary_content_entries = []
+    for model in sdk_manifest_top_versions_for_models.values():
+        model_spec = get_proprietary_sdk_spec(model["spec_key"])
+        proprietary_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
+        proprietary_content_entries.append("     - {}\n".format(False))  # TODO: support training
+        proprietary_content_entries.append("     - {}\n".format(model["version"]))
+        proprietary_content_entries.append("     - {}\n".format(model["min_version"]))
+        proprietary_content_entries.append(
+            "     - `{} <{}>`__ |external-link|\n".format("Source", model_spec.get("url"))
+        )
+    return proprietary_content_intro + proprietary_content_entries + ["\n"]
+
+
 def create_jumpstart_model_table():
     sdk_manifest = get_jumpstart_sdk_manifest()
     sdk_manifest_top_versions_for_models = {}
@@ -249,19 +299,19 @@ def create_jumpstart_model_table():
     file_content_intro.append("     - Source\n")
 
     dynamic_table_files = []
-    file_content_entries = []
+    open_weight_content_entries = []
 
     for model in sdk_manifest_top_versions_for_models.values():
         model_spec = get_jumpstart_sdk_spec(model["spec_key"])
         model_task = get_model_task(model_spec["model_id"])
         string_model_task = get_string_model_task(model_spec["model_id"])
         model_source = get_model_source(model_spec["url"])
-        file_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
-        file_content_entries.append("     - {}\n".format(model_spec["training_supported"]))
-        file_content_entries.append("     - {}\n".format(model["version"]))
-        file_content_entries.append("     - {}\n".format(model["min_version"]))
-        file_content_entries.append("     - {}\n".format(model_task))
-        file_content_entries.append(
+        open_weight_content_entries.append("   * - {}\n".format(model_spec["model_id"]))
+        open_weight_content_entries.append("     - {}\n".format(model_spec["training_supported"]))
+        open_weight_content_entries.append("     - {}\n".format(model["version"]))
+        open_weight_content_entries.append("     - {}\n".format(model["min_version"]))
+        open_weight_content_entries.append("     - {}\n".format(model_task))
+        open_weight_content_entries.append(
             "     - `{} <{}>`__ |external-link|\n".format(model_source, model_spec["url"])
         )
 
@@ -299,7 +349,10 @@ def create_jumpstart_model_table():
             f.writelines(file_content_single_entry)
             f.close()
 
+    proprietary_content_entries = create_proprietary_model_table()
+
     f = open("doc_utils/pretrainedmodels.rst", "a")
     f.writelines(file_content_intro)
-    f.writelines(file_content_entries)
+    f.writelines(open_weight_content_entries)
+    f.writelines(proprietary_content_entries)
     f.close()
diff --git a/doc/requirements.txt b/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
diff --git a/requirements/extras/huggingface_requirements.txt b/requirements/extras/huggingface_requirements.txt
@@ -0,0 +1 @@
+accelerate>=0.24.1,<=0.27.0
diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
@@ -39,3 +39,4 @@ tritonclient[http]<2.37.0
 onnx==1.14.1
 # tf2onnx==1.15.1
 nbformat>=5.9,<6
+accelerate>=0.24.1,<=0.27.0