From 5846583b2c671bb98dd6cbb00a8c65ecd13771b8 Mon Sep 17 00:00:00 2001
From: ChaonengQuan <46586569+ChaonengQuan@users.noreply.github.com>
Date: Fri, 27 Dec 2024 13:40:40 -0800
Subject: [PATCH] feat: Ensure micromamba meets minimum required version during
 build (#533)

---
 .gitignore             |  1 +
 template/v0/Dockerfile | 11 +++++++++++
 template/v1/Dockerfile | 11 +++++++++++
 template/v2/Dockerfile | 12 ++++++++++++
 template/v3/Dockerfile | 14 ++++++++++++++
 5 files changed, 49 insertions(+)

diff --git a/.gitignore b/.gitignore
index 120ee635..88a2a622 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
 __pycache__
 .idea
 .DS_Store
+.vscode
\ No newline at end of file
diff --git a/template/v0/Dockerfile b/template/v0/Dockerfile
index 8c89eb81..614268d4 100644
--- a/template/v0/Dockerfile
+++ b/template/v0/Dockerfile
@@ -18,8 +18,19 @@ ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
 ENV STUDIO_LOGGING_DIR="/var/log/studio/"
 ENV EDITOR="nano"
 ENV IMAGE_VERSION=$IMAGE_VERSION
+ENV PINNED_MICROMAMBA_MINOR_VERSION="1.5.*"
 
 USER root
+# Upgrade micromamba to the latest patch version in the pinned minor version range, if applicable
+RUN CURRENT_MICROMAMBA_VERSION=$(micromamba --version) && \
+    echo "Current micromamba version: $CURRENT_MICROMAMBA_VERSION" && \
+    if [[ "$CURRENT_MICROMAMBA_VERSION" == $PINNED_MICROMAMBA_MINOR_VERSION ]]; then \
+        echo "Upgrading micromamba to the latest $PINNED_MICROMAMBA_MINOR_VERSION version..." && \
+        micromamba self-update -c conda-forge --version "$PINNED_MICROMAMBA_MINOR_VERSION"; \
+    else \
+        echo "Micromamba is already at version $CURRENT_MICROMAMBA_VERSION (outside $PINNED_MICROMAMBA_MINOR_VERSION). No upgrade performed."; \
+    fi
+
 RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
     groupmod "--new-name=${NB_USER}" --non-unique "-g ${NB_GID}" "${MAMBA_USER}" && \
     # Update the expected value of MAMBA_USER for the
diff --git a/template/v1/Dockerfile b/template/v1/Dockerfile
index 2178e544..86ec6bee 100644
--- a/template/v1/Dockerfile
+++ b/template/v1/Dockerfile
@@ -21,8 +21,19 @@ ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
 ENV STUDIO_LOGGING_DIR="/var/log/studio/"
 ENV EDITOR="nano"
 ENV IMAGE_VERSION=$IMAGE_VERSION
+ENV PINNED_MICROMAMBA_MINOR_VERSION="1.5.*"
 
 USER root
+# Upgrade micromamba to the latest patch version in the pinned minor version range, if applicable
+RUN CURRENT_MICROMAMBA_VERSION=$(micromamba --version) && \
+    echo "Current micromamba version: $CURRENT_MICROMAMBA_VERSION" && \
+    if [[ "$CURRENT_MICROMAMBA_VERSION" == $PINNED_MICROMAMBA_MINOR_VERSION ]]; then \
+        echo "Upgrading micromamba to the latest $PINNED_MICROMAMBA_MINOR_VERSION version..." && \
+        micromamba self-update -c conda-forge --version "$PINNED_MICROMAMBA_MINOR_VERSION"; \
+    else \
+        echo "Micromamba is already at version $CURRENT_MICROMAMBA_VERSION (outside $PINNED_MICROMAMBA_MINOR_VERSION). No upgrade performed."; \
+    fi
+    
 RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
     groupmod "--new-name=${NB_USER}" --non-unique "-g ${NB_GID}" "${MAMBA_USER}" && \
     # Update the expected value of MAMBA_USER for the
diff --git a/template/v2/Dockerfile b/template/v2/Dockerfile
index b596638a..38322750 100644
--- a/template/v2/Dockerfile
+++ b/template/v2/Dockerfile
@@ -17,13 +17,25 @@ ARG NB_GID=100
 
 # https://www.openssl.org/source/
 ARG FIPS_VALIDATED_SSL=3.0.8
+ARG MIN_REQUIRED_MICROMAMBA_VERSION=1.5.11
 
 ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
 ENV STUDIO_LOGGING_DIR="/var/log/studio/"
 ENV EDITOR="nano"
 ENV IMAGE_VERSION=$IMAGE_VERSION
+ENV PINNED_MICROMAMBA_MINOR_VERSION="1.5.*"
 
 USER root
+# Upgrade micromamba to the latest patch version in the pinned minor version range, if applicable
+RUN CURRENT_MICROMAMBA_VERSION=$(micromamba --version) && \
+    echo "Current micromamba version: $CURRENT_MICROMAMBA_VERSION" && \
+    if [[ "$CURRENT_MICROMAMBA_VERSION" == $PINNED_MICROMAMBA_MINOR_VERSION ]]; then \
+        echo "Upgrading micromamba to the latest $PINNED_MICROMAMBA_MINOR_VERSION version..." && \
+        micromamba self-update -c conda-forge --version "$PINNED_MICROMAMBA_MINOR_VERSION"; \
+    else \
+        echo "Micromamba is already at version $CURRENT_MICROMAMBA_VERSION (outside $PINNED_MICROMAMBA_MINOR_VERSION). No upgrade performed."; \
+    fi
+
 RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
     groupmod "--new-name=${NB_USER}" --non-unique "-g ${NB_GID}" "${MAMBA_USER}" && \
     # Update the expected value of MAMBA_USER for the
diff --git a/template/v3/Dockerfile b/template/v3/Dockerfile
index 26881dc8..4cfb9e30 100644
--- a/template/v3/Dockerfile
+++ b/template/v3/Dockerfile
@@ -15,12 +15,26 @@ ARG NB_GID=100
 
 # https://www.openssl.org/source/
 ARG FIPS_VALIDATED_SSL=3.0.8
+ARG MIN_REQUIRED_MICROMAMBA_VERSION=1.5.11
 
 ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
 ENV STUDIO_LOGGING_DIR="/var/log/studio/"
 ENV EDITOR="nano"
+ENV MIN_REQUIRED_MICROMAMBA_VERSION=${MIN_REQUIRED_MICROMAMBA_VERSION}
 
 USER root
+
+# Compare the current micromamba version with the minimum required version, and only upgrade if strictly lower to avoid downgrades.
+RUN CURRENT_MICROMAMBA_VERSION=$(micromamba --version) && \
+    LOWEST_MICROMAMBA_VERSION=$(printf '%s\n%s' "$MIN_REQUIRED_MICROMAMBA_VERSION" "$CURRENT_MICROMAMBA_VERSION" | sort -V | head -n 1) && \
+    if [ "$LOWEST_MICROMAMBA_VERSION" = "$CURRENT_MICROMAMBA_VERSION" ] && [ "$CURRENT_MICROMAMBA_VERSION" != "$MIN_REQUIRED_MICROMAMBA_VERSION" ]; then \
+        echo "Upgrading micromamba from $CURRENT_MICROMAMBA_VERSION to $MIN_REQUIRED_MICROMAMBA_VERSION..." && \
+        micromamba self-update -c conda-forge --version "$MIN_REQUIRED_MICROMAMBA_VERSION"; \
+    else \
+        echo "Micromamba is already $MIN_REQUIRED_MICROMAMBA_VERSION or higher (current: $CURRENT_MICROMAMBA_VERSION). No update needed."; \
+    fi
+
+
 RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
     groupmod "--new-name=${NB_USER}" --non-unique "-g ${NB_GID}" "${MAMBA_USER}" && \
     # Update the expected value of MAMBA_USER for the