fix: update cuda base image and base image docs (#5492)

* fix: update cuda base image and base image docs

Signed-off-by: Frost Ming <[email protected]>

* fix: replace xla_bridge import with backend for consistency

Signed-off-by: Frost Ming <[email protected]>

* fix: update expected output values for cancer model predictions

Signed-off-by: Frost Ming <[email protected]>

---------

Signed-off-by: Frost Ming <[email protected]>

Author: frostming

src/bentoml/_internal/bento/build_config.py

@@ -67,7 +67,9 @@ def _convert_python_version(py_version: str | None) -> str | None:
     target_python_version = f"{major}.{minor}"
     if target_python_version != py_version:
         logger.warning(
-            "BentoML will install the latest 'python%s' instead of the specified 'python%s'. To use the exact python version, use a custom docker base image. See https://docs.bentoml.com/en/latest/concepts/bento.html#custom-base-image-advanced",
+            "BentoML will install the latest 'python%s' instead of the specified 'python%s'. "
+            "To use the exact python version, use a custom docker base image. "
+            "See https://docs.bentoml.com/en/latest/reference/bentoml/bento-build-options.html#docker-options-table",
             target_python_version,
             py_version,
         )

src/bentoml/_internal/container/frontend/dockerfile/__init__.py

@@ -98,7 +98,7 @@
             "supported_architectures": SUPPORTED_ARCHITECTURES,
         },
         "cuda": {
-            "image": "nvidia/cuda:{spec_version}-cudnn8-runtime-ubuntu20.04",
+            "image": "nvidia/cuda:{spec_version}-cudnn-runtime-ubuntu24.04",
             "supported_architectures": ["amd64", "arm64"],
         },
         "miniconda": {

src/bentoml/_internal/frameworks/flax.py

@@ -28,14 +28,14 @@
     from flax import struct
     from flax.core import FrozenDict
     from jax._src.lib.xla_bridge import XlaBackend
-    from jax.lib import xla_bridge
+    from jax.extend import backend
 
     from ...types import ModelSignature
     from .. import external_typing as ext
     from ..models.model import ModelSignaturesType
     from ..tag import Tag
 else:
-    xla_bridge = LazyLoader("xla_bridge", globals(), "jax.lib.xla_bridge")
+    backend = LazyLoader("backend", globals(), "jax.extend.backend")
 
 try:
     from flax import linen as nn
@@ -267,7 +267,7 @@ class FlaxRunnable(bentoml.legacy.Runnable):
 
         def __init__(self):
             super().__init__()
-            self.device = xla_bridge.get_backend().platform
+            self.device = backend.get_backend().platform
 
             self.model, self.state_dict = load_model(bento_model, device=self.device)
             self.params = self.state_dict["params"]

tests/integration/frameworks/models/xgboost.py

@@ -63,11 +63,11 @@ def check_output(out: ext.NpNDArray):
                 "predict": [
                     Input(
                         input_args=[np.array([cancer_data[0]])],
-                        expected=close_to([[0.9181239, 0.0818761]]),
+                        expected=close_to([[0.9075233, 0.09247671]]),
                     ),
                     Input(
                         input_args=[np.array([cancer_data[1]])],
-                        expected=close_to([[0.97593766, 0.02406229]]),
+                        expected=close_to([[[0.9710317, 0.02896825]]]),
                     ),
                 ],
             },
@@ -78,11 +78,11 @@ def check_output(out: ext.NpNDArray):
                 "predict": [
                     Input(
                         input_args=[pd.DataFrame([cancer_data[0]])],
-                        expected=close_to([[0.9181239, 0.0818761]]),
+                        expected=close_to([[0.9075233, 0.09247671]]),
                     ),
                     Input(
                         input_args=[pd.DataFrame([cancer_data[1]])],
-                        expected=close_to([[0.97593766, 0.02406229]]),
+                        expected=close_to([[[0.9710317, 0.02896825]]]),
                     ),
                 ],
             },