diff --git a/.github/actions/push-oci-artifact/action.yml b/.github/actions/push-oci-artifact/action.yml index 676bf26..7d4d94a 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,12 +83,18 @@ 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}" . - 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}')" @@ -101,3 +128,82 @@ 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 + + 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}" + "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 + + 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}" + + if [ "${file_count}" -eq 0 ]; then + echo "Error: no OCI artifact files were provided" >&2 + exit 1 + fi + + ( + cd "${workdir}" + 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..75d8479 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -89,6 +89,22 @@ jobs: run: | set -euo pipefail ./bin/build + sudo chown -R "$(id -u):$(id -g)" dist + + - 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 -6 -f \ + -o "dist/python${PYTHON_VERSION}.tar.zst" \ + "dist/python${PYTHON_VERSION}.tar" # STAGING UPLOAD - name: Upload to staging @@ -128,6 +144,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 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..99c0a2c 100644 --- a/dockerfiles/cache-downloader/cache-downloader.py +++ b/dockerfiles/cache-downloader/cache-downloader.py @@ -4,70 +4,201 @@ 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" +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" +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" ) - return - # Create the version directory if it doesn't exist - os.makedirs(version_path, exist_ok=True) +def _read_file(file: BinaryIO) -> bytes: + file.seek(0) + return file.read() + + +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, - 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() - ) # Wait for tar_process to complete - 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: {aws_process.stderr.read()}, 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}" ) - # Check for errors - 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( 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}") + + 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_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: {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: {zstd_process_stderr_bytes}, tar stderr: {tar_process_stderr_bytes}" + ) + if tar_process_returncode != 0: + raise Exception( + f"Error downloading {release_name} (tar command failed): {tar_process_stderr_bytes}" + ) + + 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 +209,8 @@ def submit_downloading( release_name, python_version, toolkit_index_bucket_name, + download_method, + oci_repository, ) for python_version in python_versions ] @@ -151,9 +284,27 @@ 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 + ) + + 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, 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..20d33a5 100644 --- a/tests/unit/test_cache_downloader.py +++ b/tests/unit/test_cache_downloader.py @@ -21,7 +21,9 @@ _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 +validate_download_config = _mod.validate_download_config class TestCleanupOldVersions: @@ -173,6 +175,34 @@ 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( + "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" @@ -202,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)), @@ -214,4 +243,76 @@ 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.wait.return_value = 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()