diff --git a/docker/3.0.5/base/Dockerfile.cpu b/docker/3.0.5/base/Dockerfile.cpu
new file mode 100644
index 00000000..a75b0d00
--- /dev/null
+++ b/docker/3.0.5/base/Dockerfile.cpu
@@ -0,0 +1,199 @@
+ARG UBUNTU_VERSION=20.04
+ARG CUDA_VERSION=11.6.1
+ARG IMAGE_DIGEST=c2d95c9c6ff77da41cf0f2f9e8c5088f5b4db20c16a7566b808762f05b9032ef
+
+# Build stage for SQLite compilation
+FROM ubuntu:${UBUNTU_VERSION} as sqlite-builder
+RUN apt-get update && apt-get install -y --no-install-recommends \
+        build-essential \
+        wget \
+        ca-certificates \
+        && \
+    cd /tmp && \
+    wget https://www.sqlite.org/2025/sqlite-autoconf-3500200.tar.gz && \
+    tar xzf sqlite-autoconf-3500200.tar.gz && \
+    cd sqlite-autoconf-3500200 && \
+    ./configure --prefix=/usr/local && \
+    make && \
+    make install && \
+    ldconfig && \
+    cd / && \
+    rm -rf /tmp/sqlite-autoconf-3500200 /tmp/sqlite-autoconf-3500200.tar.gz && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Main image
+FROM nvidia/cuda:${CUDA_VERSION}-base-ubuntu${UBUNTU_VERSION}@sha256:${IMAGE_DIGEST}
+
+ARG MINICONDA_VERSION=24.7.1
+ARG CONDA_CHECKSUM=684cda724bc37e3bbbb342e440fc4cac515c92e91a489eb4359feca35382894b
+ARG CONDA_PY_VERSION=310
+ARG CONDA_PKG_VERSION=24.7.1
+ARG PYTHON_VERSION=3.10
+ARG PYARROW_VERSION=17.0.0
+ARG MLIO_VERSION=0.9.0
+ARG XGBOOST_VERSION=3.0.5
+
+ENV DEBIAN_FRONTEND=noninteractive
+ENV LANG=C.UTF-8
+ENV LC_ALL=C.UTF-8
+
+# Python won’t try to write .pyc or .pyo files on the import of source modules
+# Force stdin, stdout and stderr to be totally unbuffered. Good for logging
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+ENV PYTHONIOENCODING='utf-8'
+
+RUN apt-key del 7fa2af80 && \
+        apt-get update && apt-get install -y --no-install-recommends wget && \
+        wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-keyring_1.0-1_all.deb && \
+        dpkg -i cuda-keyring_1.0-1_all.deb && \
+        apt-get update && \
+    apt-get -y upgrade && \
+    apt-get -y install --no-install-recommends \
+        build-essential \
+        curl \
+        git \
+        jq \
+        libatlas-base-dev \
+        expat \
+        nginx \
+        openjdk-8-jdk-headless \
+        unzip \
+        wget \
+        apparmor \
+        linux-libc-dev \
+        libxml2 \
+        libgstreamer1.0-0 \
+        linux-libc-dev \
+        && \
+    # MLIO build dependencies
+    # Official Ubuntu APT repositories do not contain an up-to-date version of CMake required to build MLIO.
+    # Kitware contains the latest version of CMake.
+    wget http://es.archive.ubuntu.com/ubuntu/pool/main/libf/libffi/libffi7_3.3-4_amd64.deb && \
+    dpkg -i libffi7_3.3-4_amd64.deb && \
+    apt-get -y install --no-install-recommends \
+        apt-transport-https \
+        ca-certificates \
+        gnupg \
+        software-properties-common \
+        && \
+    wget -O - https://apt.kitware.com/keys/kitware-archive-latest.asc 2>/dev/null | \
+        gpg --dearmor - | \
+        tee /usr/share/keyrings/kitware-archive-keyring.gpg >/dev/null && \
+    echo 'deb [signed-by=/usr/share/keyrings/kitware-archive-keyring.gpg] https://apt.kitware.com/ubuntu/ bionic main' | tee /etc/apt/sources.list.d/kitware.list >/dev/null && \
+    apt-get update && \
+    rm /usr/share/keyrings/kitware-archive-keyring.gpg && \
+    apt-get install -y --no-install-recommends \
+        autoconf \
+        automake \
+        build-essential \
+        cmake \
+        cmake-data \
+        doxygen \
+        kitware-archive-keyring \
+        libcurl4-openssl-dev \
+        libssl-dev \
+        libtool \
+        ninja-build \
+        python3-dev \
+        python3-distutils \
+        python3-pip \
+        zlib1g-dev \
+        libxml2 \
+        zstd \
+        libsqlite3-0 \
+        && \
+    python3 -m pip install --upgrade pip && \
+    python3 -m pip install --upgrade certifi && \
+    apt-get clean && \
+    # Node.js setup
+    mkdir -p /etc/apt/keyrings && \
+    curl -fsSL https://deb.nodesource.com/gpgkey/nodesource-repo.gpg.key | \
+        gpg --dearmor -o /etc/apt/keyrings/nodesource.gpg && \
+    echo "deb [signed-by=/etc/apt/keyrings/nodesource.gpg] https://deb.nodesource.com/node_20.x nodistro main" | \
+        tee /etc/apt/sources.list.d/nodesource.list && \
+    apt-get update && \
+    apt-get install -y nodejs && \
+    npm install -g npm@latest && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install conda
+RUN cd /tmp && \
+    curl -L --output /tmp/Miniconda3.sh https://repo.anaconda.com/miniconda/Miniconda3-py${CONDA_PY_VERSION}_${MINICONDA_VERSION}-0-Linux-x86_64.sh && \
+    echo "${CONDA_CHECKSUM} /tmp/Miniconda3.sh" | sha256sum -c - && \
+    bash /tmp/Miniconda3.sh -bfp /miniconda3 && \
+    rm /tmp/Miniconda3.sh
+
+ENV PATH=/miniconda3/bin:${PATH}
+
+# Install MLIO with Apache Arrow integration
+
+# We could install mlio-py from conda, but it comes  with extra support such as image reader that increases image size
+# which increases training time. We build from source to minimize the image size.
+RUN echo "conda ${CONDA_PKG_VERSION}" >> /miniconda3/conda-meta/pinned && \
+    # Conda configuration see https://conda.io/projects/conda/en/latest/configuration.html
+    conda config --system --set auto_update_conda false && \
+    conda config --system --set show_channel_urls true && \
+    echo "python ${PYTHON_VERSION}.*" >> /miniconda3/conda-meta/pinned && \
+    conda install -c conda-forge python=${PYTHON_VERSION} --solver classic && \
+    pip install requests==2.32.3 && \
+    conda install conda=${CONDA_PKG_VERSION} --solver classic  && \
+    conda update -y conda && \
+    conda install -c conda-forge pyarrow=${PYARROW_VERSION} --solver classic && \
+    cd /miniconda3/pkgs/libgrpc-*/info/test/examples/node && \
+    npm install minimist@latest protobufjs@latest && \
+    # Remove Node.js, npm, and their dependencies
+    apt-get purge -y nodejs npm && \
+    apt-get autoremove -y && \
+    # Final cleanup
+    rm -rf /etc/apt/sources.list.d/nodesource.list \
+           /etc/apt/keyrings/nodesource.gpg \
+           /etc/apt/sources.list.d/kitware.list && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/* && \
+    # Continue with the rest of the build process
+    cd /tmp && \
+    git clone --branch v${MLIO_VERSION} https://github.com/awslabs/ml-io.git mlio && \
+    cd mlio && \
+    sed -i 's/find_package(Arrow 14.0.1 REQUIRED/find_package(Arrow 17.0.0 REQUIRED/g' CMakeLists.txt && \
+    sed -i 's/pyarrow==14.0.1/pyarrow==17.0.0/g' src/mlio-py/setup.py && \
+    build-tools/build-dependency build/third-party all && \
+    mkdir -p build/release && \
+    cd build/release && \
+    cmake -GNinja -DCMAKE_BUILD_TYPE=RelWithDebInfo -DCMAKE_PREFIX_PATH="$(pwd)/../third-party" ../.. && \
+    cmake --build . && \
+    cmake --build . --target install && \
+    cmake -DMLIO_INCLUDE_PYTHON_EXTENSION=ON -DPYTHON_EXECUTABLE="/miniconda3/bin/python3" \
+        -DMLIO_INCLUDE_ARROW_INTEGRATION=ON ../.. && \
+    cmake --build . --target mlio-py && \
+    cmake --build . --target mlio-arrow && \
+    cd ../../src/mlio-py && \
+    python3 setup.py bdist_wheel && \
+    python3 -m pip install typing && \
+    python3 -m pip install --upgrade pip && \
+    python3 -m pip install dist/*.whl && \
+    cp -r /tmp/mlio/build/third-party/lib/libtbb* /usr/local/lib/ && \
+    ldconfig && \
+    rm -rf /tmp/mlio
+
+# Copy compiled SQLite from builder stage
+COPY --from=sqlite-builder /usr/local/bin/sqlite3 /usr/local/bin/sqlite3
+COPY --from=sqlite-builder /usr/local/lib/libsqlite3.* /usr/local/lib/
+COPY --from=sqlite-builder /usr/local/include/sqlite3*.h /usr/local/include/
+
+# Update library cache and ensure /usr/local/bin is in PATH
+RUN ldconfig && \
+    echo "/usr/local/lib" > /etc/ld.so.conf.d/sqlite3.conf && \
+    ldconfig
+
+ENV PATH="/usr/local/bin:${PATH}"
+
+RUN echo "sqlite3 " 
+# This command will check the version and print it to the build logs
+RUN sqlite3 --version
+
+RUN apt list --installed
+
+# Install latest version of XGBoost
+RUN python3 -m pip install --no-cache -I xgboost==${XGBOOST_VERSION} numpy==2.1.0 pyarrow==17.0.0 pandas==2.2.3 
diff --git a/docker/3.0.5/final/Dockerfile.cpu b/docker/3.0.5/final/Dockerfile.cpu
new file mode 100644
index 00000000..6a48f280
--- /dev/null
+++ b/docker/3.0.5/final/Dockerfile.cpu
@@ -0,0 +1,96 @@
+ARG SAGEMAKER_XGBOOST_VERSION=3.0-5
+ARG PYTHON_VERSION=3.10
+
+FROM xgboost-container-base:${SAGEMAKER_XGBOOST_VERSION}-cpu-py3
+
+ARG SAGEMAKER_XGBOOST_VERSION=3.0.5
+
+########################
+# Install dependencies #
+########################
+COPY requirements.txt /requirements.txt
+RUN python3 -m pip install -r /requirements.txt && rm /requirements.txt
+
+# Fix Python 3.10 compatibility for sagemaker-containers
+# RUN python3 -c "import sys; sys.path.insert(0, '/miniconda3/lib/python3.10/site-packages'); \
+#     import sagemaker_containers._mapping as m; \
+#     import collections.abc; \
+#     setattr(collections, 'Mapping', collections.abc.Mapping); \
+#     exec(open('/miniconda3/lib/python3.10/site-packages/sagemaker_containers/_mapping.py').read().replace('collections.Mapping', 'collections.abc.Mapping'))" || \
+#     sed -i 's/collections\.Mapping/collections.abc.Mapping/g' /miniconda3/lib/python3.10/site-packages/sagemaker_containers/_mapping.py
+
+RUN sed -i 's/collections\.Mapping/collections.abc.Mapping/g' /miniconda3/lib/python3.10/site-packages/sagemaker_containers/_mapping.py
+
+# Install smdebug from source
+RUN python3 -m pip install git+https://github.com/awslabs/sagemaker-debugger.git@1.0.29
+
+
+###########################
+# Copy wheel to container #
+###########################
+COPY dist/sagemaker_xgboost_container-2.0-py2.py3-none-any.whl /sagemaker_xgboost_container-1.0-py2.py3-none-any.whl
+RUN rm -rf /miniconda3/lib/python${PYTHON_VERSION}/site-packages/numpy-1.21.2.dist-info && \
+    python3 -m pip install --force-reinstall PyYAML==6.0.1 && \
+    python3 -m pip install --no-cache --no-deps /sagemaker_xgboost_container-1.0-py2.py3-none-any.whl && \
+    python3 -m pip uninstall -y typing && \
+    rm /sagemaker_xgboost_container-1.0-py2.py3-none-any.whl
+
+##############
+# DMLC PATCH #
+##############
+# TODO: remove after making contributions back to xgboost for tracker.py
+# COPY src/sagemaker_xgboost_container/dmlc_patch/tracker.py \
+#    /miniconda3/lib/python${PYTHON_VERSION}/site-packages/xgboost/dmlc-core/tracker/dmlc_tracker/tracker.py
+
+# # Include DMLC python code in PYTHONPATH to use RabitTracker
+# ENV PYTHONPATH=$PYTHONPATH:/miniconda3/lib/python${PYTHON_VERSION}/site-packages/xgboost/dmlc-core/tracker
+
+#######
+# MMS #
+#######
+# Create MMS user directory
+RUN useradd -m model-server
+RUN mkdir -p /home/model-server/tmp && chown -R model-server /home/model-server
+
+# Copy MMS configs
+COPY docker/${SAGEMAKER_XGBOOST_VERSION}/resources/mms/config.properties.tmp /home/model-server
+ENV XGBOOST_MMS_CONFIG=/home/model-server/config.properties
+
+# Copy execution parameters endpoint plugin for MMS
+RUN mkdir -p /tmp/plugins
+COPY docker/${SAGEMAKER_XGBOOST_VERSION}/resources/mms/endpoints-1.0.jar /tmp/plugins
+RUN chmod +x /tmp/plugins/endpoints-1.0.jar
+
+# Create directory for models
+RUN mkdir -p /opt/ml/models
+RUN chmod +rwx /opt/ml/models
+
+# Copy Dask configs
+RUN mkdir /etc/dask
+COPY docker/configs/dask_configs.yaml /etc/dask/
+
+# Required label for multi-model loading
+LABEL com.amazonaws.sagemaker.capabilities.multi-models=true
+
+#####################
+# Required ENV vars #
+#####################
+# Set SageMaker training environment variables
+ENV SM_INPUT /opt/ml/input
+ENV SM_INPUT_TRAINING_CONFIG_FILE $SM_INPUT/config/hyperparameters.json
+ENV SM_INPUT_DATA_CONFIG_FILE $SM_INPUT/config/inputdataconfig.json
+ENV SM_CHECKPOINT_CONFIG_FILE $SM_INPUT/config/checkpointconfig.json
+# See: https://github.com/dmlc/xgboost/issues/7982#issuecomment-1379390906 https://github.com/dmlc/xgboost/pull/8257
+ENV NCCL_SOCKET_IFNAME eth 
+                            
+
+# Set SageMaker serving environment variables
+ENV SM_MODEL_DIR /opt/ml/model
+
+# Set SageMaker entrypoints
+ENV SAGEMAKER_TRAINING_MODULE sagemaker_xgboost_container.training:main
+ENV SAGEMAKER_SERVING_MODULE sagemaker_xgboost_container.serving:main
+
+EXPOSE 8080
+ENV TEMP=/home/model-server/tmp
+LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
diff --git a/docker/3.0.5/resources/mms/ExecutionParameters.java b/docker/3.0.5/resources/mms/ExecutionParameters.java
new file mode 100644
index 00000000..65134a8b
--- /dev/null
+++ b/docker/3.0.5/resources/mms/ExecutionParameters.java
@@ -0,0 +1,98 @@
+package software.amazon.ai.mms.plugins.endpoint;
+
+import com.google.gson.GsonBuilder;
+import com.google.gson.annotations.SerializedName;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.util.Properties;
+import software.amazon.ai.mms.servingsdk.Context;
+import software.amazon.ai.mms.servingsdk.ModelServerEndpoint;
+import software.amazon.ai.mms.servingsdk.annotations.Endpoint;
+import software.amazon.ai.mms.servingsdk.annotations.helpers.EndpointTypes;
+import software.amazon.ai.mms.servingsdk.http.Request;
+import software.amazon.ai.mms.servingsdk.http.Response;
+
+/**
+The modified endpoint source code for the jar used in this container.
+You can create this endpoint by moving it by cloning the MMS repo:
+> git clone https://github.com/awslabs/mxnet-model-server.git
+
+Copy this file into plugins/endpoints/src/main/java/software/amazon/ai/mms/plugins/endpoints/
+and then from the plugins directory, run:
+
+> ./gradlew fJ
+
+Modify file in plugins/endpoint/resources/META-INF/services/* to specify this file location
+
+Then build the JAR:
+
+> ./gradlew build
+
+The jar should be available in plugins/endpoints/build/libs as endpoints-1.0.jar
+**/
+@Endpoint(
+        urlPattern = "execution-parameters",
+        endpointType = EndpointTypes.INFERENCE,
+        description = "Execution parameters endpoint")
+public class ExecutionParameters extends ModelServerEndpoint {
+
+    @Override
+    public void doGet(Request req, Response rsp, Context ctx) throws IOException {
+        Properties prop = ctx.getConfig();
+        // 6 * 1024 * 1024
+        int maxRequestSize = Integer.parseInt(prop.getProperty("max_request_size", "6291456"));
+        SagemakerXgboostResponse response = new SagemakerXgboostResponse();
+        response.setMaxConcurrentTransforms(Integer.parseInt(prop.getProperty("NUM_WORKERS", "1")));
+        response.setBatchStrategy("MULTI_RECORD");
+        response.setMaxPayloadInMB(maxRequestSize / (1024 * 1024));
+        rsp.getOutputStream()
+                .write(
+                        new GsonBuilder()
+                                .setPrettyPrinting()
+                                .create()
+                                .toJson(response)
+                                .getBytes(StandardCharsets.UTF_8));
+    }
+
+    /** Response for Model server endpoint */
+    public static class SagemakerXgboostResponse {
+        @SerializedName("MaxConcurrentTransforms")
+        private int maxConcurrentTransforms;
+
+        @SerializedName("BatchStrategy")
+        private String batchStrategy;
+
+        @SerializedName("MaxPayloadInMB")
+        private int maxPayloadInMB;
+
+        public SagemakerXgboostResponse() {
+            maxConcurrentTransforms = 4;
+            batchStrategy = "MULTI_RECORD";
+            maxPayloadInMB = 6;
+        }
+
+        public int getMaxConcurrentTransforms() {
+            return maxConcurrentTransforms;
+        }
+
+        public String getBatchStrategy() {
+            return batchStrategy;
+        }
+
+        public int getMaxPayloadInMB() {
+            return maxPayloadInMB;
+        }
+
+        public void setMaxConcurrentTransforms(int newMaxConcurrentTransforms) {
+            maxConcurrentTransforms = newMaxConcurrentTransforms;
+        }
+
+        public void setBatchStrategy(String newBatchStrategy) {
+            batchStrategy = newBatchStrategy;
+        }
+
+        public void setMaxPayloadInMB(int newMaxPayloadInMB) {
+            maxPayloadInMB = newMaxPayloadInMB;
+        }
+    }
+}
diff --git a/docker/3.0.5/resources/mms/config.properties.tmp b/docker/3.0.5/resources/mms/config.properties.tmp
new file mode 100644
index 00000000..0abfed93
--- /dev/null
+++ b/docker/3.0.5/resources/mms/config.properties.tmp
@@ -0,0 +1,11 @@
+model_store=$$SAGEMAKER_MMS_MODEL_STORE$$
+load_models=$$SAGEMAKER_MMS_LOAD_MODELS$$
+plugins_path=/tmp/plugins
+inference_address=http://0.0.0.0:$$SAGEMAKER_BIND_TO_PORT$$
+management_address=http://0.0.0.0:$$SAGEMAKER_BIND_TO_PORT$$
+default_workers_per_model=$$SAGEMAKER_NUM_MODEL_WORKERS$$
+max_request_size=$$SAGEMAKER_MAX_REQUEST_SIZE$$
+decode_input_request=false
+default_service_handler=$$SAGEMAKER_MMS_DEFAULT_HANDLER$$
+job_queue_size=$$SAGEMAKER_MODEL_JOB_QUEUE_SIZE$$
+preload_model=true
diff --git a/docker/3.0.5/resources/mms/endpoints-1.0.jar b/docker/3.0.5/resources/mms/endpoints-1.0.jar
new file mode 100644
index 00000000..b5f4416d
Binary files /dev/null and b/docker/3.0.5/resources/mms/endpoints-1.0.jar differ
diff --git a/pyproject.toml b/pyproject.toml
index 5d7bf33d..12f88473 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,2 +1,6 @@
 [tool.isort]
 profile = "black"
+
+[build-system]
+requires = ["setuptools>=61.0,<81"]
+build-backend = "setuptools.build_meta"
diff --git a/requirements.txt b/requirements.txt
index 15ae634b..960734ab 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,30 +3,30 @@ PyYAML==6.0.1
 Pillow==9.1.1
 boto3==1.17.52
 botocore==1.20.52
-cryptography==39.0.1
-dask==2022.11.1
-dask-cuda==22.12.0
+cryptography==45.0.5
+dask==2024.9.0
+dask-cuda==24.10.0
 gunicorn==23.0.0
 itsdangerous==2.0.1
-matplotlib==3.6.3
+matplotlib==3.9.2
 multi-model-server==1.1.2
-numpy==1.24.1
-pandas==1.4.4
+numpy==2.1.0
+pandas==2.2.3
 protobuf==3.20.1
-psutil==5.6.7  # sagemaker-containers requires psutil 5.6.7
+psutil==5.8.0  # sagemaker-containers requires psutil 5.6.7
 pynvml==11.4.1
-python-dateutil==2.8.1
+python-dateutil==2.8.2
 retrying==1.3.3
-requests==2.29.0
+requests==2.32.3
 sagemaker-containers==2.8.6.post2
 sagemaker-inference==1.5.5
-scikit-learn==1.0.2
-scipy==1.9.3
+scipy==1.15.0
+scikit-learn==1.5.2
 urllib3==1.26.5
-wheel==0.36.2
+wheel==0.45.1
 jinja2==2.11.3
 MarkupSafe==1.1.1
 Werkzeug==0.15.6
 certifi==2023.7.22
 gevent==23.9.1
-numba==0.58.1
\ No newline at end of file
+numba==0.61.0
diff --git a/src/sagemaker_algorithm_toolkit/hyperparameter_validation.py b/src/sagemaker_algorithm_toolkit/hyperparameter_validation.py
index ddc64cdf..75c7c60a 100644
--- a/src/sagemaker_algorithm_toolkit/hyperparameter_validation.py
+++ b/src/sagemaker_algorithm_toolkit/hyperparameter_validation.py
@@ -374,8 +374,8 @@ def _format_range_value(self, open_, closed, default):
         return str(open_ if open_ is not None else closed if closed is not None else default)
 
     def format_as_integer(self):
-        max_neg_signed_int = -(2 ** 31)
-        max_signed_int = 2 ** 31 - 1
+        max_neg_signed_int = -(2**31)
+        max_signed_int = 2**31 - 1
         return (
             self._format_range_value(self.min_open, self.min_closed, max_neg_signed_int),
             self._format_range_value(self.max_open, self.max_closed, max_signed_int),
diff --git a/src/sagemaker_xgboost_container/algorithm_mode/hyperparameter_validation.py b/src/sagemaker_xgboost_container/algorithm_mode/hyperparameter_validation.py
index e42c8e7b..02f7bde5 100644
--- a/src/sagemaker_xgboost_container/algorithm_mode/hyperparameter_validation.py
+++ b/src/sagemaker_xgboost_container/algorithm_mode/hyperparameter_validation.py
@@ -321,7 +321,7 @@ def interaction_constraints_validator(value, dependencies):
             required=False,
         ),
         hpv.IntegerHyperparameter(
-            name="seed", range=hpv.Interval(min_open=-(2 ** 31), max_open=2 ** 31 - 1), required=False
+            name="seed", range=hpv.Interval(min_open=-(2**31), max_open=2**31 - 1), required=False
         ),
         hpv.IntegerHyperparameter(name="num_parallel_tree", range=hpv.Interval(min_closed=1), required=False),
         hpv.CategoricalHyperparameter(name="save_model_on_termination", range=["true", "false"], required=False),
diff --git a/src/sagemaker_xgboost_container/algorithm_mode/serve.py b/src/sagemaker_xgboost_container/algorithm_mode/serve.py
index 877cb48c..1812bb3c 100644
--- a/src/sagemaker_xgboost_container/algorithm_mode/serve.py
+++ b/src/sagemaker_xgboost_container/algorithm_mode/serve.py
@@ -149,7 +149,7 @@ def execution_parameters():
         parameters = {
             "MaxConcurrentTransforms": number_of_workers(),
             "BatchStrategy": "MULTI_RECORD",
-            "MaxPayloadInMB": int(PARSED_MAX_CONTENT_LENGTH / (1024 ** 2)),
+            "MaxPayloadInMB": int(PARSED_MAX_CONTENT_LENGTH / (1024**2)),
         }
     except Exception as e:
         return flask.Response(
diff --git a/src/sagemaker_xgboost_container/algorithm_mode/serve_utils.py b/src/sagemaker_xgboost_container/algorithm_mode/serve_utils.py
index 45dd5d08..756a3b9b 100644
--- a/src/sagemaker_xgboost_container/algorithm_mode/serve_utils.py
+++ b/src/sagemaker_xgboost_container/algorithm_mode/serve_utils.py
@@ -221,11 +221,33 @@ def predict(model, model_format, dtest, input_content_type, objective=None):
         else:
             raise ValueError("Content type {} is not supported".format(content_type))
 
+    def _predict_with_compat(booster, dtest):
+        """Predict with compatibility for both old and new XGBoost versions."""
+        best_iteration = getattr(booster, "best_ntree_limit", 0)
+
+        # Handle MagicMock objects in tests
+        try:
+            best_iteration = int(best_iteration) if best_iteration is not None else 0
+        except (TypeError, ValueError):
+            best_iteration = 0
+
+        # Check XGBoost version to determine which API to use
+        import inspect
+
+        predict_signature = inspect.signature(booster.predict)
+
+        if "ntree_limit" in predict_signature.parameters:
+            # Old XGBoost API (< 2.0)
+            return booster.predict(dtest, ntree_limit=best_iteration, validate_features=False)
+        else:
+            # New XGBoost API (>= 2.0)
+            if best_iteration > 0:
+                return booster.predict(dtest, iteration_range=(0, best_iteration), validate_features=False)
+            else:
+                return booster.predict(dtest, validate_features=False)
+
     if isinstance(model, list):
-        ensemble = [
-            booster.predict(dtest, ntree_limit=getattr(booster, "best_ntree_limit", 0), validate_features=False)
-            for booster in model
-        ]
+        ensemble = [_predict_with_compat(booster, dtest) for booster in model]
 
         if objective in [MULTI_SOFTMAX, BINARY_HINGE]:
             logging.info(f"Vote ensemble prediction of {objective} with {len(model)} models")
@@ -234,7 +256,7 @@ def predict(model, model_format, dtest, input_content_type, objective=None):
             logging.info(f"Average ensemble prediction of {objective} with {len(model)} models")
             return np.mean(ensemble, axis=0)
     else:
-        return model.predict(dtest, ntree_limit=getattr(model, "best_ntree_limit", 0), validate_features=False)
+        return _predict_with_compat(model, dtest)
 
 
 def is_selectable_inference_output():
diff --git a/src/sagemaker_xgboost_container/algorithm_mode/train.py b/src/sagemaker_xgboost_container/algorithm_mode/train.py
index 07f05d85..30a8233d 100755
--- a/src/sagemaker_xgboost_container/algorithm_mode/train.py
+++ b/src/sagemaker_xgboost_container/algorithm_mode/train.py
@@ -267,6 +267,9 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
     :param model_dir: Directory where model will be saved
     :param is_master: True if single node training, or the current node is the master node in distributed training.
     """
+    logging.info(f"TRAIN_JOB_DEBUG: Received is_master={is_master}")
+    print(f"TRAIN_JOB_DEBUG: Received is_master={is_master}")
+
     # Parse arguments for train() API
     num_round = train_cfg.pop("num_round")
     # Parse arguments for intermediate model callback
@@ -321,7 +324,7 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
                 train_dmatrix,
                 num_boost_round=num_round - iteration,
                 evals=watchlist,
-                feval=configured_feval,
+                custom_metric=configured_feval,
                 callbacks=callbacks,
                 xgb_model=xgb_model,
                 verbose_eval=False,
@@ -386,7 +389,7 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
                     cv_train_dmatrix,
                     num_boost_round=num_round - iteration,
                     evals=watchlist,
-                    feval=configured_feval,
+                    custom_metric=configured_feval,
                     evals_result=evals_result,
                     callbacks=callbacks,
                     xgb_model=xgb_model,
@@ -417,7 +420,12 @@ def train_job(train_cfg, train_dmatrix, val_dmatrix, train_val_dmatrix, model_di
     if not os.path.exists(model_dir):
         os.makedirs(model_dir)
 
+    logging.info(f"FINAL_MODEL_DEBUG: is_master={is_master}, model_dir={model_dir}")
+    print(f"FINAL_MODEL_DEBUG: is_master={is_master}, model_dir={model_dir}")
+
     if is_master:
+        logging.info("FINAL_MODEL_SAVE: Saving final model as master")
+        print("FINAL_MODEL_SAVE: Saving final model as master")
         if type(bst) is not list:
             model_location = os.path.join(model_dir, MODEL_NAME)
             bst.save_model(model_location)
diff --git a/src/sagemaker_xgboost_container/callback.py b/src/sagemaker_xgboost_container/callback.py
index 50cde3b5..89d935bd 100644
--- a/src/sagemaker_xgboost_container/callback.py
+++ b/src/sagemaker_xgboost_container/callback.py
@@ -79,17 +79,23 @@ def get_callbacks(
 
     callbacks = []
     callbacks.append(xgb.callback.EvaluationMonitor())
-    if checkpoint_dir:
+
+    if checkpoint_dir and is_master:
         save_checkpoint = xgb.callback.TrainingCheckPoint(
-            directory=checkpoint_dir, iterations=iteration, name=checkpointing.CHECKPOINT_FILENAME
-         )
+            directory=checkpoint_dir, interval=iteration, name=checkpointing.CHECKPOINT_FILENAME
+        )
         callbacks.append(save_checkpoint)
 
-    if save_model_on_termination == "true":
+    logging.info(f"CALLBACK_SETUP_DEBUG: save_model_on_termination={save_model_on_termination}, is_master={is_master}")
+
+    if save_model_on_termination == "true" and is_master:
+        logging.info("CALLBACK_ADDING: Adding SaveIntermediateModelCallBack on master")
         model_name = f"{MODEL_NAME}-{fold}" if fold is not None else MODEL_NAME
         save_intermediate_model = checkpointing.SaveIntermediateModelCallBack(model_dir, model_name, is_master)
         callbacks.append(save_intermediate_model)
         add_sigterm_handler(model_dir, is_master)
+    else:
+        logging.info(f"CALLBACK_SKIPPING save_model_on_termination={save_model_on_termination}, is_master={is_master})")
 
     if early_stopping_data_name and early_stopping_metric and early_stopping_rounds:
         maximize = early_stopping_metric in XGB_MAXIMIZE_METRICS
@@ -98,7 +104,7 @@ def get_callbacks(
             data_name=early_stopping_data_name,
             metric_name=early_stopping_metric,
             maximize=maximize,
-            save_best=True,
+            save_best=is_master,
         )
         callbacks.append(early_stop)
 
diff --git a/src/sagemaker_xgboost_container/checkpointing.py b/src/sagemaker_xgboost_container/checkpointing.py
index a9f3a664..79fd755b 100644
--- a/src/sagemaker_xgboost_container/checkpointing.py
+++ b/src/sagemaker_xgboost_container/checkpointing.py
@@ -7,7 +7,8 @@
 
 import xgboost as xgb
 from typing import Optional
-from xgboost import rabit
+
+# from xgboost import rabit
 from xgboost.callback import EvaluationMonitor
 from xgboost.core import XGBoostError
 
@@ -54,10 +55,17 @@ def train(train_args, checkpoint_dir):
         logging.info("Resuming from iteration %s", start_iteration)
 
     callbacks = train_args.get("callbacks", [])
-    callbacks.append(print_checkpointed_evaluation(start_iteration=start_iteration,
-                                                   end_iteration=train_args["num_boost_round"]))
-    callbacks.append(save_checkpoint(checkpoint_dir, start_iteration=start_iteration, iteration=start_iteration,
-                                     end_iteration=train_args["num_boost_round"]))
+    callbacks.append(
+        print_checkpointed_evaluation(start_iteration=start_iteration, end_iteration=train_args["num_boost_round"])
+    )
+    callbacks.append(
+        save_checkpoint(
+            checkpoint_dir,
+            start_iteration=start_iteration,
+            iteration=start_iteration,
+            end_iteration=train_args["num_boost_round"],
+        )
+    )
 
     train_args["verbose_eval"] = False  # suppress xgboost's print_evaluation()
     train_args["xgb_model"] = xgb_model
@@ -116,7 +124,7 @@ def after_iteration(self, model, epoch=0, evals_log=None):
                         score = log[-1]
                     msg += evaluation_monitor._fmt_metric(data, metric_name, score, stdv)
             msg += "\n"
-            rabit.tracker_print("[%d]\t%s\n" % (i + self.start_iteration, msg))
+            # rabit.tracker_print("[%d]\t%s\n" % (i + self.start_iteration, msg))
 
 
 def print_checkpointed_evaluation(end_iteration, iteration=0, rank=0, period=1, show_stdv=True, start_iteration=0):
@@ -164,16 +172,21 @@ def _sort_checkpoints(checkpoint_files):
     return checkpoint_files
 
 
-def save_checkpoint(checkpoint_dir, start_iteration=0, max_to_keep=5, num_round=None, rank=0, iteration=0,
-                    end_iteration=None):
+def save_checkpoint(
+    checkpoint_dir, start_iteration=0, max_to_keep=5, num_round=None, rank=0, iteration=0, end_iteration=None
+):
     """A callback function that saves checkpoints to disk.
 
     This is a wrapper function around SaveCheckpoint.
     For details, see SaveCheckpoint.
     """
     return SaveCheckpointCallBack(
-        checkpoint_dir=checkpoint_dir, start_iteration=start_iteration, max_to_keep=max_to_keep, num_round=num_round,
-        iteration=iteration, end_iteration=end_iteration
+        checkpoint_dir=checkpoint_dir,
+        start_iteration=start_iteration,
+        max_to_keep=max_to_keep,
+        num_round=num_round,
+        iteration=iteration,
+        end_iteration=end_iteration,
     )
 
 
@@ -220,12 +233,13 @@ class SaveCheckpointCallBack(xgb.callback.TrainingCallback):
         Example:
             >>> save_checkpoint = SaveCheckpoint("/opt/ml/checkpoints")
             >>> xgboost.train(prams, dtrain, callbacks=[save_checkpoint])
-        """
+    """
 
     SENTINEL = None
 
-    def __init__(self, checkpoint_dir, start_iteration=0, max_to_keep=5, num_round=None, rank=0, iteration=0,
-                 end_iteration=None):
+    def __init__(
+        self, checkpoint_dir, start_iteration=0, max_to_keep=5, num_round=None, rank=0, iteration=0, end_iteration=None
+    ):
         """Init SaveCheckpoint with checkpoint_dir"""
         self.checkpoint_dir = checkpoint_dir
         self.max_to_keep = max_to_keep
@@ -295,6 +309,7 @@ def start(self):
         When training is complete, we put SENTINEL on the queue, and when we
         see the SENTINEL, we clean up and exit the thread.
         """
+
         def _is_uploading(path):
             uploading = os.path.isfile(path + FILE_LOCK_SUFFIX)
             uploaded = os.path.isfile(path + FILE_SAFE_SUFFIX)
@@ -344,9 +359,7 @@ def _delete_uploaded_files_and_cleanup():
             _delete_uploaded_files()
             _cleanup()
 
-        self.thread = threading.Thread(
-            target=_delete_uploaded_files_and_cleanup,
-            daemon=True)
+        self.thread = threading.Thread(target=_delete_uploaded_files_and_cleanup, daemon=True)
         self.thread.start()
 
     def stop(self):
diff --git a/src/sagemaker_xgboost_container/data_utils.py b/src/sagemaker_xgboost_container/data_utils.py
index ae49a677..83dac24e 100644
--- a/src/sagemaker_xgboost_container/data_utils.py
+++ b/src/sagemaker_xgboost_container/data_utils.py
@@ -395,7 +395,7 @@ def get_libsvm_dmatrix(files_path, is_pipe=False):
         raise exc.UserError("Pipe mode not supported for LibSVM.")
 
     try:
-        dmatrix = xgb.DMatrix(files_path)
+        dmatrix = xgb.DMatrix(f"{files_path}?format=libsvm")
     except Exception as e:
         raise exc.UserError("Failed to load libsvm data with exception:\n{}".format(e))
 
@@ -531,7 +531,7 @@ def _get_pipe_mode_files_path(data_path: Union[List[str], str]) -> List[str]:
 
 
 def _make_symlinks_from_a_folder(dest_path: str, data_path: str, depth: int):
-    if (depth > MAX_FOLDER_DEPTH):
+    if depth > MAX_FOLDER_DEPTH:
         raise exc.UserError(f"Folder depth exceed the limit: {MAX_FOLDER_DEPTH}.")
 
     if os.path.isfile(data_path):
@@ -560,7 +560,7 @@ def _make_symlinks_from_a_folder_with_warning(dest_path: str, data_path: str):
     if (not os.path.exists(dest_path)) or (not os.path.exists(data_path)):
         raise exc.AlgorithmError(f"Unable to create symlinks as {data_path} or {dest_path} doesn't exist ")
 
-    if (not os.path.isdir(dest_path)):
+    if not os.path.isdir(dest_path):
         raise exc.AlgorithmError(f"Unable to create symlinks as dest_path {dest_path} is not a dir")
 
     try:
@@ -571,7 +571,7 @@ def _make_symlinks_from_a_folder_with_warning(dest_path: str, data_path: str):
                 f"The depth of folder {data_path} exceed the limit {MAX_FOLDER_DEPTH}."
                 f" Files in deeper sub dirs won't be loaded."
                 f" Please adjust the folder structure accordingly."
-                )
+            )
 
 
 def _get_file_mode_files_path(data_path: Union[List[str], str]) -> List[str]:
diff --git a/src/sagemaker_xgboost_container/distributed.py b/src/sagemaker_xgboost_container/distributed.py
index 91e2e241..a10880c9 100644
--- a/src/sagemaker_xgboost_container/distributed.py
+++ b/src/sagemaker_xgboost_container/distributed.py
@@ -18,13 +18,12 @@
 import logging
 import socket
 import sys
-import time
+import json
 
+from threading import Thread
 from retrying import retry
-from xgboost import rabit
-
-# This should point to xgb when the tracker is updated upstream
-from sagemaker_xgboost_container.dmlc_patch import tracker
+from xgboost.tracker import RabitTracker
+from xgboost import collective
 
 LOCAL_HOSTNAME = "127.0.0.1"
 
@@ -53,7 +52,7 @@ def rabit_run(
     first_port=None,
     second_port=None,
     max_connect_attempts=None,
-    connect_retry_timeout=3,
+    connect_retry_timeout=10,
     update_rabit_args=False,
 ):
     """Run execution function after initializing dmlc/rabit.
@@ -83,12 +82,14 @@ def rabit_run(
         port=first_port,
         max_connect_attempts=max_connect_attempts,
         connect_retry_timeout=connect_retry_timeout,
-    ) as rabit:
-        hosts_with_data = rabit.synchronize({"host": rabit.current_host, "include_in_training": include_in_training})
+    ) as rabit_ctx:
+        hosts_with_data = rabit_ctx.synchronize(
+            {"host": rabit_ctx.current_host, "include_in_training": include_in_training}
+        )
         hosts_with_data = [record["host"] for record in hosts_with_data if record["include_in_training"]]
 
         # Keep track of port used, so that hosts trying to shutdown know when server is not available
-        previous_port = rabit.master_port
+        previous_port = rabit_ctx.master_port
 
     if not include_in_training:
         logging.warning("Host {} not being used for distributed training.".format(current_host))
@@ -99,6 +100,8 @@ def rabit_run(
     if len(hosts_with_data) > 1:
         # Set up rabit with nodes that have data and an unused port so that previous slaves don't confuse it
         # with the previous rabit configuration
+        logging.info(f"SECOND_RABIT_DEBUG: hosts_with_data={hosts_with_data}, current_host={current_host}")
+
         with Rabit(
             hosts=hosts_with_data,
             current_host=current_host,
@@ -107,6 +110,12 @@ def rabit_run(
             connect_retry_timeout=connect_retry_timeout,
         ) as cluster:
             if update_rabit_args:
+                logging.info(
+                    f"RABIT_DEBUG: \
+                             cluster.is_master={cluster.is_master}, \
+                            current_host={current_host}"
+                )
+
                 args.update({"is_master": cluster.is_master})
             exec_fun(**args)
 
@@ -130,10 +139,23 @@ def __init__(self, is_master, current_host, master_port):
         :param current_host:
         :param master_port:
         """
-        self.is_master = is_master
-        self.rank = rabit.get_rank()
+        import time
+
+        self.is_master = is_master  # Store hostname-based master determination
         self.current_host = current_host
         self.master_port = master_port
+        self._id = int(time.time() * 1000000) % 1000000  # Unique ID for debugging
+        logging.info(
+            f"RABIT_HELPER_INIT: Created RabitHelper {self._id} with is_master={self.is_master} for host={current_host}"
+        )
+
+        try:
+            self.rank = collective.get_rank()
+            self.world_size = collective.get_world_size()
+        except Exception:
+            logging.error("collective init failed", exc_info=True)
+            self.rank = 0
+            self.world_size = 1
 
     def synchronize(self, data):
         """Synchronize data with the cluster.
@@ -144,15 +166,27 @@ def synchronize(self, data):
         :param data: data to send to the cluster
         :return: aggregated data from the all the nodes in the cluster
         """
+        # For single node or when collective is not initialized, just return the data
+        if self.world_size == 1:
+            return [data]
+
+        try:
+            collective.get_rank()  # Test if collective is initialized
+        except Exception:
+            logging.error("collective get_rank failed", exc_info=True)
+            return [data]
+
         results = []
-        for i in range(rabit.get_world_size()):
+        data_str = json.dumps(data)
+        for i in range(self.world_size):
             if self.rank == i:
-                logging.debug("Broadcasting data from self ({}) to others".format(self.rank))
-                rabit.broadcast(data, i)
+                logging.info("Broadcasting data from self ({}) to others".format(self.rank))
+                collective.broadcast(data_str, i)
                 results.append(data)
             else:
-                logging.debug("Receiving data from {}".format(i))
-                message = rabit.broadcast(None, i)
+                logging.info("Receiving data from {}".format(i))
+                message_str = collective.broadcast("", i)
+                message = json.loads(message_str) if message_str else None
                 results.append(message)
         return results
 
@@ -178,10 +212,6 @@ def __init__(
         :param connect_retry_timeout: Timeout value when attempting to connect to RabitTracker.
                             This will be ignored if max_connect_attempt is None
         """
-        # Get the host information. This is used to identify the master host
-        # that will run the RabitTracker and also to work out how many clients/slaves
-        # exist (this will ensure that all-reduce is set up correctly and that
-        # it blocks whilst waiting for those hosts to process the data).
         if not current_host:
             current_host = LOCAL_HOSTNAME
         self.current_host = current_host
@@ -192,7 +222,6 @@ def __init__(
         self.n_workers = len(self.hosts)
         self.logger.debug("Found hosts: {} [{}]".format(self.hosts, self.n_workers))
 
-        # We use the first lexicographically named host as the master if not indicated otherwise
         if not master_host:
             master_host = self.hosts[0]
         self.master_host = master_host
@@ -201,9 +230,6 @@ def __init__(
         self.logger.debug("Is Master: {}".format(self.is_master_host))
         self.logger.debug("Master: {}".format(self.master_host))
 
-        # We start the RabitTracker on a known port on the first host. We can
-        # do this since SageMaker Training instances are single tenent and we
-        # don't need to worry about port contention.
         if port is None:
             port = 9099
             self.logger.debug("No port specified using: {}".format(port))
@@ -218,116 +244,118 @@ def __init__(
         self.connect_retry_timeout = connect_retry_timeout
 
     def start(self):
-        """Start the rabit process.
+        """Start the collective process.
 
-        If current host is master host, initialize and start the Rabit Tracker in the background. All hosts then connect
-        to the master host to set up Rabit rank.
+        Initialize XGBoost collective for distributed training.
 
         :return: Initialized RabitHelper, which includes helpful information such as is_master and port
         """
-        self.rabit_context = None
-        if self.is_master_host:
-            self.logger.debug("Master host. Starting Rabit Tracker.")
-            # The Rabit Tracker is a Python script that is responsible for
-            # allowing each instance of rabit to find its peers and organize
-            # itself in to a ring for all-reduce. It supports primitive failure
-            # recovery modes.
-            #
-            # It runs on a master node that each of the individual Rabit instances
-            # talk to.
-            self.rabit_context = tracker.RabitTracker(
-                hostIP=self.current_host, nslave=self.n_workers, port=self.port, port_end=self.port + 1
+        self.logger.debug("Starting collective communication.")
+        self.tracker = None
+        self.tracker_thread = None
+
+        # For single node, skip collective initialization
+        if self.n_workers == 1:
+            self.logger.debug("Single worker detected, skipping collective init")
+            return RabitHelper(True, self.current_host, self.port)
+
+        try:
+            # Launch tracker on master only
+            if self.is_master_host:
+                self.tracker = RabitTracker(
+                    host_ip=str(_dns_lookup(self.master_host)),
+                    n_workers=int(self.n_workers),
+                    port=int(self.port),
+                    sortby="task",
+                )
+                self.tracker.start()
+                self.tracker_thread = Thread(target=self.tracker.wait_for)
+                self.tracker_thread.daemon = True
+                self.tracker_thread.start()
+                self.logger.info(f"RabitTracker worker_args: {self.tracker.worker_args()}")
+
+            self.logger.info(
+                f"MASTER_DEBUG_FIXED: Using hostname logic: \
+                    current_host={self.current_host}, \
+                        master_host={self.master_host}, \
+                            is_master={self.is_master_host}, \
+                                port={self.port}"
             )
 
-            # Useful logging to ensure that the tracker has started.
-            # These are the key-value config pairs that each of the rabit slaves
-            # should be initialized with. Since we have deterministically allocated
-            # the master host, its port, and the number of workers, we don't need
-            # to pass these out-of-band to each slave; but rely on the fact
-            # that each slave will calculate the exact same config as the server.
-            #
-            # TODO: should probably check that these match up what we pass below.
-            self.logger.info("Rabit slave environment: {}".format(self.rabit_context.slave_envs()))
-
-            # This actually starts the RabitTracker in a background/daemon thread
-            # that will automatically exit when the main process has finished.
-            self.rabit_context.start(self.n_workers)
-
-        # Start each parameter server that connects to the master.
-        self.logger.debug("Starting parameter server.")
-
-        # Rabit runs as an in-process singleton library that can be configured once.
-        # Calling this multiple times will cause a seg-fault (without calling finalize).
-        # We pass it the environment variables that match up with the RabitTracker
-        # so that this instance can discover its peers (and recover from failure).
-        #
-        # First we check that the RabitTracker is up and running. Rabit actually
-        # breaks (at least on Mac OS X) if the server is not running before it
-        # begins to try to connect (its internal retries fail because they reuse
-        # the same socket instead of creating a new one).
-        #
-        # if self.max_connect_attempts is None, this will loop indefinitely.
-        attempt = 0
-        successful_connection = False
-        while not successful_connection and (self.max_connect_attempts is None or attempt < self.max_connect_attempts):
-            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-                try:
-                    self.logger.debug("Checking if RabitTracker is available.")
-                    s.connect((self.master_host, self.port))
-                    successful_connection = True
-                    self.logger.debug("Successfully connected to RabitTracker.")
-                except OSError:
-                    self.logger.info("Failed to connect to RabitTracker on attempt {}".format(attempt))
-                    attempt += 1
-                    self.logger.info("Sleeping for {} sec before retrying".format(self.connect_retry_timeout))
-                    time.sleep(self.connect_retry_timeout)
-
-        if not successful_connection:
-            self.logger.error("Failed to connect to Rabit Tracker after %s attempts", self.max_connect_attempts)
-            raise Exception("Failed to connect to Rabit Tracker")
-        else:
-            self.logger.info("Connected to RabitTracker.")
-
-        rabit.init(
-            [
-                "DMLC_NUM_WORKER={}".format(self.n_workers).encode(),
-                "DMLC_TRACKER_URI={}".format(self.master_host).encode(),
-                "DMLC_TRACKER_PORT={}".format(self.port).encode(),
-            ]
-        )
+            import time
+
+            attempt = 0
+            successful_connection = False
+            while not successful_connection and (
+                self.max_connect_attempts is None or attempt < self.max_connect_attempts
+            ):
+                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                    try:
+                        self.logger.debug("Checking if RabitTracker is available.")
+                        s.connect((self.master_host, self.port))
+                        successful_connection = True
+                        self.logger.debug("Successfully connected to RabitTracker.")
+                    except OSError:
+                        self.logger.info("Failed to connect to RabitTracker on attempt {}".format(attempt))
+                        attempt += 1
+                        self.logger.info("Sleeping for {} sec before retrying".format(self.connect_retry_timeout))
+                        time.sleep(self.connect_retry_timeout)
+
+            if not successful_connection:
+                self.logger.error("Failed to connect to Rabit Tracker after %s attempts", self.max_connect_attempts)
+                raise Exception(f"Failed to connect to Rabit Tracker, current_host={self.current_host}")
+            else:
+                self.logger.info(f"Connected to RabitTracker, current_host={self.current_host}")
+
+            # Initialize collective for synchronization
+            collective.init(
+                dmlc_tracker_uri=str(_dns_lookup(self.master_host)),
+                dmlc_tracker_port=int(self.port),
+                dmlc_task_id=str(self.hosts.index(self.current_host)),
+                dmlc_retry=self.max_connect_attempts,
+                dmlc_timeout=self.connect_retry_timeout,
+            )
 
-        # We can check that the rabit instance has successfully connected to the
-        # server by getting the rank of the server (e.g. its position in the ring).
-        # This should be unique for each instance.
-        self.logger.debug("Rabit started - Rank {}".format(rabit.get_rank()))
-        self.logger.debug("Executing user code")
-
-        # We can now run user-code. Since XGBoost runs in the same process space
-        # it will use the same instance of rabit that we have configured. It has
-        # a number of checks throughout the learning process to see if it is running
-        # in distributed mode by calling rabit APIs. If it is it will do the
-        # synchronization automatically.
-        #
-        # Hence we can now execute any XGBoost specific training code and it
-        # will be distributed automatically.
+        except Exception as e:
+            self.logger.error(f"{self.current_host} collective init failed", exc_info=True)
+            self._cleanup_tracker()
+            raise e
+
+        self.logger.info(f"RABIT_START_DEBUG: Creating RabitHelper with is_master={self.is_master_host}")
         return RabitHelper(self.is_master_host, self.current_host, self.port)
 
     def stop(self):
-        """Shutdown parameter server.
-
-        If current host is master host, also join the background thread that is running the master host.
-        """
-        self.logger.debug("Shutting down parameter server.")
-
-        # This is the call that actually shuts down the rabit server; and when
-        # all of the slaves have been shut down then the RabitTracker will close
-        # /shutdown itself.
-        rabit.finalize()
-        if self.is_master_host:
-            self.rabit_context.join()
+        """Shutdown collective communication."""
+        self.logger.info(f"Shutting down collective, current_host={self.current_host}")
+
+        try:
+            collective.finalize()
+        except Exception:
+            self.logger.error(f"{self.current_host} collective finalize failed", exc_info=True)
+
+        # Wait for tracker thread to finish
+        if self.tracker_thread is not None:
+            try:
+                self.tracker_thread.join(timeout=1.0)
+            except Exception as e:
+                self.logger.debug("Tracker thread join failed: {}".format(e))
+            finally:
+                self.tracker_thread = None
+
+        self._cleanup_tracker()
+
+    def _cleanup_tracker(self):
+        """Clean up tracker safely."""
+        if self.tracker is not None:
+            try:
+                self.tracker.free()
+            except Exception as e:
+                self.logger.debug("Tracker cleanup failed: {}".format(e))
+            finally:
+                self.tracker = None
 
     def __enter__(self):
         return self.start()
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
-        return self.stop()
+        self.stop()
diff --git a/src/sagemaker_xgboost_container/distributed_gpu/distributed_gpu_training.py b/src/sagemaker_xgboost_container/distributed_gpu/distributed_gpu_training.py
index 99f95e3c..1abd3462 100644
--- a/src/sagemaker_xgboost_container/distributed_gpu/distributed_gpu_training.py
+++ b/src/sagemaker_xgboost_container/distributed_gpu/distributed_gpu_training.py
@@ -167,7 +167,7 @@ def run_training_with_dask(
                     dtrain=dtrain,
                     num_boost_round=num_round,
                     evals=watchlist,
-                    feval=configured_feval,
+                    custom_metric=configured_feval,
                     callbacks=callbacks,
                 )
                 booster = output["booster"]
diff --git a/src/sagemaker_xgboost_container/dmlc_patch/tracker.py b/src/sagemaker_xgboost_container/dmlc_patch/tracker.py
index c4b782ec..8c0130a8 100644
--- a/src/sagemaker_xgboost_container/dmlc_patch/tracker.py
+++ b/src/sagemaker_xgboost_container/dmlc_patch/tracker.py
@@ -17,6 +17,7 @@
 
 Tianqi Chen
 """
+
 # pylint: disable=invalid-name, missing-docstring, too-many-arguments, too-many-locals
 # pylint: disable=too-many-branches, too-many-statements
 from __future__ import absolute_import
diff --git a/src/sagemaker_xgboost_container/encoder.py b/src/sagemaker_xgboost_container/encoder.py
index cd11ee90..1b5dac0d 100644
--- a/src/sagemaker_xgboost_container/encoder.py
+++ b/src/sagemaker_xgboost_container/encoder.py
@@ -69,7 +69,7 @@ def libsvm_to_dmatrix(string_like):  # type: (bytes) -> xgb.DMatrix
             temp_file_location = libsvm_file.name
             libsvm_file.write(string_like)
 
-        dmatrix = xgb.DMatrix(temp_file_location)
+        dmatrix = xgb.DMatrix(f"{temp_file_location}?format=libsvm")
     finally:
         if temp_file_location and os.path.exists(temp_file_location):
             os.remove(temp_file_location)
diff --git a/src/sagemaker_xgboost_container/prediction_utils.py b/src/sagemaker_xgboost_container/prediction_utils.py
index 92c7c225..12b0c6c9 100644
--- a/src/sagemaker_xgboost_container/prediction_utils.py
+++ b/src/sagemaker_xgboost_container/prediction_utils.py
@@ -91,7 +91,11 @@ def _aggregate_predictions(self) -> np.ndarray:
         if self.classification:
             columns.append(self.y_prob.mean(axis=-1))
             # mode always returns same number of dimensions of output as for input
-            columns.append(stats.mode(self.y_pred, axis=1).mode[:, 0])
+            model_result = stats.mode(self.y_pred, axis=1, keepdims=True)
+            model_values = model_result.mode
+            if model_values.ndim > 1:
+                model_values = model_values[:, 0]
+            columns.append(model_values)
         else:
             columns.append(self.y_pred.mean(axis=-1))
 
diff --git a/src/sagemaker_xgboost_container/serving_mms.py b/src/sagemaker_xgboost_container/serving_mms.py
index c247aec8..70b8d20d 100644
--- a/src/sagemaker_xgboost_container/serving_mms.py
+++ b/src/sagemaker_xgboost_container/serving_mms.py
@@ -31,8 +31,8 @@
 USER_HANDLER_SERVICE = user_module_handler_service.__name__
 
 PORT = 8080
-DEFAULT_MAX_CONTENT_LEN = 6 * 1024 ** 2
-MAX_CONTENT_LEN_LIMIT = 20 * 1024 ** 2
+DEFAULT_MAX_CONTENT_LEN = 6 * 1024**2
+MAX_CONTENT_LEN_LIMIT = 20 * 1024**2
 MMS_NUM_MODEL_WORKERS_INIT = 1
 MMS_MODEL_JOB_QUEUE_SIZE_DEFAULT = 100
 
@@ -85,7 +85,7 @@ def _set_mms_configs(is_multi_model, handler):
     max_job_queue_size = 2 * max_workers
 
     # Max heap size = (max workers + max job queue size) * max payload size * 1.2 (20% buffer) + 128 (base amount)
-    max_heap_size = ceil((max_workers + max_job_queue_size) * (int(max_content_length) / 1024 ** 2) * 1.2) + 128
+    max_heap_size = ceil((max_workers + max_job_queue_size) * (int(max_content_length) / 1024**2) * 1.2) + 128
 
     os.environ["SAGEMAKER_MMS_MODEL_STORE"] = "/"
     os.environ["SAGEMAKER_MMS_LOAD_MODELS"] = ""
@@ -104,8 +104,10 @@ def _set_mms_configs(is_multi_model, handler):
     _set_default_if_not_exist("SAGEMAKER_MAX_DIRECT_MEMORY_SIZE", os.environ["SAGEMAKER_MAX_HEAP_SIZE"])
 
     disable_container_support_flag = ""
-    if "SAGEMAKER_DISABLE_CONTAINER_SUPPORT" in os.environ \
-            and os.environ["SAGEMAKER_DISABLE_CONTAINER_SUPPORT"] == "true":
+    if (
+        "SAGEMAKER_DISABLE_CONTAINER_SUPPORT" in os.environ
+        and os.environ["SAGEMAKER_DISABLE_CONTAINER_SUPPORT"] == "true"
+    ):
         disable_container_support_flag = " -XX:-UseContainerSupport"
 
     MMS_CONFIG_FILE_PATH = get_mms_config_file_path()
diff --git a/test/resources/boston/single_machine_customer_script.py b/test/resources/boston/single_machine_customer_script.py
index f323dcf5..4d11baf3 100644
--- a/test/resources/boston/single_machine_customer_script.py
+++ b/test/resources/boston/single_machine_customer_script.py
@@ -18,7 +18,7 @@
 import numpy as np
 import pandas as pd
 import xgboost as xgb
-from sklearn.datasets import load_boston
+from sklearn.datasets import fetch_california_housing
 from sklearn.metrics import mean_squared_error
 from sklearn.model_selection import train_test_split
 
@@ -37,11 +37,11 @@
 
     args = parser.parse_args()
 
-    # Load the Boston housing data into pandas data frame
-    boston = load_boston()
-    data = pd.DataFrame(boston.data)
-    data.columns = boston.feature_names
-    data["PRICE"] = boston.target
+    # Load the California housing data into pandas data frame (replacement for deprecated Boston dataset)
+    california = fetch_california_housing()
+    data = pd.DataFrame(california.data)
+    data.columns = california.feature_names
+    data["PRICE"] = california.target
 
     # Convert Pandas dataframe to XGBoost DMatrix for better performance (used later).
     X, y = data.iloc[:, :-1], data.iloc[:, -1]
@@ -74,10 +74,13 @@
     if not os.path.exists(args.output_data_dir):
         os.makedirs(args.output_data_dir)
 
-    ax = xgb.plot_importance(xg_reg)
-    fig = ax.figure
-    fig.set_size_inches(5, 5)
-    fig.savefig(os.path.join(args.output_data_dir, "feature-importance-plot.png"))
+    try:
+        ax = xgb.plot_importance(xg_reg)
+        fig = ax.figure
+        fig.set_size_inches(5, 5)
+        fig.savefig(os.path.join(args.output_data_dir, "feature-importance-plot.png"))
+    except Exception as e:
+        print(f"Warning: Could not create feature importance plot: {e}")
 
     # Finally, lets do a bit of cross-validation by using native XGB functionality (keeping some parameters constant, so
     # that we don't have a huge input list for this simple example.
diff --git a/test/resources/versions/train.py b/test/resources/versions/train.py
index 39cf17f6..f0fbb8e3 100644
--- a/test/resources/versions/train.py
+++ b/test/resources/versions/train.py
@@ -11,23 +11,23 @@
 boto3==1.17.52
 botocore==1.20.52
 conda==24.7.1
-cryptography==39.0.1
+cryptography==45.0.5
 gunicorn==23.0.0
-matplotlib==3.6.3
+matplotlib==3.9.2
 multi-model-server==1.1.2
-numpy==1.24.1
-pandas==1.4.4
-psutil==5.6.7
-pyarrow==14.0.1
-python-dateutil==2.8.1
+numpy==2.1.0
+pandas==2.2.3
+psutil==5.8.0
+pyarrow==17.0.0
+python-dateutil==2.8.2
 retrying==1.3.3
 sagemaker-containers==2.8.6.post2
 sagemaker-inference==1.5.5
-scikit-learn==1.0.2
-scipy==1.9.3
+scipy==1.15.0
+scikit-learn==1.5.2
 smdebug==1.0.29
 urllib3==1.26.5
-wheel==0.36.2
+wheel==0.45.1
 jinja2==2.11.3
 MarkupSafe==1.1.1
 Werkzeug==0.15.6
diff --git a/test/unit/algorithm_mode/test_serve.py b/test/unit/algorithm_mode/test_serve.py
index 6dd7b7f1..452777b5 100644
--- a/test/unit/algorithm_mode/test_serve.py
+++ b/test/unit/algorithm_mode/test_serve.py
@@ -28,7 +28,7 @@ def test_default_execution_parameters():
     assert parsed_exec_params_response["BatchStrategy"] == "MULTI_RECORD"
 
 
-@patch("sagemaker_xgboost_container.algorithm_mode.serve.PARSED_MAX_CONTENT_LENGTH", 19 * 1024 ** 2)
+@patch("sagemaker_xgboost_container.algorithm_mode.serve.PARSED_MAX_CONTENT_LENGTH", 19 * 1024**2)
 def test_max_execution_parameters():
     execution_parameters_response = serve.execution_parameters()
 
diff --git a/test/unit/algorithm_mode/test_serve_utils.py b/test/unit/algorithm_mode/test_serve_utils.py
index de54af48..ccb8a45b 100644
--- a/test/unit/algorithm_mode/test_serve_utils.py
+++ b/test/unit/algorithm_mode/test_serve_utils.py
@@ -164,8 +164,12 @@ def test_get_selected_content_keys_error():
     [
         (TEST_RAW_PREDICTIONS, TEST_KEYS_BINARY_LOG, serve_utils.BINARY_LOG, TEST_PREDICTIONS_BINARY_LOG),
         (TEST_RAW_PREDICTIONS_REG_LOG, TEST_KEYS_REG_LOG, serve_utils.REG_LOG, TEST_PREDICTIONS_REG_LOG),
-        (TEST_RAW_PREDICTIONS_REG_ABSOLUTEERR, TEST_KEYS_REG_ABSOLUTEERR, serve_utils.REG_ABSOLUTEERR,
-         TEST_PREDICTIONS_REG_ABSOLUTEERR),
+        (
+            TEST_RAW_PREDICTIONS_REG_ABSOLUTEERR,
+            TEST_KEYS_REG_ABSOLUTEERR,
+            serve_utils.REG_ABSOLUTEERR,
+            TEST_PREDICTIONS_REG_ABSOLUTEERR,
+        ),
     ],
 )
 def test_get_selected_predictions_all_keys(test_raw_predictions, selected_keys, objective, expected_predictions):
diff --git a/test/unit/distributed_gpu/test_dask_data_utils.py b/test/unit/distributed_gpu/test_dask_data_utils.py
index 571247fc..c74f3ada 100644
--- a/test/unit/distributed_gpu/test_dask_data_utils.py
+++ b/test/unit/distributed_gpu/test_dask_data_utils.py
@@ -40,7 +40,7 @@ def test_read_data_csv(self):
         x, y = read_data(self.data_path_csv, CSV)
         assert x.shape[0].compute() == self.NUM_ROWS_IN_EACH_FILE
         assert x.shape[1] == self.NUM_COLS_IN_EACH_FILE - 1
-        assert y.shape[0].compute() == self.NUM_ROWS_IN_EACH_FILE
+        assert len(y) == self.NUM_ROWS_IN_EACH_FILE
 
     def test_read_data_csv_malformed_path(self):
         x, y = read_data(self.data_path_csv + "/", CSV)
@@ -54,7 +54,7 @@ def test_read_data_parquet(self):
         x, y = read_data(self.data_path_parquet, PARQUET)
         assert x.shape[0].compute() == self.NUM_ROWS_IN_EACH_FILE * 2
         assert x.shape[1] == self.NUM_COLS_IN_EACH_FILE - 1
-        assert y.shape[0].compute() == self.NUM_ROWS_IN_EACH_FILE * 2
+        assert len(y) == self.NUM_ROWS_IN_EACH_FILE * 2
 
     def test_read_data_unsupported_content(self):
         with self.assertRaises(UserError):
diff --git a/test/unit/test_checkpointing.py b/test/unit/test_checkpointing.py
index aea12f64..2297b800 100644
--- a/test/unit/test_checkpointing.py
+++ b/test/unit/test_checkpointing.py
@@ -40,8 +40,9 @@ def test_SaveCheckpoint_single_iteration(self, model):
         iteration = 42
         end_iteration = 100
 
-        callback = SaveCheckpointCallBack(checkpoint_dir=self.test_dir, rank=rank, iteration=iteration,
-                                          end_iteration=end_iteration)
+        callback = SaveCheckpointCallBack(
+            checkpoint_dir=self.test_dir, rank=rank, iteration=iteration, end_iteration=end_iteration
+        )
 
         callback(model)
 
@@ -57,8 +58,9 @@ def test_SaveCheckpoint_multiple_from_scratch(self, model):
         rank = 0
         end_iteration = 100
 
-        callback = SaveCheckpointCallBack(checkpoint_dir=self.test_dir, max_to_keep=3, rank=rank,
-                                          end_iteration=end_iteration)
+        callback = SaveCheckpointCallBack(
+            checkpoint_dir=self.test_dir, max_to_keep=3, rank=rank, end_iteration=end_iteration
+        )
 
         for iteration in range(end_iteration):
             callback(model)
@@ -110,8 +112,9 @@ def test_SaveCheckpoint_uploading(self, model):
         rank = 0
         end_iteration = 100
 
-        callback = SaveCheckpointCallBack(checkpoint_dir=self.test_dir, max_to_keep=1, rank=rank,
-                                          end_iteration=end_iteration)
+        callback = SaveCheckpointCallBack(
+            checkpoint_dir=self.test_dir, max_to_keep=1, rank=rank, end_iteration=end_iteration
+        )
 
         # For iteration 0
         callback(model)
diff --git a/test/unit/test_distributed.py b/test/unit/test_distributed.py
index 5fca3836..a92ba44f 100644
--- a/test/unit/test_distributed.py
+++ b/test/unit/test_distributed.py
@@ -32,7 +32,7 @@ def synchronize_fn(host_count, port, master, idx, q):
 
 
 def rabit_run_fn(
-    host_count, is_run, first_port, second_port, master, idx, q, max_connect_attempts=None, connect_retry_timeout=3
+    host_count, is_run, first_port, second_port, master, idx, q, max_connect_attempts=None, connect_retry_timeout=60
 ):
     hosts = ["127.0.0.1"] + ["localhost" for _ in range(host_count - 1)]
     current_host = "127.0.0.1" if master else "localhost"
@@ -74,6 +74,7 @@ def test_integration_rabit_synchronize():
     q = Queue()
 
     port, _ = find_two_open_ports()
+    print(f"test_integration_rabit_synchronize, port={port}")
 
     host_count = 5
     host_list = range(host_count)
@@ -85,7 +86,7 @@ def test_integration_rabit_synchronize():
 
     num_responses = 0
     while num_responses < host_count:
-        host_aggregated_result = q.get(timeout=10)
+        host_aggregated_result = q.get(timeout=30)
         for host_individual_result in host_aggregated_result:
             assert host_individual_result in expected_results
         num_responses += 1
@@ -106,7 +107,7 @@ def test_rabit_run_all_hosts_run():
 
     num_responses = 0
     while num_responses < host_count:
-        response = q.get(timeout=15)
+        response = q.get(timeout=120)
         expected_results.remove(response)
         num_responses += 1
 
@@ -132,7 +133,7 @@ def test_rabit_run_exclude_one_host():
 
     num_responses = 0
     while num_responses < host_count - 1:
-        response = q.get(timeout=15)
+        response = q.get(timeout=300)
         expected_results.remove(response)
         num_responses += 1
 
@@ -150,13 +151,13 @@ def test_rabit_delay_master():
 
     for idx in host_list:
         p = Process(
-            target=rabit_run_delay_master, args=(host_count, True, first_port, second_port, idx == 0, idx, q, None)
+            target=rabit_run_delay_master, args=(host_count, True, first_port, second_port, idx == 0, idx, q, 3)
         )
         p.start()
 
     num_responses = 0
     while num_responses < host_count:
-        response = q.get(timeout=20)
+        response = q.get(timeout=300)
         expected_results.remove(response)
         num_responses += 1
 
@@ -181,6 +182,6 @@ def test_rabit_run_fail_bad_max_retry_attempts(bad_max_retry_attempts):
 
     num_responses = 0
     while num_responses < host_count:
-        host_result = q.get(timeout=10)
+        host_result = q.get(timeout=30)
         assert "max_connect_attempts must be None or an integer greater than 0." in host_result
         num_responses += 1
diff --git a/test/utils/local_mode.py b/test/utils/local_mode.py
index 914208f7..dfa8ff76 100644
--- a/test/utils/local_mode.py
+++ b/test/utils/local_mode.py
@@ -146,7 +146,7 @@ def train(
     entrypoint=None,
     source_dir=None,
     early_stopping=False,
-    train_time=30,
+    train_time=20,
 ):
     additional_env_vars = additional_env_vars or []
     additional_volumes = additional_volumes or []
@@ -426,7 +426,7 @@ def read_hyperparameters(customer_script, additonal_hyperparameters=None):
 
 def create_input_data_config(data_path, customer_script):
     channels = []
-    for (_, dirs, _) in os.walk(data_path):
+    for _, dirs, _ in os.walk(data_path):
         channels.extend(dirs)
         del dirs
 
diff --git a/tox.ini b/tox.ini
index 51a849e2..066ff0e6 100644
--- a/tox.ini
+++ b/tox.ini
@@ -15,12 +15,13 @@ deps =
     xgboost1.3: xgboost==1.3.3
     xgboost1.5: xgboost==1.5.2
     xgboost1.7: xgboost==1.7.4
+    xgboost3.0.5: xgboost==3.0.5
     xgboostlatest: xgboost
     -r{toxinidir}/requirements.txt
     -r{toxinidir}/test-requirements.txt
 conda_deps=
-    pyarrow==14.0.1
-    tbb==2020.2
+    pyarrow==17.0.0
+    tbb==2022.2.0
     mlio-py==0.9.0
 conda_channels=
     conda-forge