Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feature: Support maintaining an environment variable containing the image version #479

Merged
merged 1 commit into from
Sep 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,7 @@ def _build_local_images(

for image_generator_config in _image_generator_configs[target_version.major]:
config = _get_config_for_image(target_version_dir, image_generator_config, force)
config["build_args"]["IMAGE_VERSION"] = config["image_tag_generator"].format(image_version=str(target_version))
try:
image, log_gen = _docker_client.images.build(
path=target_version_dir, rm=True, pull=True, buildargs=config["build_args"]
Expand Down
2 changes: 2 additions & 0 deletions template/v0/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ FROM mambaorg/micromamba:$TAG_FOR_BASE_MICROMAMBA_IMAGE
ARG CUDA_MAJOR_MINOR_VERSION=''
ARG ENV_IN_FILENAME
ARG ARG_BASED_ENV_IN_FILENAME
ARG IMAGE_VERSION


ARG NB_USER="sagemaker-user"
Expand All @@ -16,6 +17,7 @@ ARG FIPS_VALIDATED_SSL=3.0.8
ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
ENV STUDIO_LOGGING_DIR="/var/log/studio/"
ENV EDITOR="nano"
ENV IMAGE_VERSION=$IMAGE_VERSION

USER root
RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
Expand Down
2 changes: 2 additions & 0 deletions template/v1/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ FROM mambaorg/micromamba:$TAG_FOR_BASE_MICROMAMBA_IMAGE
ARG CUDA_MAJOR_MINOR_VERSION=''
ARG ENV_IN_FILENAME
ARG ARG_BASED_ENV_IN_FILENAME
ARG IMAGE_VERSION

ARG AMZN_BASE="/opt/amazon/sagemaker"
ARG DB_ROOT_DIR="/opt/db"
Expand All @@ -19,6 +20,7 @@ ARG FIPS_VALIDATED_SSL=3.0.8
ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
ENV STUDIO_LOGGING_DIR="/var/log/studio/"
ENV EDITOR="nano"
ENV IMAGE_VERSION=$IMAGE_VERSION

USER root
RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
Expand Down
2 changes: 2 additions & 0 deletions template/v2/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ FROM mambaorg/micromamba:$TAG_FOR_BASE_MICROMAMBA_IMAGE
ARG CUDA_MAJOR_MINOR_VERSION=''
ARG ENV_IN_FILENAME
ARG ARG_BASED_ENV_IN_FILENAME
ARG IMAGE_VERSION

ARG AMZN_BASE="/opt/amazon/sagemaker"
ARG DB_ROOT_DIR="/opt/db"
Expand All @@ -19,6 +20,7 @@ ARG FIPS_VALIDATED_SSL=3.0.8
ENV SAGEMAKER_LOGGING_DIR="/var/log/sagemaker/"
ENV STUDIO_LOGGING_DIR="/var/log/studio/"
ENV EDITOR="nano"
ENV IMAGE_VERSION=$IMAGE_VERSION

USER root
RUN usermod "--login=${NB_USER}" "--home=/home/${NB_USER}" --move-home "-u ${NB_UID}" "${MAMBA_USER}" && \
Expand Down
10 changes: 10 additions & 0 deletions test/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -505,6 +505,16 @@ def mock_get_dir_for_version(base_version):
actual_output.add(f.read())
expected_output = {"container_logs1", "container_logs2"}
assert actual_output == expected_output
assert mock_docker_from_env.images.build.call_count == 2
call_args_list = mock_docker_from_env.images.build.call_args_list

# Assert the value for the first invocation
first_call_args = call_args_list[0].kwargs
assert first_call_args["buildargs"]["IMAGE_VERSION"] == "1.124.5-gpu"

# Assert the value for the second invocation
second_call_args = call_args_list[1].kwargs
assert second_call_args["buildargs"]["IMAGE_VERSION"] == "1.124.5-cpu"


@patch("os.path.exists")
Expand Down