chore(docker): support more cuda versions (#3035)

bentoml · Sep 23, 2022 · 9ac664f · 9ac664f
1 parent 1714e5b
commit 9ac664f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/bentoml/_internal/bento/docker/__init__.py b/bentoml/_internal/bento/docker/__init__.py
@@ -18,12 +18,16 @@
 # Python supported versions
 SUPPORTED_PYTHON_VERSIONS = ["3.7", "3.8", "3.9", "3.10"]
 # CUDA supported versions
-SUPPORTED_CUDA_VERSIONS = ["11.6.2"]
+SUPPORTED_CUDA_VERSIONS = ["11.6.2", "11.4.3", "11.2.2"]
 # Mapping from user provided version argument to the full version target to install
 ALLOWED_CUDA_VERSION_ARGS = {
     "11": "11.6.2",
     "11.6": "11.6.2",
     "11.6.2": "11.6.2",
+    "11.4": "11.4.3",
+    "11.4.3": "11.4.3",
+    "11.2": "11.2.2",
+    "11.2.2": "11.2.2",
 }
 
 # Supported supported_architectures

diff --git a/docs/source/concepts/bento.rst b/docs/source/concepts/bento.rst
@@ -726,9 +726,11 @@ support matrix for all distros:
 GPU support
 """""""""""
 
-The :code:`cuda_version` field specifies the target cuda version to install on the
-generated docker image. Currently, the only supported cuda version is :code:`"11.6.2"`.
-BentoML also installs additional packages required for the target cuda version.
+The :code:`cuda_version` field specifies the target cuda version to
+install on the generated docker image. Currently, the only supported
+cuda versions are :code:`"11.6.2"`, :code:`"11.4.3"` and
+:code:`"11.2.2"`.  BentoML also installs additional packages required
+for the target cuda version.
 
 .. code:: yaml