From 200cd4f9dc7278014631c7ad9472ea1d4994896b Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 15:06:03 +0200 Subject: [PATCH 01/10] feat(build): Push toolkit bundle as OCI zstd artifact --- .github/actions/push-oci-artifact/action.yml | 97 +++++++++++++++++++- .github/workflows/cd.yml | 27 ++++++ 2 files changed, 122 insertions(+), 2 deletions(-) diff --git a/.github/actions/push-oci-artifact/action.yml b/.github/actions/push-oci-artifact/action.yml index 676bf26..5f32e08 100644 --- a/.github/actions/push-oci-artifact/action.yml +++ b/.github/actions/push-oci-artifact/action.yml @@ -1,5 +1,5 @@ -name: 'Push OCI image volume artifact' -description: 'Push one or more files as an OCI image-volume-compatible artifact' +name: 'Push OCI artifact' +description: 'Push one or more files as either an image-volume-compatible OCI artifact or direct-file OCI artifact' inputs: artifact_ref: description: 'Fully qualified OCI image reference' @@ -16,6 +16,9 @@ inputs: annotations: description: 'Additional newline-separated OCI annotations in key=value format' default: '' + publish_mode: + description: 'How to publish files: image-volume wraps files in a gzipped rootfs layer, direct-files pushes files directly.' + default: 'image-volume' runs: using: "composite" @@ -24,7 +27,24 @@ runs: - name: Set up ORAS uses: oras-project/setup-oras@38de303aac69abb66f3e6255b7198bff35f323e3 # v2.0.0 + - name: Validate publish mode + shell: bash + env: + PUBLISH_MODE: ${{ inputs.publish_mode }} + run: | + set -euo pipefail + + case "${PUBLISH_MODE}" in + image-volume|direct-files) + ;; + *) + echo "Error: unsupported publish_mode: ${PUBLISH_MODE}" >&2 + exit 1 + ;; + esac + - name: Push OCI image volume artifact + if: ${{ inputs.publish_mode == 'image-volume' }} shell: bash env: ARTIFACT_REF: ${{ inputs.artifact_ref }} @@ -53,6 +73,7 @@ runs: manifest_annotations+=("${annotation}") done <<< "${EXTRA_ANNOTATIONS}" + file_count=0 while IFS= read -r file_descriptor; do [ -n "${file_descriptor}" ] || continue file_path="${file_descriptor%%:*}" @@ -62,8 +83,14 @@ runs: exit 1 fi cp "${file_path}" "${rootfs_dir}/$(basename "${file_path}")" + file_count=$((file_count + 1)) done <<< "${FILES}" + if [ "${file_count}" -eq 0 ]; then + echo "Error: no OCI artifact files were provided" >&2 + exit 1 + fi + layer_tar="${workdir}/rootfs.tar" layer_tar_gz="${workdir}/rootfs.tar.gz" tar -C "${rootfs_dir}" -cf "${layer_tar}" . @@ -101,3 +128,69 @@ runs: ) oras manifest fetch "${ARTIFACT_REF}" >/dev/null echo "Pushed ${ARTIFACT_REF}" + + - name: Push OCI direct file artifact + if: ${{ inputs.publish_mode == 'direct-files' }} + shell: bash + env: + ARTIFACT_REF: ${{ inputs.artifact_ref }} + ARTIFACT_TYPE: ${{ inputs.artifact_type }} + FILES: ${{ inputs.files }} + TOOLKIT_VERSION: ${{ inputs.toolkit_version }} + EXTRA_ANNOTATIONS: ${{ inputs.annotations }} + run: | + set -euo pipefail + + manifest_annotations=( + "org.opencontainers.image.source=https://github.com/${GITHUB_REPOSITORY}" + "org.opencontainers.image.revision=${GITHUB_SHA}" + "org.opencontainers.image.version=${TOOLKIT_VERSION}" + "com.deepnote.toolkit.artifact-type=${ARTIFACT_TYPE}" + ) + + while IFS= read -r annotation; do + [ -n "${annotation}" ] || continue + manifest_annotations+=("${annotation}") + done <<< "${EXTRA_ANNOTATIONS}" + + oras_args=( + push + "${ARTIFACT_REF}" + --artifact-type "${ARTIFACT_TYPE}" + ) + + for annotation in "${manifest_annotations[@]}"; do + oras_args+=(--annotation "${annotation}") + done + + file_count=0 + while IFS= read -r file_descriptor; do + [ -n "${file_descriptor}" ] || continue + file_path="${file_descriptor%%:*}" + media_type="${file_descriptor#*:}" + if [ "${file_path}" = "${media_type}" ]; then + echo "Error: invalid OCI artifact file descriptor: ${file_descriptor}" >&2 + exit 1 + fi + if [[ "${file_path}" = /* ]]; then + echo "Error: direct-files mode requires relative file paths, got ${file_path}" >&2 + exit 1 + fi + if [ ! -f "${file_path}" ]; then + echo "Error: OCI artifact file not found at ${file_path}" >&2 + ls -la "$(dirname "${file_path}")" >&2 || true + exit 1 + fi + + oras_args+=("${file_path}:${media_type}") + file_count=$((file_count + 1)) + done <<< "${FILES}" + + if [ "${file_count}" -eq 0 ]; then + echo "Error: no OCI artifact files were provided" >&2 + exit 1 + fi + + oras "${oras_args[@]}" + oras manifest fetch "${ARTIFACT_REF}" >/dev/null + echo "Pushed ${ARTIFACT_REF}" diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 59553cf..5ec92c2 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -90,6 +90,21 @@ jobs: set -euo pipefail ./bin/build + - name: Create zstd bundle artifact + env: + PYTHON_VERSION: ${{ matrix.python_version }} + run: | + set -euo pipefail + + if ! command -v zstd >/dev/null 2>&1; then + sudo apt-get update + sudo apt-get install --no-install-recommends -y zstd + fi + + zstd -T0 -3 -f \ + -o "dist/python${PYTHON_VERSION}.tar.zst" \ + "dist/python${PYTHON_VERSION}.tar" + # STAGING UPLOAD - name: Upload to staging uses: ./.github/actions/upload-bundle-artifacts @@ -128,6 +143,18 @@ jobs: toolkit_version: ${{ steps.version.outputs.VERSION }} annotations: com.deepnote.toolkit.python-version=${{ matrix.python_version }} + - name: Push toolkit bundle OCI zstd artifact + uses: ./.github/actions/push-oci-artifact + with: + artifact_ref: docker.io/deepnote/toolkit-bundle:${{ steps.version.outputs.VERSION }}-python${{ matrix.python_version }}-tar-zst + artifact_type: application/vnd.deepnote.toolkit.bundle.v1 + files: dist/python${{ matrix.python_version }}.tar.zst:application/vnd.deepnote.toolkit.python-bundle.v1.tar+zstd + toolkit_version: ${{ steps.version.outputs.VERSION }} + annotations: | + com.deepnote.toolkit.python-version=${{ matrix.python_version }} + com.deepnote.toolkit.bundle-format=tar-zst + publish_mode: direct-files + - name: Push toolkit installer OCI artifact if: matrix.python_version == '3.10' uses: ./.github/actions/push-oci-artifact From d553736ca80b71baef872f4f596e4ee45a450742 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 15:17:35 +0200 Subject: [PATCH 02/10] fix(build): Permissions from bind-mounted Docker build in the runner --- .github/workflows/cd.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 5ec92c2..fab23f3 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -89,6 +89,7 @@ jobs: run: | set -euo pipefail ./bin/build + sudo chown -R "$(id -u):$(id -g)" dist - name: Create zstd bundle artifact env: From 87bff23365d76c5dddf75aa729eebfec9e716711 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 16:20:59 +0200 Subject: [PATCH 03/10] fix(build): Align paths in the published artifacts --- .github/actions/push-oci-artifact/action.yml | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/.github/actions/push-oci-artifact/action.yml b/.github/actions/push-oci-artifact/action.yml index 5f32e08..49907ba 100644 --- a/.github/actions/push-oci-artifact/action.yml +++ b/.github/actions/push-oci-artifact/action.yml @@ -141,6 +141,9 @@ runs: run: | set -euo pipefail + workdir="$(mktemp -d)" + trap 'rm -rf "${workdir}"' EXIT + manifest_annotations=( "org.opencontainers.image.source=https://github.com/${GITHUB_REPOSITORY}" "org.opencontainers.image.revision=${GITHUB_SHA}" @@ -182,7 +185,14 @@ runs: exit 1 fi - oras_args+=("${file_path}:${media_type}") + file_name="$(basename "${file_path}")" + staged_path="${workdir}/${file_name}" + if [ -e "${staged_path}" ]; then + echo "Error: multiple OCI artifact files have the same basename: ${file_name}" >&2 + exit 1 + fi + cp "${file_path}" "${staged_path}" + oras_args+=("${file_name}:${media_type}") file_count=$((file_count + 1)) done <<< "${FILES}" @@ -191,6 +201,9 @@ runs: exit 1 fi - oras "${oras_args[@]}" + ( + cd "${workdir}" + oras "${oras_args[@]}" + ) oras manifest fetch "${ARTIFACT_REF}" >/dev/null echo "Pushed ${ARTIFACT_REF}" From 4029e7949ac56acad98e95e59bb0878f637499a5 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 16:21:20 +0200 Subject: [PATCH 04/10] test(build): Increase zstd compression ratio --- .github/workflows/cd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index fab23f3..75d8479 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -102,7 +102,7 @@ jobs: sudo apt-get install --no-install-recommends -y zstd fi - zstd -T0 -3 -f \ + zstd -T0 -6 -f \ -o "dist/python${PYTHON_VERSION}.tar.zst" \ "dist/python${PYTHON_VERSION}.tar" From ac9585045ad8b73e5b0b00e088a5f64b9848e9a7 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 16:40:40 +0200 Subject: [PATCH 05/10] test(build): Increase zstd compression ratio --- .github/workflows/cd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 75d8479..a82263a 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -102,7 +102,7 @@ jobs: sudo apt-get install --no-install-recommends -y zstd fi - zstd -T0 -6 -f \ + zstd -T0 -12 -f \ -o "dist/python${PYTHON_VERSION}.tar.zst" \ "dist/python${PYTHON_VERSION}.tar" From b9341f6d300427f2345b9e008421a5309bef11a6 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 17:04:18 +0200 Subject: [PATCH 06/10] test(build): Increase zstd compression ratio --- .github/workflows/cd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index a82263a..6d1f6a7 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -102,7 +102,7 @@ jobs: sudo apt-get install --no-install-recommends -y zstd fi - zstd -T0 -12 -f \ + zstd -T0 --ultra -20 -f \ -o "dist/python${PYTHON_VERSION}.tar.zst" \ "dist/python${PYTHON_VERSION}.tar" From fde08a4e994d8ef9613c1b344502a6e26d50c38b Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 17:33:25 +0200 Subject: [PATCH 07/10] chore(build): Use -6 compression ratio for zstd --- .github/workflows/cd.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 6d1f6a7..75d8479 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -102,7 +102,7 @@ jobs: sudo apt-get install --no-install-recommends -y zstd fi - zstd -T0 --ultra -20 -f \ + zstd -T0 -6 -f \ -o "dist/python${PYTHON_VERSION}.tar.zst" \ "dist/python${PYTHON_VERSION}.tar" From bbf93968765fb7612b04da7a1d8523cd8d1ec15a Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 17:33:50 +0200 Subject: [PATCH 08/10] chore(build): Switch to pigz to speed up the gzip compression --- .github/actions/push-oci-artifact/action.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/actions/push-oci-artifact/action.yml b/.github/actions/push-oci-artifact/action.yml index 49907ba..7d4d94a 100644 --- a/.github/actions/push-oci-artifact/action.yml +++ b/.github/actions/push-oci-artifact/action.yml @@ -94,7 +94,7 @@ runs: layer_tar="${workdir}/rootfs.tar" layer_tar_gz="${workdir}/rootfs.tar.gz" tar -C "${rootfs_dir}" -cf "${layer_tar}" . - gzip -n -c "${layer_tar}" > "${layer_tar_gz}" + pigz -n -6 -p "$(nproc)" -c "${layer_tar}" > "${layer_tar_gz}" layer_diff_id="$(sha256sum "${layer_tar}" | awk '{print $1}')" From 54d13bf018cf01ee2bc30c5a868770f0c5705940 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Fri, 3 Jul 2026 17:57:05 +0200 Subject: [PATCH 09/10] feat(downloader): Support pulling Toolkit bundle from OCI registry --- dockerfiles/cache-downloader/Dockerfile | 6 +- .../cache-downloader/cache-downloader.py | 155 +++++++++++++++--- tests/unit/test_cache_downloader.py | 82 +++++++++ 3 files changed, 220 insertions(+), 23 deletions(-) diff --git a/dockerfiles/cache-downloader/Dockerfile b/dockerfiles/cache-downloader/Dockerfile index d4685f0..31de1b9 100644 --- a/dockerfiles/cache-downloader/Dockerfile +++ b/dockerfiles/cache-downloader/Dockerfile @@ -1,6 +1,8 @@ FROM alpine:3.22 -RUN apk add --no-cache bash aws-cli python3 +RUN apk add --no-cache bash aws-cli ca-certificates python3 zstd + +COPY --from=docker.io/regclient/regctl:v0.11.5 /regctl /usr/local/bin/regctl COPY cache-downloader.py /usr/bin/cache-downloader RUN chmod 755 /usr/bin/cache-downloader @@ -15,4 +17,4 @@ RUN if [ -z "$PYTHON_VERSIONS" ]; then echo "PYTHON_VERSIONS argument not provid ENV PYTHONUNBUFFERED=1 -CMD [ "cache-downloader" ] \ No newline at end of file +CMD [ "cache-downloader" ] diff --git a/dockerfiles/cache-downloader/cache-downloader.py b/dockerfiles/cache-downloader/cache-downloader.py index 0e6fd6d..a74af16 100644 --- a/dockerfiles/cache-downloader/cache-downloader.py +++ b/dockerfiles/cache-downloader/cache-downloader.py @@ -9,29 +9,34 @@ from typing import List BASE_PATH = "/deepnote-toolkit/" +DEFAULT_DOWNLOAD_METHOD = "s3" +OCI_ZSTD_DOWNLOAD_METHOD = "oci-zstd" +DEFAULT_TOOLKIT_BUNDLE_OCI_REPOSITORY = "docker.io/deepnote/toolkit-bundle" -def download_dependency( - release_name: str, python_version: str, toolkit_index_bucket_name: str -): - """Download the dependencies for the given Python version and release name.""" +def build_toolkit_bundle_ref( + release_name: str, + python_version: str, + oci_repository: str = DEFAULT_TOOLKIT_BUNDLE_OCI_REPOSITORY, +) -> str: + return f"{oci_repository}:{release_name}-python{python_version}-tar-zst" - version_path = os.path.join(BASE_PATH, release_name, f"python{python_version}") - done_file = os.path.join(version_path, f"{python_version}-done") - if Path(done_file).is_file(): - print( - f"{datetime.datetime.now()}: {release_name} python{python_version} already cached, skipping download" - ) - return +def _read_stderr(process: subprocess.Popen) -> bytes: + if process.stderr is None: + return b"" + return process.stderr.read() - # Create the version directory if it doesn't exist - os.makedirs(version_path, exist_ok=True) +def download_dependency_from_s3( + release_name: str, + python_version: str, + toolkit_index_bucket_name: str, + version_path: str, +) -> None: s3_path = f"s3://{toolkit_index_bucket_name}/deepnote-toolkit/{release_name}/python{python_version}.tar" print(f"{datetime.datetime.now()}: Downloading {release_name} {s3_path}") - # Use Popen to stream the data aws_process = subprocess.Popen( ["aws", "s3", "cp", "--no-sign-request", s3_path, "-"], stdout=subprocess.PIPE, @@ -44,16 +49,13 @@ def download_dependency( stderr=subprocess.PIPE, ) aws_process.stdout.close() # Allow aws_process to receive a SIGPIPE if tar_process exits. - _, tar_process_stderr = ( - tar_process.communicate() - ) # Wait for tar_process to complete + _, tar_process_stderr = tar_process.communicate() aws_process_returncode = aws_process.wait() if aws_process_returncode != 0: raise Exception( - f"Error downloading {release_name} (aws s3 command failed): aws stderr: {aws_process.stderr.read()}, tar stderr: {tar_process_stderr}" + f"Error downloading {release_name} (aws s3 command failed): aws stderr: {_read_stderr(aws_process)}, tar stderr: {tar_process_stderr}" ) - # Check for errors if tar_process.returncode != 0: raise Exception( f"Error downloading {release_name} (tar command failed): {tar_process_stderr}" @@ -62,12 +64,111 @@ def download_dependency( print( f"{datetime.datetime.now()}: Done downloading {release_name} {s3_path} and extracting to {version_path}" ) + + +def download_dependency_from_oci_zstd( + release_name: str, + python_version: str, + oci_repository: str, + version_path: str, +) -> None: + artifact_file = f"python{python_version}.tar.zst" + artifact_ref = build_toolkit_bundle_ref( + release_name, python_version, oci_repository + ) + print(f"{datetime.datetime.now()}: Downloading {release_name} {artifact_ref}") + + regctl_process = subprocess.Popen( + ["regctl", "artifact", "get", "--file", artifact_file, artifact_ref], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + zstd_process = subprocess.Popen( + ["zstd", "-dc"], + stdin=regctl_process.stdout, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + regctl_process.stdout.close() + tar_process = subprocess.Popen( + ["tar", "-xf", "-", "-C", version_path], + stdin=zstd_process.stdout, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + zstd_process.stdout.close() + + _, tar_process_stderr = tar_process.communicate() + zstd_process_returncode = zstd_process.wait() + regctl_process_returncode = regctl_process.wait() + + if regctl_process_returncode != 0: + raise Exception( + f"Error downloading {release_name} (regctl command failed): regctl stderr: {_read_stderr(regctl_process)}, zstd stderr: {_read_stderr(zstd_process)}, tar stderr: {tar_process_stderr}" + ) + if zstd_process_returncode != 0: + raise Exception( + f"Error downloading {release_name} (zstd command failed): zstd stderr: {_read_stderr(zstd_process)}, tar stderr: {tar_process_stderr}" + ) + if tar_process.returncode != 0: + raise Exception( + f"Error downloading {release_name} (tar command failed): {tar_process_stderr}" + ) + + print( + f"{datetime.datetime.now()}: Done downloading {release_name} {artifact_ref} and extracting to {version_path}" + ) + + +def download_dependency( + release_name: str, + python_version: str, + toolkit_index_bucket_name: str, + download_method: str = DEFAULT_DOWNLOAD_METHOD, + oci_repository: str = DEFAULT_TOOLKIT_BUNDLE_OCI_REPOSITORY, +): + """Download the dependencies for the given Python version and release name.""" + + version_path = os.path.join(BASE_PATH, release_name, f"python{python_version}") + done_file = os.path.join(version_path, f"{python_version}-done") + + if Path(done_file).is_file(): + print( + f"{datetime.datetime.now()}: {release_name} python{python_version} already cached, skipping download" + ) + return + + os.makedirs(version_path, exist_ok=True) + + if download_method == DEFAULT_DOWNLOAD_METHOD: + download_dependency_from_s3( + release_name, + python_version, + toolkit_index_bucket_name, + version_path, + ) + elif download_method == OCI_ZSTD_DOWNLOAD_METHOD: + download_dependency_from_oci_zstd( + release_name, + python_version, + oci_repository, + version_path, + ) + else: + raise ValueError( + f"Unsupported TOOLKIT_DOWNLOAD_METHOD={download_method!r}; expected one of: {DEFAULT_DOWNLOAD_METHOD}, {OCI_ZSTD_DOWNLOAD_METHOD}" + ) + # Create the "done" file Path(done_file).touch() def submit_downloading( - python_versions: List[str], release_name: str, toolkit_index_bucket_name: str + python_versions: List[str], + release_name: str, + toolkit_index_bucket_name: str, + download_method: str = DEFAULT_DOWNLOAD_METHOD, + oci_repository: str = DEFAULT_TOOLKIT_BUNDLE_OCI_REPOSITORY, ): """Download the dependencies for the given Python versions and release name.""" @@ -78,6 +179,8 @@ def submit_downloading( release_name, python_version, toolkit_index_bucket_name, + download_method, + oci_repository, ) for python_version in python_versions ] @@ -151,9 +254,19 @@ def main(): sys.exit(1) toolkit_index_bucket_name = os.getenv("TOOLKIT_INDEX_BUCKET_NAME") + download_method = os.getenv("TOOLKIT_DOWNLOAD_METHOD", DEFAULT_DOWNLOAD_METHOD) + oci_repository = os.getenv( + "TOOLKIT_BUNDLE_OCI_REPOSITORY", DEFAULT_TOOLKIT_BUNDLE_OCI_REPOSITORY + ) cleanup_old_versions(BASE_PATH, release_name) - submit_downloading(python_versions, release_name, toolkit_index_bucket_name) + submit_downloading( + python_versions, + release_name, + toolkit_index_bucket_name, + download_method, + oci_repository, + ) end_time = datetime.datetime.now() print("End time:", end_time) diff --git a/tests/unit/test_cache_downloader.py b/tests/unit/test_cache_downloader.py index 2e09a58..6cbd6a2 100644 --- a/tests/unit/test_cache_downloader.py +++ b/tests/unit/test_cache_downloader.py @@ -21,6 +21,7 @@ _spec.loader.exec_module(_mod) cleanup_old_versions = _mod.cleanup_old_versions +build_toolkit_bundle_ref = _mod.build_toolkit_bundle_ref download_dependency = _mod.download_dependency @@ -173,6 +174,14 @@ def _fail_on_v1(path, *args, **kwargs): class TestDownloadDependency: + def test_build_toolkit_bundle_ref(self): + assert ( + build_toolkit_bundle_ref( + "v1.0.0", "3.11", "registry.example.com/deepnote/toolkit-bundle" + ) + == "registry.example.com/deepnote/toolkit-bundle:v1.0.0-python3.11-tar-zst" + ) + def test_skips_download_when_done_file_exists(self, tmp_path): """Already-cached version should not trigger a download.""" release = "v1.0.0" @@ -214,4 +223,77 @@ def test_downloads_when_done_file_missing(self, tmp_path): download_dependency(release, py_ver, "fake-bucket") assert mock_popen.call_count == 2 + assert mock_popen.call_args_list[0][0][0] == [ + "aws", + "s3", + "cp", + "--no-sign-request", + "s3://fake-bucket/deepnote-toolkit/v1.0.0/python3.11.tar", + "-", + ] + assert mock_popen.call_args_list[1][0][0] == [ + "tar", + "-xf", + "-", + "-C", + str(version_path), + ] + assert (version_path / f"{py_ver}-done").exists() + + def test_downloads_from_oci_zstd_when_enabled(self, tmp_path): + """The OCI zstd feature flag should stream regctl through zstd into tar.""" + release = "v1.0.0" + py_ver = "3.11" + version_path = tmp_path / release / f"python{py_ver}" + version_path.mkdir(parents=True) + + mock_regctl = MagicMock() + mock_regctl.stdout = MagicMock() + mock_regctl.stderr = MagicMock(read=MagicMock(return_value=b"")) + mock_regctl.wait.return_value = 0 + + mock_zstd = MagicMock() + mock_zstd.stdout = MagicMock() + mock_zstd.stderr = MagicMock(read=MagicMock(return_value=b"")) + mock_zstd.wait.return_value = 0 + + mock_tar = MagicMock() + mock_tar.communicate.return_value = (b"", b"") + mock_tar.returncode = 0 + + with ( + patch.object(_mod, "BASE_PATH", str(tmp_path)), + patch.object( + _mod.subprocess, + "Popen", + side_effect=[mock_regctl, mock_zstd, mock_tar], + ) as mock_popen, + ): + download_dependency( + release, + py_ver, + "fake-bucket", + download_method="oci-zstd", + oci_repository="registry.example.com/deepnote/toolkit-bundle", + ) + + assert mock_popen.call_count == 3 + assert mock_popen.call_args_list[0][0][0] == [ + "regctl", + "artifact", + "get", + "--file", + "python3.11.tar.zst", + "registry.example.com/deepnote/toolkit-bundle:v1.0.0-python3.11-tar-zst", + ] + assert mock_popen.call_args_list[1][0][0] == ["zstd", "-dc"] + assert mock_popen.call_args_list[2][0][0] == [ + "tar", + "-xf", + "-", + "-C", + str(version_path), + ] + mock_regctl.stdout.close.assert_called_once() + mock_zstd.stdout.close.assert_called_once() assert (version_path / f"{py_ver}-done").exists() From a2d70bb91edcd026161980412267cfacfee603c5 Mon Sep 17 00:00:00 2001 From: Michal Baumgartner Date: Mon, 6 Jul 2026 13:53:49 +0200 Subject: [PATCH 10/10] chore: Incorporate PR review suggestions --- .../cache-downloader/cache-downloader.py | 134 +++++++++++------- tests/unit/test_cache_downloader.py | 27 +++- 2 files changed, 109 insertions(+), 52 deletions(-) diff --git a/dockerfiles/cache-downloader/cache-downloader.py b/dockerfiles/cache-downloader/cache-downloader.py index a74af16..99c0a2c 100644 --- a/dockerfiles/cache-downloader/cache-downloader.py +++ b/dockerfiles/cache-downloader/cache-downloader.py @@ -4,9 +4,10 @@ import shutil import subprocess import sys +import tempfile from concurrent.futures import ThreadPoolExecutor, as_completed from pathlib import Path -from typing import List +from typing import BinaryIO, List BASE_PATH = "/deepnote-toolkit/" DEFAULT_DOWNLOAD_METHOD = "s3" @@ -22,10 +23,25 @@ def build_toolkit_bundle_ref( return f"{oci_repository}:{release_name}-python{python_version}-tar-zst" -def _read_stderr(process: subprocess.Popen) -> bytes: - if process.stderr is None: - return b"" - return process.stderr.read() +def validate_download_config( + download_method: str, + toolkit_index_bucket_name: str | None, + oci_repository: str | None, +) -> None: + valid_download_methods = {DEFAULT_DOWNLOAD_METHOD, OCI_ZSTD_DOWNLOAD_METHOD} + if download_method not in valid_download_methods: + raise ValueError(f"unsupported TOOLKIT_DOWNLOAD_METHOD={download_method!r}") + if download_method == DEFAULT_DOWNLOAD_METHOD and not toolkit_index_bucket_name: + raise ValueError("TOOLKIT_INDEX_BUCKET_NAME environment variable is not set") + if download_method == OCI_ZSTD_DOWNLOAD_METHOD and not oci_repository: + raise ValueError( + "TOOLKIT_BUNDLE_OCI_REPOSITORY environment variable is not set" + ) + + +def _read_file(file: BinaryIO) -> bytes: + file.seek(0) + return file.read() def download_dependency_from_s3( @@ -37,28 +53,34 @@ def download_dependency_from_s3( s3_path = f"s3://{toolkit_index_bucket_name}/deepnote-toolkit/{release_name}/python{python_version}.tar" print(f"{datetime.datetime.now()}: Downloading {release_name} {s3_path}") - aws_process = subprocess.Popen( - ["aws", "s3", "cp", "--no-sign-request", s3_path, "-"], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - tar_process = subprocess.Popen( - ["tar", "-xf", "-", "-C", version_path], - stdin=aws_process.stdout, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - aws_process.stdout.close() # Allow aws_process to receive a SIGPIPE if tar_process exits. - _, tar_process_stderr = tar_process.communicate() - aws_process_returncode = aws_process.wait() + with ( + tempfile.TemporaryFile() as aws_process_stderr, + tempfile.TemporaryFile() as tar_process_stderr, + ): + aws_process = subprocess.Popen( + ["aws", "s3", "cp", "--no-sign-request", s3_path, "-"], + stdout=subprocess.PIPE, + stderr=aws_process_stderr, + ) + tar_process = subprocess.Popen( + ["tar", "-xf", "-", "-C", version_path], + stdin=aws_process.stdout, + stdout=subprocess.DEVNULL, + stderr=tar_process_stderr, + ) + aws_process.stdout.close() # Allow aws_process to receive a SIGPIPE if tar_process exits. + tar_process_returncode = tar_process.wait() + aws_process_returncode = aws_process.wait() + aws_process_stderr_bytes = _read_file(aws_process_stderr) + tar_process_stderr_bytes = _read_file(tar_process_stderr) if aws_process_returncode != 0: raise Exception( - f"Error downloading {release_name} (aws s3 command failed): aws stderr: {_read_stderr(aws_process)}, tar stderr: {tar_process_stderr}" + f"Error downloading {release_name} (aws s3 command failed): aws stderr: {aws_process_stderr_bytes}, tar stderr: {tar_process_stderr_bytes}" ) - if tar_process.returncode != 0: + if tar_process_returncode != 0: raise Exception( - f"Error downloading {release_name} (tar command failed): {tar_process_stderr}" + f"Error downloading {release_name} (tar command failed): {tar_process_stderr_bytes}" ) print( @@ -78,41 +100,49 @@ def download_dependency_from_oci_zstd( ) print(f"{datetime.datetime.now()}: Downloading {release_name} {artifact_ref}") - regctl_process = subprocess.Popen( - ["regctl", "artifact", "get", "--file", artifact_file, artifact_ref], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - zstd_process = subprocess.Popen( - ["zstd", "-dc"], - stdin=regctl_process.stdout, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - regctl_process.stdout.close() - tar_process = subprocess.Popen( - ["tar", "-xf", "-", "-C", version_path], - stdin=zstd_process.stdout, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - zstd_process.stdout.close() + with ( + tempfile.TemporaryFile() as regctl_process_stderr, + tempfile.TemporaryFile() as zstd_process_stderr, + tempfile.TemporaryFile() as tar_process_stderr, + ): + regctl_process = subprocess.Popen( + ["regctl", "artifact", "get", "--file", artifact_file, artifact_ref], + stdout=subprocess.PIPE, + stderr=regctl_process_stderr, + ) + zstd_process = subprocess.Popen( + ["zstd", "-dc"], + stdin=regctl_process.stdout, + stdout=subprocess.PIPE, + stderr=zstd_process_stderr, + ) + regctl_process.stdout.close() + tar_process = subprocess.Popen( + ["tar", "-xf", "-", "-C", version_path], + stdin=zstd_process.stdout, + stdout=subprocess.DEVNULL, + stderr=tar_process_stderr, + ) + zstd_process.stdout.close() - _, tar_process_stderr = tar_process.communicate() - zstd_process_returncode = zstd_process.wait() - regctl_process_returncode = regctl_process.wait() + tar_process_returncode = tar_process.wait() + zstd_process_returncode = zstd_process.wait() + regctl_process_returncode = regctl_process.wait() + regctl_process_stderr_bytes = _read_file(regctl_process_stderr) + zstd_process_stderr_bytes = _read_file(zstd_process_stderr) + tar_process_stderr_bytes = _read_file(tar_process_stderr) if regctl_process_returncode != 0: raise Exception( - f"Error downloading {release_name} (regctl command failed): regctl stderr: {_read_stderr(regctl_process)}, zstd stderr: {_read_stderr(zstd_process)}, tar stderr: {tar_process_stderr}" + f"Error downloading {release_name} (regctl command failed): regctl stderr: {regctl_process_stderr_bytes}, zstd stderr: {zstd_process_stderr_bytes}, tar stderr: {tar_process_stderr_bytes}" ) if zstd_process_returncode != 0: raise Exception( - f"Error downloading {release_name} (zstd command failed): zstd stderr: {_read_stderr(zstd_process)}, tar stderr: {tar_process_stderr}" + f"Error downloading {release_name} (zstd command failed): zstd stderr: {zstd_process_stderr_bytes}, tar stderr: {tar_process_stderr_bytes}" ) - if tar_process.returncode != 0: + if tar_process_returncode != 0: raise Exception( - f"Error downloading {release_name} (tar command failed): {tar_process_stderr}" + f"Error downloading {release_name} (tar command failed): {tar_process_stderr_bytes}" ) print( @@ -259,6 +289,14 @@ def main(): "TOOLKIT_BUNDLE_OCI_REPOSITORY", DEFAULT_TOOLKIT_BUNDLE_OCI_REPOSITORY ) + try: + validate_download_config( + download_method, toolkit_index_bucket_name, oci_repository + ) + except ValueError as exc: + print(f"Error: {exc}") + sys.exit(1) + cleanup_old_versions(BASE_PATH, release_name) submit_downloading( python_versions, diff --git a/tests/unit/test_cache_downloader.py b/tests/unit/test_cache_downloader.py index 6cbd6a2..20d33a5 100644 --- a/tests/unit/test_cache_downloader.py +++ b/tests/unit/test_cache_downloader.py @@ -23,6 +23,7 @@ cleanup_old_versions = _mod.cleanup_old_versions build_toolkit_bundle_ref = _mod.build_toolkit_bundle_ref download_dependency = _mod.download_dependency +validate_download_config = _mod.validate_download_config class TestCleanupOldVersions: @@ -174,6 +175,26 @@ def _fail_on_v1(path, *args, **kwargs): class TestDownloadDependency: + def test_validate_download_config_accepts_s3_bucket(self): + validate_download_config("s3", "fake-bucket", None) + + def test_validate_download_config_accepts_oci_repository(self): + validate_download_config( + "oci-zstd", None, "registry.example.com/deepnote/toolkit-bundle" + ) + + def test_validate_download_config_rejects_unknown_method(self): + with pytest.raises(ValueError, match="unsupported TOOLKIT_DOWNLOAD_METHOD"): + validate_download_config("unknown", "fake-bucket", None) + + def test_validate_download_config_rejects_missing_s3_bucket(self): + with pytest.raises(ValueError, match="TOOLKIT_INDEX_BUCKET_NAME"): + validate_download_config("s3", None, None) + + def test_validate_download_config_rejects_missing_oci_repository(self): + with pytest.raises(ValueError, match="TOOLKIT_BUNDLE_OCI_REPOSITORY"): + validate_download_config("oci-zstd", None, "") + def test_build_toolkit_bundle_ref(self): assert ( build_toolkit_bundle_ref( @@ -211,8 +232,7 @@ def test_downloads_when_done_file_missing(self, tmp_path): mock_aws.wait.return_value = 0 mock_tar = MagicMock() - mock_tar.communicate.return_value = (b"", b"") - mock_tar.returncode = 0 + mock_tar.wait.return_value = 0 with ( patch.object(_mod, "BASE_PATH", str(tmp_path)), @@ -258,8 +278,7 @@ def test_downloads_from_oci_zstd_when_enabled(self, tmp_path): mock_zstd.wait.return_value = 0 mock_tar = MagicMock() - mock_tar.communicate.return_value = (b"", b"") - mock_tar.returncode = 0 + mock_tar.wait.return_value = 0 with ( patch.object(_mod, "BASE_PATH", str(tmp_path)),