diff --git a/.github/config/image/whisperx/ec2-amzn2023.yml b/.github/config/image/whisperx/ec2-amzn2023.yml new file mode 100644 index 000000000000..36a1eb927b22 --- /dev/null +++ b/.github/config/image/whisperx/ec2-amzn2023.yml @@ -0,0 +1,38 @@ +image: + name: "whisperx-ec2-amzn2023" + description: "WhisperX ASR (VAD + faster-whisper transcription + wav2vec2 alignment + pyannote diarization) for EC2 instances (AL2023, GPU)" + +metadata: + framework: "whisperx" + framework_version: "3.8.6" + os_version: "amzn2023" + customer_type: "ec2" + arch_type: "x86" + device_type: "gpu" + job_type: "inference" + prod_image: "whisperx:3.8.6-ec2-gpu-cuda-amzn2023" + +build: + dockerfile: "docker/whisperx/Dockerfile.amzn2023" + target: "whisperx-ec2-amzn2023" + python_version: "3.12" + cuda_version: "12.8.2" + torch_version: "2.8.0" + torchaudio_version: "2.8.0" + whisperx_version: "3.8.6" + faster_whisper_version: "1.2.1" + ctranslate2_version: "4.8.0" + pyannote_audio_version: "4.0.7" + fastapi_version: "0.121.0" + uvicorn_version: "0.40.0" + python_multipart_version: "0.0.30" + dlc_major_version: "1" + dlc_minor_version: "0" + +release: + release: false + force_release: false + public_registry: false + private_registry: false + enable_soci: false + environment: "production" diff --git a/.github/config/image/whisperx/sagemaker-amzn2023.yml b/.github/config/image/whisperx/sagemaker-amzn2023.yml new file mode 100644 index 000000000000..bb621f4aa642 --- /dev/null +++ b/.github/config/image/whisperx/sagemaker-amzn2023.yml @@ -0,0 +1,39 @@ +image: + name: "whisperx-sagemaker-amzn2023" + description: "WhisperX ASR (VAD + faster-whisper transcription + wav2vec2 alignment + pyannote diarization) for SageMaker (AL2023, GPU)" + +metadata: + framework: "whisperx" + framework_version: "3.8.6" + os_version: "amzn2023" + customer_type: "sagemaker" + platform: "sagemaker" + arch_type: "x86" + device_type: "gpu" + job_type: "inference" + prod_image: "whisperx:3.8.6-sagemaker-gpu-cuda-amzn2023" + +build: + dockerfile: "docker/whisperx/Dockerfile.amzn2023" + target: "whisperx-sagemaker-amzn2023" + python_version: "3.12" + cuda_version: "12.8.2" + torch_version: "2.8.0" + torchaudio_version: "2.8.0" + whisperx_version: "3.8.6" + faster_whisper_version: "1.2.1" + ctranslate2_version: "4.8.0" + pyannote_audio_version: "4.0.7" + fastapi_version: "0.121.0" + uvicorn_version: "0.40.0" + python_multipart_version: "0.0.30" + dlc_major_version: "1" + dlc_minor_version: "0" + +release: + release: false + force_release: false + public_registry: false + private_registry: false + enable_soci: false + environment: "production" diff --git a/.github/workflows/whisperx.pipeline.yml b/.github/workflows/whisperx.pipeline.yml new file mode 100644 index 000000000000..2aced5a920d6 --- /dev/null +++ b/.github/workflows/whisperx.pipeline.yml @@ -0,0 +1,215 @@ +name: "Pipeline - WhisperX" + +on: + workflow_call: + inputs: + config-file: + description: "Path to image config YAML" + required: true + type: string + tag-suffix: + description: "CI tag suffix (run_id or pr-NUMBER)" + required: true + type: string + build: + description: "Whether to build the image (false = skip build, test against prod)" + required: false + type: boolean + default: true + run-sanity-test: + description: "Run sanity tests" + required: false + type: boolean + default: true + run-unit-test: + description: "Run CPU-only unit tests" + required: false + type: boolean + default: true + run-security-test: + description: "Run security tests" + required: false + type: boolean + default: true + run-telemetry-test: + description: "Run telemetry tests" + required: false + type: boolean + default: true + run-ec2-test: + description: "Run EC2 functional tests (gated to ec2 configs)" + required: false + type: boolean + default: true + run-sagemaker-test: + description: "Run SageMaker endpoint tests (gated to sagemaker configs)" + required: false + type: boolean + default: true + release: + description: "Run release job after tests pass" + required: false + type: boolean + default: false + outputs: + image-uri: + description: "Built image URI (empty if build was skipped)" + value: ${{ jobs.build.outputs.image-uri }} + +jobs: + ci-config: + runs-on: ubuntu-latest + outputs: + customer-type: ${{ steps.parse.outputs.customer-type }} + steps: + - uses: actions/checkout@v6 + - id: parse + run: | + echo "customer-type=$(yq '.metadata.customer_type' '${{ inputs.config-file }}')" >> $GITHUB_OUTPUT + + build: + if: ${{ inputs.build }} + concurrency: + group: ${{ github.workflow }}-build-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: true + runs-on: + - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }} + fleet:x86-build-runner + buildspec-override:true + timeout-minutes: 120 + outputs: + image-uri: ${{ steps.build.outputs.image-uri }} + steps: + - uses: actions/checkout@v6 + - id: build + uses: ./.github/actions/build-image + with: + config-file: ${{ inputs.config-file }} + aws-account-id: ${{ vars.CI_AWS_ACCOUNT_ID }} + aws-region: ${{ vars.AWS_REGION }} + tag-suffix: ${{ inputs.tag-suffix }} + + sanity-test: + if: ${{ always() && !failure() && !cancelled() && inputs.run-sanity-test }} + needs: [build] + concurrency: + group: ${{ github.workflow }}-sanity-test-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: true + uses: ./.github/workflows/_reusable.sanity-tests.yml + with: + config-file: ${{ inputs.config-file }} + image-uri: ${{ needs.build.outputs.image-uri || '' }} + + unit-test: + # CPU-only, image-independent tests — no `needs: [build]` so they run + # immediately in parallel with the image build and still report even if the + # build fails (they never consume the image; image-uri is unused). + if: ${{ inputs.run-unit-test }} + concurrency: + group: ${{ github.workflow }}-unit-test-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: true + uses: ./.github/workflows/whisperx.tests-unit.yml + with: + config-file: ${{ inputs.config-file }} + image-uri: '' + + security-test: + if: ${{ always() && !failure() && !cancelled() && inputs.run-security-test }} + needs: [build] + concurrency: + group: ${{ github.workflow }}-security-test-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: true + uses: ./.github/workflows/_reusable.security-tests.yml + with: + config-file: ${{ inputs.config-file }} + image-uri: ${{ needs.build.outputs.image-uri || '' }} + + telemetry-test: + if: ${{ always() && !failure() && !cancelled() && inputs.run-telemetry-test }} + needs: [build] + concurrency: + group: ${{ github.workflow }}-telemetry-test-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: false + uses: ./.github/workflows/_reusable.telemetry-tests.yml + with: + config-file: ${{ inputs.config-file }} + image-uri: ${{ needs.build.outputs.image-uri || '' }} + + ec2-test: + if: ${{ always() && !failure() && !cancelled() && inputs.run-ec2-test && needs.ci-config.outputs.customer-type == 'ec2' }} + needs: [build, ci-config] + concurrency: + group: ${{ github.workflow }}-ec2-test-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: true + uses: ./.github/workflows/whisperx.tests-ec2.yml + with: + config-file: ${{ inputs.config-file }} + image-uri: ${{ needs.build.outputs.image-uri || '' }} + + sagemaker-test: + if: ${{ always() && !failure() && !cancelled() && inputs.run-sagemaker-test && needs.ci-config.outputs.customer-type == 'sagemaker' }} + needs: [build, ci-config] + concurrency: + group: ${{ github.workflow }}-sagemaker-test-${{ inputs.config-file }}-${{ github.ref }} + cancel-in-progress: false + uses: ./.github/workflows/whisperx.tests-sagemaker.yml + with: + config-file: ${{ inputs.config-file }} + image-uri: ${{ needs.build.outputs.image-uri || '' }} + + release-gate: + if: >- + github.ref == 'refs/heads/main' && + (contains(github.workflow_ref, '.autorelease') || contains(github.workflow_ref, '.dispatch-release')) && + !contains(github.workflow_ref, '.pr') && + inputs.release && + !failure() && !cancelled() && + needs.build.result == 'success' + needs: [build, sanity-test, unit-test, security-test, telemetry-test, ec2-test, sagemaker-test] + runs-on: ubuntu-latest + outputs: + should-release: ${{ steps.check.outputs.should-release }} + release-spec: ${{ steps.spec.outputs.release-spec }} + environment: ${{ steps.check.outputs.environment }} + steps: + - uses: actions/checkout@v6 + - id: check + run: | + SHOULD_RELEASE=$(yq '.release.release' "${{ inputs.config-file }}") + ENVIRONMENT=$(yq '.release.environment' "${{ inputs.config-file }}") + echo "should-release=${SHOULD_RELEASE}" >> $GITHUB_OUTPUT + echo "environment=${ENVIRONMENT}" >> $GITHUB_OUTPUT + - id: spec + if: steps.check.outputs.should-release == 'true' + uses: ./.github/actions/generate-release-spec + with: + config-file: ${{ inputs.config-file }} + + release: + if: ${{ !failure() && !cancelled() && needs.release-gate.outputs.should-release == 'true' }} + needs: [build, release-gate] + uses: ./.github/workflows/_reusable.release-image.yml + with: + source-image-uri: ${{ needs.build.outputs.image-uri }} + release-spec: ${{ needs.release-gate.outputs.release-spec }} + environment: ${{ needs.release-gate.outputs.environment }} + aws-region: ${{ vars.AWS_REGION }} + secrets: inherit + + notify-failure: + if: >- + failure() && + github.ref == 'refs/heads/main' && + (contains(github.workflow_ref, '.autorelease') || contains(github.workflow_ref, '.dispatch-release')) && + !contains(github.workflow_ref, '.pr') && + inputs.release + needs: [build, sanity-test, unit-test, security-test, telemetry-test, ec2-test, sagemaker-test, release-gate, release] + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + - uses: ./.github/actions/notify-slack-failure + with: + webhook-url: ${{ secrets.SLACK_RELEASE_FAILURE_WEBHOOK_URL }} + workflow-name: ${{ github.workflow }} + config-file: ${{ inputs.config-file }} + run-url: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }} diff --git a/.github/workflows/whisperx.pr-amzn2023.yml b/.github/workflows/whisperx.pr-amzn2023.yml new file mode 100644 index 000000000000..a20936bd29e3 --- /dev/null +++ b/.github/workflows/whisperx.pr-amzn2023.yml @@ -0,0 +1,108 @@ +name: "PR - WhisperX AL2023" + +on: + workflow_dispatch: + pull_request: + branches: [main] + types: [opened, reopened, synchronize] + paths: + - ".github/config/image/whisperx/*.yml" + - ".github/workflows/whisperx.pipeline.yml" + - ".github/workflows/whisperx.pr-amzn2023.yml" + - ".github/workflows/whisperx.tests-ec2.yml" + - ".github/workflows/whisperx.tests-sagemaker.yml" + - "docker/whisperx/**" + - "scripts/docker/whisperx/**" + - "scripts/docker/common/**" + - "scripts/docker/telemetry/**" + - "scripts/ci/build/whisperx/**" + - "test/sanity/**" + - "test/telemetry/**" + - "test/whisperx/**" + - "test/security/data/ecr_scan_allowlist/whisperx/**" + - "!docs/**" + +permissions: + contents: read + pull-requests: read + +jobs: + gatekeeper: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + with: + ref: ${{ github.event.pull_request.base.sha }} + fetch-depth: 1 + - uses: ./.github/actions/pr-permission-gate + + check-changes: + needs: [gatekeeper] + runs-on: ubuntu-latest + outputs: + build-change: ${{ steps.changes.outputs.build-change }} + sanity-test-change: ${{ steps.changes.outputs.sanity-test-change }} + unit-test-change: ${{ steps.changes.outputs.unit-test-change }} + telemetry-test-change: ${{ steps.changes.outputs.telemetry-test-change }} + ec2-test-change: ${{ steps.changes.outputs.ec2-test-change }} + sagemaker-test-change: ${{ steps.changes.outputs.sagemaker-test-change }} + steps: + - uses: actions/checkout@v6 + - uses: dorny/paths-filter@v4 + id: changes + with: + filters: | + build-change: + - ".github/config/image/whisperx/*.yml" + - "docker/whisperx/**" + - "scripts/docker/whisperx/**" + - "scripts/docker/common/**" + - "scripts/docker/telemetry/**" + - "scripts/ci/build/whisperx/**" + - "test/security/data/ecr_scan_allowlist/whisperx/**" + sanity-test-change: + - "test/sanity/**" + unit-test-change: + - "scripts/docker/whisperx/**" + - "test/whisperx/test_*.py" + - ".github/workflows/whisperx.tests-unit.yml" + telemetry-test-change: + - "test/telemetry/**" + ec2-test-change: + - "test/whisperx/ec2/**" + - ".github/workflows/whisperx.tests-ec2.yml" + sagemaker-test-change: + - "test/whisperx/sagemaker/**" + - ".github/workflows/whisperx.tests-sagemaker.yml" + + discover: + needs: [gatekeeper] + runs-on: ubuntu-latest + outputs: + configs: ${{ steps.discover.outputs.configs }} + steps: + - uses: actions/checkout@v6 + - id: discover + uses: ./.github/actions/discover-configs + with: + pattern: ".github/config/image/whisperx/*.yml" + + pipeline: + needs: [check-changes, discover] + strategy: + fail-fast: false + matrix: + include: ${{ fromJson(needs.discover.outputs.configs) }} + uses: ./.github/workflows/whisperx.pipeline.yml + with: + config-file: ${{ matrix.config_file }} + tag-suffix: pr-${{ github.event.pull_request.number }} + build: ${{ needs.check-changes.outputs.build-change == 'true' }} + run-sanity-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.sanity-test-change == 'true' }} + run-unit-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.unit-test-change == 'true' }} + run-security-test: ${{ needs.check-changes.outputs.build-change == 'true' }} + run-telemetry-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.telemetry-test-change == 'true' }} + run-ec2-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.ec2-test-change == 'true' }} + run-sagemaker-test: ${{ needs.check-changes.outputs.build-change == 'true' || needs.check-changes.outputs.sagemaker-test-change == 'true' }} + release: false + secrets: inherit diff --git a/.github/workflows/whisperx.tests-ec2.yml b/.github/workflows/whisperx.tests-ec2.yml new file mode 100644 index 000000000000..14a2ae71796b --- /dev/null +++ b/.github/workflows/whisperx.tests-ec2.yml @@ -0,0 +1,85 @@ +name: Reusable WhisperX EC2 Tests + +on: + workflow_call: + inputs: + config-file: + description: "Path to image config YAML" + required: true + type: string + image-uri: + description: "Image URI to test. If empty, derives prod URI from config." + required: false + type: string + default: "" + +permissions: + contents: read + +jobs: + preflight: + runs-on: + - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }} + fleet:default-runner + buildspec-override:true + outputs: + image-uri: ${{ steps.image.outputs.image-uri }} + aws-account-id: ${{ steps.image.outputs.aws-account-id }} + image-exists: ${{ steps.check.outputs.exists }} + device-type: ${{ steps.config.outputs.device-type }} + steps: + - uses: actions/checkout@v6 + + - name: Resolve image URI + id: image + uses: ./.github/actions/resolve-image-uri + with: + image-uri: ${{ inputs.image-uri }} + ci-aws-account-id: ${{ vars.CI_AWS_ACCOUNT_ID }} + prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID }} + aws-region: ${{ vars.AWS_REGION }} + config-file: ${{ inputs.config-file }} + + - name: Check image exists + id: check + uses: ./.github/actions/check-image-exists + with: + image-uri: ${{ steps.image.outputs.image-uri }} + + - name: Parse config + id: config + run: | + echo "device-type=$(yq '.metadata.device_type' '${{ inputs.config-file }}')" >> $GITHUB_OUTPUT + + ec2-gpu: + if: ${{ needs.preflight.outputs.image-exists == 'true' && needs.preflight.outputs.device-type == 'gpu' }} + needs: [preflight] + timeout-minutes: 60 + runs-on: + - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }} + fleet:x86-g6xl-runner + buildspec-override:true + steps: + - uses: actions/checkout@v6 + + - name: ECR login + uses: ./.github/actions/ecr-authenticate + with: + aws-account-id: ${{ needs.preflight.outputs.aws-account-id }} + aws-region: ${{ vars.AWS_REGION }} + image-uri: ${{ needs.preflight.outputs.image-uri }} + + - name: Install test dependencies + run: | + uv venv --python 3.12 + source .venv/bin/activate + uv pip install -r test/requirements.txt + uv pip install -r test/whisperx/ec2/requirements.txt + + - name: Run EC2 GPU tests + run: | + source .venv/bin/activate + cd test/ + python3 -m pytest -v --tb=short -rA --log-cli-level=INFO \ + --image-uri ${{ needs.preflight.outputs.image-uri }} \ + whisperx/ec2 diff --git a/.github/workflows/whisperx.tests-sagemaker.yml b/.github/workflows/whisperx.tests-sagemaker.yml new file mode 100644 index 000000000000..ebb8d551238c --- /dev/null +++ b/.github/workflows/whisperx.tests-sagemaker.yml @@ -0,0 +1,79 @@ +name: Reusable WhisperX SageMaker Tests + +on: + workflow_call: + inputs: + config-file: + description: "Path to image config YAML" + required: true + type: string + image-uri: + description: "Image URI to test. If empty, derives prod URI from config." + required: false + type: string + default: "" + +permissions: + contents: read + +jobs: + preflight: + runs-on: + - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }} + fleet:default-runner + buildspec-override:true + outputs: + image-uri: ${{ steps.image.outputs.image-uri }} + aws-account-id: ${{ steps.image.outputs.aws-account-id }} + image-exists: ${{ steps.check.outputs.exists }} + steps: + - uses: actions/checkout@v6 + + - name: Resolve image URI + id: image + uses: ./.github/actions/resolve-image-uri + with: + image-uri: ${{ inputs.image-uri }} + ci-aws-account-id: ${{ vars.CI_AWS_ACCOUNT_ID }} + prod-aws-account-id: ${{ vars.PROD_AWS_ACCOUNT_ID }} + aws-region: ${{ vars.AWS_REGION }} + config-file: ${{ inputs.config-file }} + + - name: Check image exists + id: check + uses: ./.github/actions/check-image-exists + with: + image-uri: ${{ steps.image.outputs.image-uri }} + + endpoint-test: + if: ${{ needs.preflight.outputs.image-exists == 'true' }} + needs: [preflight] + timeout-minutes: 120 + runs-on: + - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }} + fleet:default-runner + buildspec-override:true + steps: + - uses: actions/checkout@v6 + + - name: ECR login + uses: ./.github/actions/ecr-authenticate + with: + aws-account-id: ${{ needs.preflight.outputs.aws-account-id }} + aws-region: ${{ vars.AWS_REGION }} + image-uri: ${{ needs.preflight.outputs.image-uri }} + + - name: Install test dependencies + run: | + uv venv --python 3.12 + source .venv/bin/activate + uv pip install -r test/requirements.txt + uv pip install -r test/whisperx/sagemaker/requirements.txt + + - name: Run SageMaker endpoint tests (sync + async) + run: | + source .venv/bin/activate + cd test/ + python3 -m pytest -vs -rA \ + --image-uri ${{ needs.preflight.outputs.image-uri }} \ + whisperx/sagemaker diff --git a/.github/workflows/whisperx.tests-unit.yml b/.github/workflows/whisperx.tests-unit.yml new file mode 100644 index 000000000000..e521c16650bf --- /dev/null +++ b/.github/workflows/whisperx.tests-unit.yml @@ -0,0 +1,43 @@ +name: Reusable WhisperX Unit Tests + +on: + workflow_call: + inputs: + config-file: + description: "Path to image config YAML" + required: true + type: string + image-uri: + description: "Unused (unit tests are image-independent); kept for signature parity." + required: false + type: string + default: "" + +permissions: + contents: read + +jobs: + unit-test: + runs-on: + - codebuild-runner-${{ github.run_id }}-${{ github.run_attempt }} + fleet:default-runner + buildspec-override:true + timeout-minutes: 15 + steps: + - uses: actions/checkout@v6 + + - name: Install test dependencies + run: | + uv venv --python 3.12 + source .venv/bin/activate + uv pip install pytest anyio + + - name: Run WhisperX unit tests (CPU-only, image-independent) + run: | + source .venv/bin/activate + cd test/ + # --noconftest: skip the root test/conftest.py, which imports boto3/fabric + # (heavy AWS deps) that these image-independent tests don't need or install. + python3 -m pytest -v --tb=short -rA --noconftest \ + whisperx/test_server_options.py \ + whisperx/test_transcribe_behavior.py diff --git a/.gitignore b/.gitignore index 06a54d75c6b4..25ff5e2f80b7 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ site/ tutorials/ .sisyphus/ docker/xgboost/prebuilt.whl +docker/whisperx/pyannote-diarization.tar.gz # tmp .github/archive diff --git a/docker/whisperx/Dockerfile.amzn2023 b/docker/whisperx/Dockerfile.amzn2023 new file mode 100644 index 000000000000..53c29a7904f6 --- /dev/null +++ b/docker/whisperx/Dockerfile.amzn2023 @@ -0,0 +1,294 @@ +# ============================================================================= +# WhisperX DLC image — AL2023, CUDA 12.8, Python 3.12 +# +# Two BuildKit targets, one Dockerfile: +# whisperx-ec2-amzn2023 (docker buildx build --target whisperx-ec2-amzn2023 …) +# whisperx-sagemaker-amzn2023 (docker buildx build --target whisperx-sagemaker-amzn2023 …) +# +# Multi-stage layout (design §7): +# builder — CUDA-runtime base, installs everything into /opt/venv, prunes. +# runtime — CUDA-base image (no /usr/local/cuda/targets), copies venv + +# models across. Final image is ~4 GB smaller than a single-stage +# build on the runtime base. +# +# In DLC CI the build-image action passes these as --build-args from the image +# config. For a standalone build the caller sources versions.env, e.g.: +# set -a; . versions.env; set +a +# docker buildx build --build-arg CUDA_VERSION --build-arg PYTHON_VERSION … +# ============================================================================= +ARG CUDA_VERSION=12.8.2 +ARG PYTHON_VERSION=3.12 + +# ============================================================================= +# STAGE: builder — installs everything into /opt/venv on the CUDA runtime +# base. The runtime base is needed here because pip-installing torch cu12x +# probes the base's CUDA installation during resolution; the final image +# uses the smaller -base- variant. +# ============================================================================= +FROM nvidia/cuda:${CUDA_VERSION}-runtime-amzn2023 AS builder + +ARG CUDA_VERSION +ARG PYTHON_VERSION + +# System packages the *build* needs. Split into "needed at runtime too" and +# "builder only" — we don't try to install them here selectively; runtime +# stage installs its own subset. +# No CVE patch here: this stage's OS layer is discarded (only /opt/venv + +# /opt/models cross into runtime), so patching it would add build time without +# affecting any shipped image. +RUN dnf upgrade -y --releasever=latest --setopt=install_weak_deps=False system-release \ + && dnf install -y spal-release \ + && dnf install -y --setopt=install_weak_deps=False \ + python${PYTHON_VERSION} python${PYTHON_VERSION}-devel python${PYTHON_VERSION}-pip \ + tar gzip git ca-certificates \ + && dnf clean all && rm -rf /var/cache/dnf + +RUN python${PYTHON_VERSION} -m venv /opt/venv +ENV PATH="/opt/venv/bin:${PATH}" +ENV VIRTUAL_ENV="/opt/venv" + +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH="/root/.local/bin:${PATH}" +ENV UV_HTTP_TIMEOUT=500 +ENV UV_INDEX_STRATEGY="unsafe-best-match" +ENV UV_LINK_MODE=copy + +ARG PYTORCH_CUDA_INDEX_BASE_URL=https://download.pytorch.org/whl +ARG TORCH_VERSION +ARG TORCHAUDIO_VERSION +RUN --mount=type=cache,target=/root/.cache/uv uv pip install \ + "torch==${TORCH_VERSION}" \ + "torchaudio==${TORCHAUDIO_VERSION}" \ + --extra-index-url ${PYTORCH_CUDA_INDEX_BASE_URL}/cu$(echo ${CUDA_VERSION} | cut -d. -f1,2 | tr -d '.') + +# ASR + diarization stack. Order matters: install ctranslate2 explicitly first +# so faster-whisper's dep resolver doesn't pull a mismatched version. +ARG WHISPERX_VERSION +ARG FASTER_WHISPER_VERSION +ARG CTRANSLATE2_VERSION +ARG PYANNOTE_AUDIO_VERSION +RUN --mount=type=cache,target=/root/.cache/uv uv pip install \ + "ctranslate2==${CTRANSLATE2_VERSION}" \ + "faster-whisper==${FASTER_WHISPER_VERSION}" \ + "whisperx==${WHISPERX_VERSION}" \ + "pyannote.audio==${PYANNOTE_AUDIO_VERSION}" + +# HTTP layer. +ARG FASTAPI_VERSION +ARG UVICORN_VERSION +ARG PYTHON_MULTIPART_VERSION +RUN --mount=type=cache,target=/root/.cache/uv uv pip install \ + "fastapi==${FASTAPI_VERSION}" \ + "uvicorn[standard]==${UVICORN_VERSION}" \ + "python-multipart==${PYTHON_MULTIPART_VERSION}" + +# ---- DLC telemetry runtime deps ---- +# The DLC telemetry script (deep_learning_container.py, run at container start via +# bashrc) imports botocore + requests. Install into the venv so the DLC telemetry +# CI check passes. Small deps; they do not affect the ASR stack. +RUN --mount=type=cache,target=/root/.cache/uv uv pip install "boto3" "requests" + +# ---- Prune ---- +# torch cu12x wheels statically reference libnccl, libcusparseLt, and a +# handful of other nvidia/* libs at import time (`from torch._C import *`). +# Uninstalling them post-install breaks `import torch`. We tried and it +# fails with "libcusparseLt.so.0: cannot open shared object file". Do NOT +# add nvidia/* packages to the uninstall list without proving torch still +# imports afterwards. +# +# Drop __pycache__ (rebuilt lazily at first import; PYTHONDONTWRITEBYTECODE +# prevents new ones) and strip debug symbols from every .so in site-packages. +# strip only removes debug info; runtime symbols stay. Saves ~500 MB. +# +# NOTE: triton (~540 MB) is torch's torch.compile JIT backend. Nothing in the +# current WhisperX stack calls torch.compile, so it's dead weight today — but +# we deliberately KEEP it: torch increasingly enables compile paths by default, +# and a future version could reach for it at *runtime* (failing only at compute, +# not import). 540 MB isn't worth that latent break. Revisit only if image size +# becomes a hard constraint AND a smoke test proves the compile path is unused. +RUN find /opt/venv -type d -name '__pycache__' -exec rm -rf {} + 2>/dev/null; \ + find /opt/venv -type f \( -name '*.so' -o -name '*.so.*' \) \ + -exec strip --strip-unneeded {} + 2>/dev/null; \ + true + +# ---- Model cache dirs ---- +# HF_HOME / TORCH_HOME are where Whisper + per-language wav2vec2 aligners cache +# lazily at runtime (design Dim 2: too many languages to bake). TORCH_HOME is +# also where whisperx's default (pyannote) VAD resolves torch.hub._get_torch_home(). +# The default pyannote VAD segmentation model ships inside the whisperx wheel +# (whisperx/assets/pytorch_model.bin, loaded from a local file — no network), so +# no VAD model is baked here. Silero VAD is not used: whisperx 3.8.6 defaults to +# vad_method="pyannote" and server.py never selects silero, so we avoid its +# torch.hub GitHub download (rate-limit-prone at build and runtime). +ENV HF_HOME=/opt/models/hf +ENV TORCH_HOME=/opt/models/torch +RUN mkdir -p ${HF_HOME} ${TORCH_HOME} + +# pyannote diarization weights: pre-downloaded from S3 by scripts/build.sh into +# the build context, extracted here into a fixed path. No HF token, no network +# at build. The tarball never reaches the final image: only /opt/models and +# /opt/venv cross into the runtime stage (this builder stage is discarded). +COPY docker/whisperx/pyannote-diarization.tar.gz /tmp/pyannote-diarization.tar.gz +RUN mkdir -p /opt/models/pyannote/speaker-diarization-community-1 \ + && tar xzf /tmp/pyannote-diarization.tar.gz \ + -C /opt/models/pyannote/speaker-diarization-community-1 \ + && rm /tmp/pyannote-diarization.tar.gz + +# ============================================================================= +# STAGE: runtime — small final image on the CUDA -base- variant. +# The -base- image ships libcuda stub + NVIDIA container hooks but NOT the +# full CUDA userspace (/usr/local/cuda/targets/*). Torch's cu128 wheel +# bundles the CUDA libs it needs under torch/lib/, so LD_LIBRARY_PATH is +# pointed there. Saves ~2.4 GB vs the -runtime- base. +# ============================================================================= +FROM nvidia/cuda:${CUDA_VERSION}-base-amzn2023 AS runtime + +ARG CUDA_VERSION +ARG PYTHON_VERSION + +# Runtime system packages only: +# python3.12 — needed to run the venv +# ffmpeg-free — audio decode (used by whisperx.load_audio → subprocess) +# espeak-ng — phoneme fallback for wav2vec2 aligner +# jack-audio-... — ffmpeg-free is linked against libjack.so.0 +# ca-certificates — needed for HF cache misses over HTTPS +# Notice: no python-devel/pip, no gcc, no git — build-only. +# Runtime OS packages. `system-release` floats the AL2023 release forward (also +# needed for the spal-release repo, which supplies ffmpeg-free + espeak-ng). +# The CVE --security sweep is NOT here: this layer's inputs are stable, so CI's +# layer cache would hit it forever and never re-patch. Instead the sweep lives +# in the ec2/sagemaker leaf stages behind CACHE_REFRESH (see below), so a +# rebuild re-patches without busting the expensive torch/model layers. +RUN dnf upgrade -y --releasever=latest --setopt=install_weak_deps=False system-release \ + && dnf install -y spal-release \ + && dnf install -y --setopt=install_weak_deps=False \ + python${PYTHON_VERSION} \ + ffmpeg-free espeak-ng jack-audio-connection-kit \ + ca-certificates \ + && dnf clean all && rm -rf /var/cache/dnf + +# Bring across the venv (pruned, stripped) and baked models. +COPY --from=builder /opt/venv /opt/venv +COPY --from=builder /opt/models /opt/models + +ENV PATH="/opt/venv/bin:${PATH}" +ENV VIRTUAL_ENV="/opt/venv" +ENV HF_HOME=/opt/models/hf +ENV TORCH_HOME=/opt/models/torch + +# No explicit LD_LIBRARY_PATH for torch's CUDA libs. Torch's cu12x wheels bake +# an $ORIGIN-relative RPATH into libtorch_cuda.so (and the nvidia/* libs use +# $ORIGIN too), so cudart/cublas/cudnn/cufft/cusolver/cusparse/nvrtc/nvjitlink — +# including the lazily-dlopen'd cuDNN engine sublibs — all resolve from inside +# the venv with no env help. Verified: `ldd` on every CUDA .so reports zero +# "not found" with LD_LIBRARY_PATH at only the base default, and +# torch.backends.cudnn.version() (which forces a real dlopen) resolves. +# The inherited base value (/usr/local/nvidia/lib{,64}) is left intact for the +# NVIDIA driver libs injected at `docker run --gpus` time. +# If a future torch wheel drops its RPATH, restore an explicit path here. + +# ---- Server code ---- +COPY scripts/docker/whisperx/server.py /opt/whisperx/server.py + +# ---- CUDA forward-compat activation (sourced by both entrypoints) ---- +# Activates /usr/local/cuda/compat when the host NVIDIA driver is older than the +# baked cu128 libs need (e.g. SageMaker ml.g4dn). No-op on new-driver / CPU +# hosts. Without this, old-driver GPU hosts fail to start with a zero-log +# CannotStartContainerError. +COPY scripts/docker/whisperx/cuda_compat.sh /opt/whisperx/cuda_compat.sh +RUN chmod +x /opt/whisperx/cuda_compat.sh + +# ---- Third-party attribution (CC-BY-4.0 for pyannote diarization weights) ---- +COPY docker/whisperx/NOTICE /NOTICE + +# ---- DLC telemetry + OSS compliance (required by DLC sanity & telemetry CI) ---- +# FRAMEWORK / FRAMEWORK_VERSION / CONTAINER_TYPE are injected automatically by the +# DLC build-image action from the image config metadata. DLC_CONTAINER_TYPE must +# equal metadata.job_type (inference). The venv is on PATH, so `python` resolves to +# /opt/venv/bin/python (where boto3/requests were installed above). +ENV DLC_CONTAINER_TYPE=inference +COPY scripts/docker/telemetry/deep_learning_container.py /usr/local/bin/deep_learning_container.py +COPY scripts/docker/telemetry/bash_telemetry.sh.template /tmp/bash_telemetry.sh.template +COPY scripts/docker/common/setup_oss_compliance.sh setup_oss_compliance.sh + +ARG FRAMEWORK +ARG FRAMEWORK_VERSION +ARG CONTAINER_TYPE=inference +RUN chmod +x /usr/local/bin/deep_learning_container.py \ + && sed -e "s/{{FRAMEWORK}}/${FRAMEWORK}/g" \ + -e "s/{{FRAMEWORK_VERSION}}/${FRAMEWORK_VERSION}/g" \ + -e "s/{{CONTAINER_TYPE}}/${CONTAINER_TYPE}/g" \ + /tmp/bash_telemetry.sh.template >/usr/local/bin/bash_telemetry.sh \ + && chmod +x /usr/local/bin/bash_telemetry.sh \ + && rm /tmp/bash_telemetry.sh.template \ + && echo 'source /usr/local/bin/bash_telemetry.sh' >>/etc/bashrc \ + && echo 'source /usr/local/bin/bash_telemetry.sh' >>/root/.bashrc \ + && bash setup_oss_compliance.sh python \ + && rm setup_oss_compliance.sh \ + && rm -rf /root/oss_compliance* /tmp/tmp* + +ENV LANG=C.UTF-8 \ + LC_ALL=C.UTF-8 \ + PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + PYTHONIOENCODING=UTF-8 \ + PYTHONPATH=/opt/whisperx + +WORKDIR /opt/whisperx + +# ============================================================================= +# TARGET: whisperx-ec2-amzn2023 +# ============================================================================= +FROM runtime AS whisperx-ec2-amzn2023 + +# ---- OS CVE patch (cache-bustable) ---- +# Drives open AL2023 security advisories to zero in the shipped image. Placed in +# the leaf stage (not runtime) on purpose: CI relies on the layer cache, and +# runtime's dnf layer has stable inputs so it would cache-hit forever and never +# re-patch. Gating this layer on CACHE_REFRESH lets CI pass +# `--build-arg CACHE_REFRESH=` to force a fresh re-patch every +# rebuild — re-applying advisories published since the cached layer — WITHOUT +# busting the expensive torch/model layers below in runtime. Default 0 keeps +# local builds cache-friendly. Runs on the base system python (dnf's own 3.9); +# it never touches /opt/venv, so no python symlink fixup is needed. +ARG CACHE_REFRESH=0 +RUN echo "CVE patch refresh token: ${CACHE_REFRESH}" \ + && dnf upgrade -y --security --releasever=latest --setopt=install_weak_deps=False \ + && dnf clean all && rm -rf /var/cache/dnf /tmp/* + +ARG DLC_MAJOR_VERSION +ARG DLC_MINOR_VERSION +LABEL maintainer="Amazon AI" +LABEL dlc_major_version="${DLC_MAJOR_VERSION}" +LABEL dlc_minor_version="${DLC_MINOR_VERSION}" +LABEL com.amazon.dlc.variant="whisperx-ec2-amzn2023" +EXPOSE 8000 + +COPY scripts/docker/whisperx/dockerd_entrypoint.sh /usr/local/bin/dockerd_entrypoint.sh +RUN chmod +x /usr/local/bin/dockerd_entrypoint.sh + +ENTRYPOINT ["/usr/local/bin/dockerd_entrypoint.sh"] + +# ============================================================================= +# TARGET: whisperx-sagemaker-amzn2023 +# ============================================================================= +FROM runtime AS whisperx-sagemaker-amzn2023 + +# ---- OS CVE patch (cache-bustable) ---- (see ec2 target above for rationale) +ARG CACHE_REFRESH=0 +RUN echo "CVE patch refresh token: ${CACHE_REFRESH}" \ + && dnf upgrade -y --security --releasever=latest --setopt=install_weak_deps=False \ + && dnf clean all && rm -rf /var/cache/dnf /tmp/* + +ARG DLC_MAJOR_VERSION +ARG DLC_MINOR_VERSION +LABEL maintainer="Amazon AI" +LABEL dlc_major_version="${DLC_MAJOR_VERSION}" +LABEL dlc_minor_version="${DLC_MINOR_VERSION}" +LABEL com.amazon.dlc.variant="whisperx-sagemaker-amzn2023" +EXPOSE 8080 + +COPY scripts/docker/whisperx/sagemaker_entrypoint.sh /usr/local/bin/sagemaker_entrypoint.sh +RUN chmod +x /usr/local/bin/sagemaker_entrypoint.sh + +ENTRYPOINT ["/usr/local/bin/sagemaker_entrypoint.sh"] \ No newline at end of file diff --git a/docker/whisperx/NOTICE b/docker/whisperx/NOTICE new file mode 100644 index 000000000000..891e25c4daaa --- /dev/null +++ b/docker/whisperx/NOTICE @@ -0,0 +1,34 @@ +This container image includes the pyannote.audio speaker diarization pipeline +"speaker-diarization-community-1" and its constituent speaker embedding model +"wespeaker-voxceleb-resnet34-LM", redistributed under the terms of the +Creative Commons Attribution 4.0 International License (CC BY 4.0). + +-------------------------------------------------------------------------- +Work: pyannote/speaker-diarization-community-1 +Authors: Hervé Bredin, Alexis Plaquet, and the pyannote.audio contributors +Source: https://huggingface.co/pyannote/speaker-diarization-community-1 +License: CC BY 4.0 — https://creativecommons.org/licenses/by/4.0/ +Cites: Plaquet & Bredin, "Powerset multi-class cross entropy loss for + neural speaker diarization", INTERSPEECH 2023; + Wang et al., "WeSpeaker: A research and production oriented speaker + embedding learning toolkit", ICASSP 2023; + Landini et al., "Bayesian HMM clustering of x-vector sequences + (VBx) in speaker diarization", Computer Speech & Language, 2022. +Changes: Model weights are redistributed unmodified; packaged into a + container image at build time. No modification to weights or + configuration files has been made. + +Work: pyannote/wespeaker-voxceleb-resnet34-LM +Authors: Hongji Wang, Chengdong Liang, Shuai Wang, Zhengyang Chen, + Binbin Zhang, Xu Xiang, Yanlei Deng, Yanmin Qian (WeSpeaker), + packaged for pyannote.audio by Hervé Bredin +Source: https://huggingface.co/pyannote/wespeaker-voxceleb-resnet34-LM +License: CC BY 4.0 — https://creativecommons.org/licenses/by/4.0/ +Changes: Redistributed unmodified. +-------------------------------------------------------------------------- + +DISCLAIMER OF WARRANTIES: The licensed material is provided "as-is" and +"as-available" by the Licensor. To the extent possible, the Licensor +offers no representations or warranties of any kind, express or implied, +statutory or otherwise. See CC BY 4.0 §5 for the full disclaimer: +https://creativecommons.org/licenses/by/4.0/legalcode.txt diff --git a/scripts/ci/build/whisperx/pre_build.sh b/scripts/ci/build/whisperx/pre_build.sh new file mode 100755 index 000000000000..5e9e320d968e --- /dev/null +++ b/scripts/ci/build/whisperx/pre_build.sh @@ -0,0 +1,33 @@ +#!/usr/bin/env bash +# Pre-build hook for WhisperX. +# Downloads the pyannote speaker-diarization weights tarball from the CI models +# bucket into the Docker build context so the Dockerfile can COPY it in. Mirrors +# the S3 step in the standalone whisperx-docker/scripts/build.sh. The CI runner +# has read access to dlc-cicd-models (same account), so no AWS creds enter the +# docker build itself. +# +# Usage: +# bash scripts/ci/build/whisperx/pre_build.sh --config-file +# +# Side effects: +# Places the tarball at docker/whisperx/pyannote-diarization.tar.gz +set -euo pipefail + +CONFIG_FILE="" +while [[ $# -gt 0 ]]; do + case "$1" in + --config-file) CONFIG_FILE="$2"; shift 2 ;; + *) echo "Unknown argument: $1" >&2; exit 1 ;; + esac +done + +[[ -n "$CONFIG_FILE" ]] || { echo "ERROR: --config-file is required" >&2; exit 1; } +[[ -f "$CONFIG_FILE" ]] || { echo "ERROR: Config file not found: $CONFIG_FILE" >&2; exit 1; } + +REPO_ROOT=$(pwd) +DEST="${REPO_ROOT}/docker/whisperx/pyannote-diarization.tar.gz" +PYANNOTE_S3_URI="${PYANNOTE_S3_URI:-s3://dlc-cicd-models/whisperx-models/speaker-diarization-community-1.tar.gz}" + +echo "Downloading pyannote weights: ${PYANNOTE_S3_URI}" +aws s3 cp "${PYANNOTE_S3_URI}" "${DEST}" +echo "pyannote weights ready: $(ls -la "${DEST}")" diff --git a/scripts/docker/telemetry/deep_learning_container.py b/scripts/docker/telemetry/deep_learning_container.py index 466911d6d51f..5554a59b873a 100755 --- a/scripts/docker/telemetry/deep_learning_container.py +++ b/scripts/docker/telemetry/deep_learning_container.py @@ -243,6 +243,7 @@ def parse_args(): "ray", "vllm_omni", "huggingface-vllm", + "whisperx", ], help="framework of container image.", required=True, diff --git a/scripts/docker/whisperx/cuda_compat.sh b/scripts/docker/whisperx/cuda_compat.sh new file mode 100755 index 000000000000..be0ad44a5751 --- /dev/null +++ b/scripts/docker/whisperx/cuda_compat.sh @@ -0,0 +1,61 @@ +#!/usr/bin/env bash +# CUDA forward-compatibility activation. Sourced by both entrypoints BEFORE they +# exec uvicorn. +# +# Why this exists: the image bakes torch's cu128 CUDA libs. On a GPU host whose +# NVIDIA driver is older than those libs need, this prepends the bundled +# cuda-compat userspace libcuda (/usr/local/cuda/compat) to LD_LIBRARY_PATH so +# CUDA still initialises — but ONLY when the host driver is actually older than +# the compat build, so newer-driver hosts (and CPU hosts with no driver) are +# untouched. Adapted from DLC scripts/docker/pytorch/{entrypoint,start_cuda_compat}.sh. +# +# IMPORTANT — SageMaker GPU endpoints need BOTH this AND an AMI pin: +# cuda-compat can only bridge from a base driver that itself meets the compat +# package's minimum. SageMaker's DEFAULT g4dn host AMI ships a driver too old +# even for that, so the container fails to start with a zero-log +# "CannotStartContainerError" REGARDLESS of this script. The endpoint-config +# MUST set InferenceAmiVersion to a CUDA-12.x AMI: +# "InferenceAmiVersion": "al2-ami-sagemaker-inference-gpu-3-1" # driver 550 +# With that AMI the host driver is 550.163.01; this script then activates +# compat (550 < 570.211.01) and CUDA 12.8 runs. VERIFIED on ml.g4dn.xlarge: +# endpoint reached InService and served GPU transcription + diarization. +# (CUDA 13.x images use al2023-ami-sagemaker-inference-gpu-4-1, driver 580.) +# See DLC test/test_utils/constants.py for the CUDA-major -> AMI mapping. +# +# This is intentionally guarded to never abort a caller running under +# `set -euo pipefail`: every probe that can legitimately fail (no GPU, no +# /proc/driver/nvidia) is `|| true`-guarded and the whole thing is a no-op when +# the compat lib is absent or the host driver is already new enough. + +_activate_cuda_forward_compat() { + local compat_file=/usr/local/cuda/compat/libcuda.so.1 + [ -f "$compat_file" ] || return 0 + + # Max driver the compat build speaks, parsed from e.g. + # libcuda.so.1 -> libcuda.so.570.211.01 => 570.211.01 + local compat_max + compat_max="$(readlink "$compat_file" | cut -d'.' -f3-)" + [ -n "$compat_max" ] || return 0 + + # Host driver: prefer /proc (present iff a GPU + kernel module exist), then + # nvidia-smi. On a CPU host neither exists -> empty -> skip (no compat needed). + local host_drv="" + host_drv="$(sed -n 's/^NVRM.*Kernel Module *\([0-9.]*\).*$/\1/p' /proc/driver/nvidia/version 2>/dev/null || true)" + if [ -z "$host_drv" ]; then + host_drv="$(nvidia-smi --query-gpu=driver_version --format=csv,noheader --id=0 2>/dev/null || true)" + fi + [ -n "$host_drv" ] || return 0 + + # Activate compat ONLY when host_drv < compat_max (strict). If the host driver + # is >= the compat build, the host driver is already sufficient — leave it. + local lowest + lowest="$(printf '%s\n%s\n' "$host_drv" "$compat_max" | sort -V | head -n1)" + if [ "$host_drv" = "$lowest" ] && [ "$host_drv" != "$compat_max" ]; then + export LD_LIBRARY_PATH="/usr/local/cuda/compat:${LD_LIBRARY_PATH:-}" + echo "INFO: host NVIDIA driver ${host_drv} < CUDA compat ${compat_max}; activated forward-compat libcuda (/usr/local/cuda/compat)" + else + echo "INFO: host NVIDIA driver ${host_drv} >= CUDA compat ${compat_max}; using host driver (no forward-compat needed)" + fi +} + +_activate_cuda_forward_compat diff --git a/scripts/docker/whisperx/dockerd_entrypoint.sh b/scripts/docker/whisperx/dockerd_entrypoint.sh new file mode 100755 index 000000000000..97b9a4c0a174 --- /dev/null +++ b/scripts/docker/whisperx/dockerd_entrypoint.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +# EC2 PID 1. Design §3: no tini, exec uvicorn directly so signals reach it. +set -euo pipefail + +# DLC telemetry: fire-and-forget IMDS ping at container start. Runs here (not +# only via bashrc) because the entrypoint exec's uvicorn directly, with no +# interactive/login shell to trigger the bashrc hook. Errors suppressed so +# telemetry never blocks or fails startup. +bash /usr/local/bin/bash_telemetry.sh >/dev/null 2>&1 || true + +# Activate CUDA forward-compat if the host driver is older than the baked CUDA +# libs need. No-op on new-driver and CPU hosts. Must run before uvicorn imports +# torch. See cuda_compat.sh. +# shellcheck source=/dev/null +. /opt/whisperx/cuda_compat.sh + +# Any extra flags the operator passes on `docker run` land in $@; forward them +# to uvicorn (e.g. --workers 2, --log-level debug). +exec uvicorn server:app \ + --host 0.0.0.0 \ + --port 8000 \ + --app-dir /opt/whisperx \ + "$@" diff --git a/scripts/docker/whisperx/sagemaker_entrypoint.sh b/scripts/docker/whisperx/sagemaker_entrypoint.sh new file mode 100755 index 000000000000..5ab12aedd001 --- /dev/null +++ b/scripts/docker/whisperx/sagemaker_entrypoint.sh @@ -0,0 +1,38 @@ +#!/usr/bin/env bash +# SageMaker PID 1. SageMaker invokes the container with `serve` as the first +# argument; anything else means the user is `docker run …`-ing this variant +# outside SageMaker, in which case we behave identically to the EC2 target. +set -euo pipefail + +# DLC telemetry: fire-and-forget IMDS ping at container start. Runs here (not +# only via bashrc) because the entrypoint exec's uvicorn directly, with no +# interactive/login shell to trigger the bashrc hook. Errors suppressed so +# telemetry never blocks or fails startup. +bash /usr/local/bin/bash_telemetry.sh >/dev/null 2>&1 || true + +# If the customer mounted a model tarball, point HF_HOME at it so any pre-baked +# Whisper / wav2vec2 caches inside the tarball are picked up. If /opt/ml/model +# is empty (SageMaker still mounts an empty dir) fall back to the image's cache. +if [ -d /opt/ml/model ] && [ -n "$(ls -A /opt/ml/model 2>/dev/null || true)" ]; then + export HF_HOME=/opt/ml/model + echo "INFO: /opt/ml/model is populated; using it as HF_HOME" +fi + +# Activate CUDA forward-compat if the host driver is older than the baked CUDA +# libs need (e.g. SageMaker ml.g4dn ships an older driver). No-op on new-driver +# and CPU hosts. Must run before uvicorn imports torch. See cuda_compat.sh. +# shellcheck source=/dev/null +. /opt/whisperx/cuda_compat.sh + +# SageMaker passes `serve` when it launches an inference endpoint. Any other +# invocation (including no args at all) is a local docker run — treat it the +# same, since the entrypoint has nothing else to do. +if [ "${1:-serve}" = "serve" ]; then + shift || true +fi + +exec uvicorn server:app \ + --host 0.0.0.0 \ + --port 8080 \ + --app-dir /opt/whisperx \ + "$@" diff --git a/scripts/docker/whisperx/server.py b/scripts/docker/whisperx/server.py new file mode 100644 index 000000000000..c2228837fefc --- /dev/null +++ b/scripts/docker/whisperx/server.py @@ -0,0 +1,735 @@ +"""WhisperX FastAPI server. + +Design ref: workspace/whisperx-docker/DESIGN.md §5. + +Four routes, all served in one process: + POST /v1/audio/transcriptions — primary, OpenAI-compatible + POST /invocations — alias for SageMaker + GET /ping — readiness health check + GET /v1/models — advertises the single served model id + +Extension fields on top of OpenAI's schema: `diarize`, `min_speakers`, +`max_speakers`. When `diarize=false` (default) output is byte-identical to +OpenAI's `verbose_json`. + +Subtitle line-formatting fields — `max_line_width`, `max_line_count`, +`highlight_words` — are WhisperX-CLI-named (not OpenAI) and apply ONLY to the +`srt` and `vtt` response formats; for json/text/verbose_json they are accepted +but ignored (matching WhisperX, whose non-subtitle writers ignore them). Setting +any of them routes srt/vtt through WhisperX's own SubtitlesWriter so output is +byte-identical to the WhisperX CLI, and forces word-level alignment internally +(the writer only wraps/highlights when segments carry word timings). + +All WhisperX model, decoding, and VAD parameters (beam size, temperature, +initial prompt, VAD thresholds, task, ...) are fixed at container launch via +WHISPERX_* env vars — see the "Launch-time WhisperX configuration" block below. +The HTTP API is deliberately lean: `temperature` and `prompt` are launch-only +knobs, not per-request fields. + +The Whisper model is one-per-container, chosen at launch via +WHISPERX_DEFAULT_MODEL; the request `model` field is accepted but ignored (an +OpenAI-compat no-op), so a single model stays resident rather than being loaded +per request. +""" + +from __future__ import annotations + +import functools +import io +import os +import tempfile +import threading +from collections import OrderedDict +from contextlib import asynccontextmanager +from pathlib import Path +from typing import Any + +import anyio +import torch +import whisperx +from fastapi import FastAPI, File, Form, HTTPException, UploadFile +from fastapi.responses import JSONResponse, PlainTextResponse +from whisperx.diarize import DiarizationPipeline +from whisperx.utils import WriteSRT, WriteVTT + + +# --------------------------------------------------------------------------- +# Launch-time configuration helpers +# --------------------------------------------------------------------------- +# All WhisperX decoding/VAD/model knobs are fixed at container launch via env +# vars and read once (below) into immutable module globals. Grouped by where +# whisperx.load_model consumes them: +# +# asr_options (decoding) — WHISPERX_BEAM_SIZE, WHISPERX_BEST_OF, +# WHISPERX_PATIENCE, WHISPERX_LENGTH_PENALTY, WHISPERX_TEMPERATURE, +# WHISPERX_TEMPERATURE_INCREMENT_ON_FALLBACK, +# WHISPERX_COMPRESSION_RATIO_THRESHOLD, WHISPERX_LOGPROB_THRESHOLD, +# WHISPERX_NO_SPEECH_THRESHOLD, WHISPERX_CONDITION_ON_PREVIOUS_TEXT, +# WHISPERX_INITIAL_PROMPT, WHISPERX_HOTWORDS, WHISPERX_SUPPRESS_TOKENS, +# WHISPERX_SUPPRESS_NUMERALS +# vad_options — WHISPERX_CHUNK_SIZE, WHISPERX_VAD_ONSET, WHISPERX_VAD_OFFSET +# pipeline/model — WHISPERX_VAD_METHOD, WHISPERX_TASK, WHISPERX_ALIGN_MODEL, +# WHISPERX_DEFAULT_MODEL, WHISPERX_COMPUTE_TYPE, WHISPERX_BATCH_SIZE, +# WHISPERX_SERVED_MODEL_NAME, ... +def _env_bool(name: str, default: bool = False) -> bool: + val = os.environ.get(name) + if val is None or not val.strip(): + return default + return val.strip().lower() in {"1", "true", "yes", "on"} + + +def _build_asr_options() -> dict[str, Any]: + """Assemble faster-whisper decoding options from WHISPERX_* env. + + Mirrors the whisperX CLI defaults; passed to load_model(asr_options=...) at + launch. Pure (env -> dict) so it is unit-testable without a GPU/container. + """ + base = float(os.environ.get("WHISPERX_TEMPERATURE") or 0.0) + inc_raw = os.environ.get("WHISPERX_TEMPERATURE_INCREMENT_ON_FALLBACK", "0.2") + inc = float(inc_raw) if inc_raw.strip() else 0.0 + if inc > 0: + # Match whisper's np.arange(base, 1.0 + 1e-6, inc) fallback schedule. + steps = int((1.0 + 1e-6 - base) / inc) + 1 + temperatures = tuple(round(base + i * inc, 4) for i in range(max(steps, 1))) + else: + temperatures = (base,) + suppress = os.environ.get("WHISPERX_SUPPRESS_TOKENS", "-1") + return { + "beam_size": int(os.environ.get("WHISPERX_BEAM_SIZE", "5")), + "best_of": int(os.environ.get("WHISPERX_BEST_OF", "5")), + "patience": float(os.environ.get("WHISPERX_PATIENCE", "1.0")), + "length_penalty": float(os.environ.get("WHISPERX_LENGTH_PENALTY", "1.0")), + "temperatures": temperatures, + "compression_ratio_threshold": float( + os.environ.get("WHISPERX_COMPRESSION_RATIO_THRESHOLD", "2.4") + ), + "log_prob_threshold": float(os.environ.get("WHISPERX_LOGPROB_THRESHOLD", "-1.0")), + "no_speech_threshold": float(os.environ.get("WHISPERX_NO_SPEECH_THRESHOLD", "0.6")), + "condition_on_previous_text": _env_bool("WHISPERX_CONDITION_ON_PREVIOUS_TEXT"), + "initial_prompt": os.environ.get("WHISPERX_INITIAL_PROMPT") or None, + "hotwords": os.environ.get("WHISPERX_HOTWORDS") or None, + "suppress_tokens": [int(x) for x in suppress.split(",") if x.strip()], + "suppress_numerals": _env_bool("WHISPERX_SUPPRESS_NUMERALS"), + } + + +def _build_vad_options() -> dict[str, Any]: + """Assemble VAD options from WHISPERX_* env; passed to load_model(vad_options=...).""" + return { + "chunk_size": int(os.environ.get("WHISPERX_CHUNK_SIZE", "30")), + "vad_onset": float(os.environ.get("WHISPERX_VAD_ONSET", "0.5")), + "vad_offset": float(os.environ.get("WHISPERX_VAD_OFFSET", "0.363")), + } + + +# --------------------------------------------------------------------------- +# Model caches +# --------------------------------------------------------------------------- +# One Whisper model per container, pinned at launch (DEFAULT_MODEL) and warmed +# in the lifespan hook. The request `model` field is ignored (OpenAI-compat +# no-op), so a single model stays resident rather than a per-name LRU. +_WHISPER_MODEL: Any = None + +# wav2vec2 align models — keyed by language code. Orthogonal to the Whisper +# model (a request's detected language selects the aligner), so this stays an +# LRU even though only one Whisper model is served. +_ALIGN_LRU: "OrderedDict[str, tuple[Any, dict[str, Any]]]" = OrderedDict() +_ALIGN_LRU_MAX = int(os.environ.get("WHISPERX_ALIGN_LRU_SIZE", "3")) + +# _transcribe runs in an anyio worker thread. _WHISPER_LOCK serializes the +# one-time model load (defense-in-depth if MAX_CONCURRENT_REQUESTS>1 ever races +# concurrent cold misses, so only the first thread loads); _ALIGN_LOCK does the +# same for the align LRU. Inference itself stays outside these locks (it runs in +# _transcribe after the getters return). +_WHISPER_LOCK = threading.Lock() +_ALIGN_LOCK = threading.Lock() + +# Fixed models loaded at startup. +_DIARIZE_PIPELINE: Any = None + +# Process readiness. Starts True; the inference path flips it False on a fatal +# CUDA/device fault (see _is_fatal_cuda_error) so /ping can report 503. +_HEALTHY = True + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" +COMPUTE_TYPE = os.environ.get("WHISPERX_COMPUTE_TYPE", "float16" if DEVICE == "cuda" else "int8") +# The Whisper model pinned at container launch and warmed in the lifespan hook. +DEFAULT_MODEL = os.environ.get("WHISPERX_DEFAULT_MODEL", "large-v2") +# Optional client-facing alias for the pinned model (e.g. "whisper-1" for an +# OpenAI-SDK drop-in). Used ONLY by /v1/models to advertise the served id; the +# request `model` field is ignored, so no request-time resolution happens. +SERVED_MODEL_NAME = os.environ.get("WHISPERX_SERVED_MODEL_NAME") +DEFAULT_BATCH_SIZE = int(os.environ.get("WHISPERX_BATCH_SIZE", "16")) +DIARIZE_MODEL_PATH = os.environ.get( + "WHISPERX_DIARIZE_MODEL_PATH", + "/opt/models/pyannote/speaker-diarization-community-1", +) +# VAD backend and transcription task are load_model args (not per-request). An +# optional align model override lets operators pin a specific wav2vec2 model. +VAD_METHOD = os.environ.get("WHISPERX_VAD_METHOD", "pyannote") +TASK = os.environ.get("WHISPERX_TASK", "transcribe") +ALIGN_MODEL = os.environ.get("WHISPERX_ALIGN_MODEL") or None +# Decoding + VAD options are global launch config: read env once here so the +# model loaded via _get_whisper uses one immutable configuration. +ASR_OPTIONS = _build_asr_options() +VAD_OPTIONS = _build_vad_options() + +# --------------------------------------------------------------------------- +# Request admission + inference concurrency +# --------------------------------------------------------------------------- +# GPU inference is serialized to MAX_CONCURRENT_REQUESTS (default 1): WhisperX's +# FasterWhisperPipeline mutates self.options/self.tokenizer mid-call and the +# shared diarization pipeline is not documented thread-safe, so running two +# transcriptions on one instance concurrently is unsafe. _INFERENCE_LIMITER caps +# in-flight inference; MAX_QUEUE bounds how many requests may wait behind it +# before new ones are shed with 503; MAX_UPLOAD_BYTES bounds per-request +# memory/temp-file retention. +MAX_CONCURRENT_REQUESTS = int(os.environ.get("WHISPERX_MAX_CONCURRENT_REQUESTS", "1")) +_INFERENCE_LIMITER = anyio.CapacityLimiter(MAX_CONCURRENT_REQUESTS) +MAX_QUEUE = int(os.environ.get("WHISPERX_MAX_QUEUE", "2")) +MAX_UPLOAD_BYTES = int(os.environ.get("WHISPERX_MAX_UPLOAD_BYTES", str(1024 * 1024 * 1024))) + + +def _lru_touch(cache: "OrderedDict[str, Any]", key: str, max_size: int) -> None: + cache.move_to_end(key) + while len(cache) > max_size: + cache.popitem(last=False) + + +def _get_whisper(model_name: str) -> Any: + # One model per container: model_name is always DEFAULT_MODEL, so a single + # resident model is cached. _WHISPER_LOCK holds across the check-load-store + # so that if MAX_CONCURRENT_REQUESTS>1 ever races two cold misses, only the + # first loads (dedupes the ~30 s load); a warm hit takes the lock briefly. + # Decoding/VAD/task config is global launch config (ASR_OPTIONS/VAD_OPTIONS). + global _WHISPER_MODEL + with _WHISPER_LOCK: + if _WHISPER_MODEL is None: + _WHISPER_MODEL = whisperx.load_model( + model_name, + device=DEVICE, + compute_type=COMPUTE_TYPE, + asr_options=ASR_OPTIONS, + vad_method=VAD_METHOD, + vad_options=VAD_OPTIONS, + task=TASK, + ) + return _WHISPER_MODEL + + +def _get_align(language: str) -> tuple[Any, dict[str, Any]]: + # Same single-critical-section rationale as _get_whisper: serialize the + # keyed-by-language check-load-store-touch so concurrent worker threads + # dedupe the load and never race move_to_end against an eviction. + with _ALIGN_LOCK: + if language not in _ALIGN_LRU: + _ALIGN_LRU[language] = whisperx.load_align_model( + language_code=language, device=DEVICE, model_name=ALIGN_MODEL + ) + _lru_touch(_ALIGN_LRU, language, _ALIGN_LRU_MAX) + return _ALIGN_LRU[language] + + +# --------------------------------------------------------------------------- +# App + lifespan +# --------------------------------------------------------------------------- +@asynccontextmanager +async def lifespan(_: FastAPI): + """Load fixed models before uvicorn binds so /ping stays shallow but honest. + + Design §6 decision 8: `/ping` reachable ⇒ the diarization pipeline and the + default Whisper model are resident. VAD uses whisperx's default (pyannote), + whose segmentation model ships inside the whisperx wheel. + """ + global _DIARIZE_PIPELINE, _HEALTHY + # Warm the default Whisper model so the first request isn't a cold ~30 s + # download-and-load. + _get_whisper(DEFAULT_MODEL) + + # Diarization: load the pyannote pipeline from the fixed local path the + # image baked at build time (COPYed from S3). Explicit local path ⇒ no HF + # token and no network needed. Absence is not fatal — non-diarized requests + # still work. This path is independent of HF_HOME, so the SageMaker + # entrypoint repointing HF_HOME (for user-mounted Whisper models) does not + # hide it. + if os.path.isdir(DIARIZE_MODEL_PATH): + try: + _DIARIZE_PIPELINE = DiarizationPipeline( + model_name=DIARIZE_MODEL_PATH, + device=DEVICE, + ) + except Exception as exc: # noqa: BLE001 — best-effort startup + print(f"WARN: diarization pipeline failed to load from {DIARIZE_MODEL_PATH}: {exc}") + _DIARIZE_PIPELINE = None + else: + print(f"WARN: diarization model dir {DIARIZE_MODEL_PATH} missing; diarization disabled") + _DIARIZE_PIPELINE = None + + # Warmup completed without a fatal load error; confirm readiness. (If + # _get_whisper raised above, we never reach here and uvicorn won't bind.) + _HEALTHY = True + yield + + +app = FastAPI(title="WhisperX DLC", lifespan=lifespan) + + +# --------------------------------------------------------------------------- +# Formatting helpers +# --------------------------------------------------------------------------- +def _seconds_to_srt_ts(t: float) -> str: + h, r = divmod(t, 3600) + m, s = divmod(r, 60) + ms = int((s - int(s)) * 1000) + return f"{int(h):02d}:{int(m):02d}:{int(s):02d},{ms:03d}" + + +def _seconds_to_vtt_ts(t: float) -> str: + return _seconds_to_srt_ts(t).replace(",", ".") + + +def _to_srt(segments: list[dict[str, Any]]) -> str: + lines: list[str] = [] + for i, seg in enumerate(segments, start=1): + lines.append(str(i)) + lines.append(f"{_seconds_to_srt_ts(seg['start'])} --> {_seconds_to_srt_ts(seg['end'])}") + lines.append(seg.get("text", "").strip()) + lines.append("") + return "\n".join(lines) + + +def _to_vtt(segments: list[dict[str, Any]]) -> str: + lines: list[str] = ["WEBVTT", ""] + for seg in segments: + lines.append(f"{_seconds_to_vtt_ts(seg['start'])} --> {_seconds_to_vtt_ts(seg['end'])}") + lines.append(seg.get("text", "").strip()) + lines.append("") + return "\n".join(lines) + + +def _render_subtitle( + result: dict[str, Any], + fmt: str, + max_line_width: int | None, + max_line_count: int | None, + highlight_words: bool, +) -> str: + """Render srt/vtt via WhisperX's own SubtitlesWriter (byte-identical to the CLI). + + Reuses whisperx.utils.WriteSRT / WriteVTT so line wrapping and word + highlighting exactly match the WhisperX CLI. output_dir is only consumed by + the writer's __call__ (file-path mode), never by write_result, so "" is safe. + All three option keys are mandatory: SubtitlesWriter.iterate_result reads + each by key and KeyErrors otherwise. `result` must carry "language" (checked + against LANGUAGES_WITHOUT_SPACES) and "segments"; pass the full _transcribe + result, not just segments. When segments lack "words" (alignment didn't run + or degraded for an unsupported language) the writer falls back to + segment-level cues automatically, so the knobs simply no-op. + """ + writer = (WriteSRT if fmt == "srt" else WriteVTT)(output_dir="") + options = { + "max_line_width": max_line_width, + "max_line_count": max_line_count, + "highlight_words": highlight_words, + } + buf = io.StringIO() + writer.write_result(result, file=buf, options=options) + return buf.getvalue() + + +def _flatten_words(segments: list[dict[str, Any]]) -> list[dict[str, Any]]: + out: list[dict[str, Any]] = [] + for seg in segments: + for w in seg.get("words", []) or []: + entry: dict[str, Any] = { + "word": w.get("word", ""), + "start": w.get("start"), + "end": w.get("end"), + } + if "speaker" in w: + entry["speaker"] = w["speaker"] + out.append(entry) + return out + + +# --------------------------------------------------------------------------- +# Core pipeline +# --------------------------------------------------------------------------- +# Fatal CUDA/device error substrings — a RuntimeError whose message contains any +# of these means the process's GPU context is corrupted and every subsequent +# request will fail, so readiness flips. Transient faults (CUDA OOM) are +# deliberately excluded: they don't poison the context. +_FATAL_CUDA_MARKERS = ( + "cuda error", + "device-side assert", + "illegal memory access", + "misaligned address", + "unspecified launch failure", +) + + +def _is_fatal_cuda_error(exc: Exception) -> bool: + """True for unrecoverable CUDA/device faults that corrupt the GPU context. + + Only RuntimeErrors carrying a known-fatal marker qualify. A plain CUDA OOM + (torch.cuda.OutOfMemoryError / "out of memory") is transient — the context + survives and later requests can succeed — so it is NOT treated as fatal. + """ + if not isinstance(exc, RuntimeError): + return False + msg = str(exc).lower() + return any(marker in msg for marker in _FATAL_CUDA_MARKERS) + + +def _transcribe( + audio_path: str, + language: str | None, + want_words: bool, + diarize: bool, + min_speakers: int | None, + max_speakers: int | None, +) -> dict[str, Any]: + audio = whisperx.load_audio(audio_path) + + model = _get_whisper(DEFAULT_MODEL) + # Decoding params (temperature/prompt/beam/...) are baked into the model at + # load time via ASR_OPTIONS. FasterWhisperPipeline.transcribe only accepts + # pipeline-level args; passing decoding kwargs here raises TypeError. `task` + # is passed per-call because load_model discards its task= when no language + # is pinned at load time, so it must reach transcribe() to take effect. + transcribe_kwargs: dict[str, Any] = { + "batch_size": DEFAULT_BATCH_SIZE, + "chunk_size": VAD_OPTIONS["chunk_size"], + "task": TASK, + } + if language: + transcribe_kwargs["language"] = language + + try: + result = model.transcribe(audio, **transcribe_kwargs) + detected_language = result.get("language", language or "") + + # Alignment (word-level timestamps). Required if the caller asked for + # word granularity OR diarization (word-level speaker assignment needs + # word timing). + if want_words or diarize: + try: + align_model, align_metadata = _get_align(detected_language) + result = whisperx.align( + result["segments"], + align_model, + align_metadata, + audio, + DEVICE, + return_char_alignments=False, + ) + except (ValueError, NotImplementedError, KeyError) as exc: + # These are the only exceptions whisperX raises when a language + # genuinely has no wav2vec2 aligner (no default align-model / + # unsupported model type / missing metadata key). Only these + # degrade gracefully; any other error (CUDA OOM RuntimeError, + # network failure fetching the aligner, ...) propagates so it + # surfaces as a real 500 instead of a silently-degraded 200. + print(f"WARN: alignment failed for language={detected_language}: {exc}") + degraded_from_diarize = diarize + want_words = False + diarize = False + result = {"segments": result["segments"]} + # Diarization needs word timing, so degradation makes the + # requested speakers impossible to produce. Rather than return + # 200 with no speakers, tell the caller explicitly. A best-effort + # want_words request keeps the silent (logged) fallback above. + if degraded_from_diarize: + raise HTTPException( + status_code=422, + detail=( + f"diarization was requested but alignment is unavailable for " + f"language '{detected_language}' (no wav2vec2 aligner); retry " + f"without diarize or with a supported language" + ), + ) from exc + + speakers: list[str] = [] + if diarize: + if _DIARIZE_PIPELINE is None: + raise HTTPException( + status_code=400, + detail="diarization requested but pyannote pipeline is not available", + ) + diar_kwargs: dict[str, Any] = {} + if min_speakers is not None: + diar_kwargs["min_speakers"] = min_speakers + if max_speakers is not None: + diar_kwargs["max_speakers"] = max_speakers + diarize_segments = _DIARIZE_PIPELINE(audio, **diar_kwargs) + result = whisperx.assign_word_speakers(diarize_segments, result) + # Collect unique speaker labels in first-seen order. + seen: set[str] = set() + for seg in result.get("segments", []): + spk = seg.get("speaker") + if spk and spk not in seen: + seen.add(spk) + speakers.append(spk) + except Exception as exc: + # A fatal CUDA/device fault corrupts this process's GPU context, so every + # later request on this instance will fail too. Flip readiness (=> /ping + # 503) so orchestration can drain/replace the host, then re-raise so the + # current request still errors. Transient failures — CUDA OOM, the + # unsupported-language degrade above, and the 422/400 HTTPExceptions — are + # NOT fatal and simply propagate without flipping health. + global _HEALTHY + if _is_fatal_cuda_error(exc): + _HEALTHY = False + raise + + segments: list[dict[str, Any]] = result.get("segments", []) + text = " ".join(seg.get("text", "").strip() for seg in segments).strip() + + return { + "task": TASK, + "language": detected_language, + "duration": float(len(audio)) / 16000.0 if hasattr(audio, "__len__") else 0.0, + "text": text, + "segments": segments, + "words": _flatten_words(segments) if want_words else None, + "speakers": speakers if diarize else None, + } + + +def _format_response( + result: dict[str, Any], + response_format: str, + want_words: bool, + diarize: bool, + max_line_width: int | None = None, + max_line_count: int | None = None, + highlight_words: bool = False, +): + segments = result["segments"] + # A subtitle knob being set routes srt/vtt through WhisperX's SubtitlesWriter + # (byte-identical to the CLI). With no knob set we keep the legacy + # _to_srt/_to_vtt so default srt/vtt output is unchanged. The knobs are + # subtitle-only: json/text/verbose_json ignore them (mirroring WhisperX). + subtitle_knob = max_line_width is not None or max_line_count is not None or highlight_words + if response_format == "text": + return PlainTextResponse(result["text"]) + if response_format == "srt": + if subtitle_knob: + body = _render_subtitle(result, "srt", max_line_width, max_line_count, highlight_words) + else: + body = _to_srt(segments) + return PlainTextResponse(body, media_type="application/x-subrip") + if response_format == "vtt": + if subtitle_knob: + body = _render_subtitle(result, "vtt", max_line_width, max_line_count, highlight_words) + else: + body = _to_vtt(segments) + return PlainTextResponse(body, media_type="text/vtt") + if response_format == "json": + return JSONResponse({"text": result["text"]}) + + # verbose_json — full payload, trimmed to what the caller asked for. + payload: dict[str, Any] = { + "task": result["task"], + "language": result["language"], + "duration": result["duration"], + "text": result["text"], + "segments": segments, + } + if want_words and result.get("words") is not None: + payload["words"] = result["words"] + if diarize and result.get("speakers"): + payload["speakers"] = result["speakers"] + return JSONResponse(payload) + + +# --------------------------------------------------------------------------- +# Routes +# --------------------------------------------------------------------------- +@app.get("/ping") +def ping(): + # Readiness reflects fatal CUDA/model state observed by the inference path. + # It is PASSIVE: a GPU that dies while the server is idle won't be caught + # here (no request has touched the device) — that's host-monitoring's job. + # Once the inference path hits a fatal CUDA fault, _HEALTHY flips and every + # /ping returns 503 so orchestration can drain/replace the instance. + if _HEALTHY: + return JSONResponse({"status": "ok"}) + return JSONResponse({"status": "unavailable"}, status_code=503) + + +@app.get("/v1/models") +def list_models() -> dict[str, Any]: + """Advertise the single served model id for OpenAI-client compatibility. + + The Whisper model is one-per-container and the request `model` field is + ignored, so this simply reports the id a client should display/send: the + SERVED_MODEL_NAME alias if configured, else the pinned DEFAULT_MODEL. + """ + return { + "object": "list", + "data": [ + { + "id": SERVED_MODEL_NAME or DEFAULT_MODEL, + "object": "model", + "owned_by": "openai-whisper", + } + ], + } + + +async def _handle_transcription( + file: UploadFile, + language: str | None, + response_format: str, + timestamp_granularities: list[str] | None, + diarize: bool, + min_speakers: int | None, + max_speakers: int | None, + max_line_width: int | None = None, + max_line_count: int | None = None, + highlight_words: bool = False, +): + valid_formats = {"json", "text", "srt", "vtt", "verbose_json"} + if response_format not in valid_formats: + raise HTTPException(400, f"response_format must be one of {sorted(valid_formats)}") + + # Validate the extension params at the boundary so a bad request fails fast + # with a 422 instead of surfacing deep in the diarization/subtitle path. + if min_speakers is not None and min_speakers < 1: + raise HTTPException(422, "min_speakers must be >= 1") + if max_speakers is not None and max_speakers < 1: + raise HTTPException(422, "max_speakers must be >= 1") + if min_speakers is not None and max_speakers is not None and min_speakers > max_speakers: + raise HTTPException(422, "min_speakers must be <= max_speakers") + if max_line_width is not None and max_line_width < 0: + raise HTTPException(422, "max_line_width must be >= 0") + if max_line_count is not None and max_line_count < 0: + raise HTTPException(422, "max_line_count must be >= 0") + + granularities = timestamp_granularities or ["segment"] + want_words = "word" in granularities + + # Subtitle line-formatting applies only to srt/vtt. WhisperX's + # SubtitlesWriter only wraps/highlights when segments carry word timings, so + # a subtitle knob forces word-level alignment internally (independent of the + # user-facing timestamp_granularities). Leave want_words as derived above and + # only OR in this internal flag for the _transcribe call. + subtitle_formatting = response_format in {"srt", "vtt"} and ( + max_line_width is not None or max_line_count is not None or highlight_words + ) + + # Admission control: when the inference queue is already full, shed load with + # 503 before the full-size `await file.read()` below, so an overload doesn't + # materialize more large byte buffers or grow the worker queue. This is a soft + # bound, not a hard cap: Starlette's File(...) has already spooled the raw + # multipart body by the time this runs, and tasks_waiting only counts requests + # that have reached the limiter, so a tight concurrent burst can admit a few + # past this check. Excess admitted requests then wait on the event loop. + if _INFERENCE_LIMITER.statistics().tasks_waiting >= MAX_QUEUE: + raise HTTPException(503, "server busy: inference queue full") + + audio_bytes = await file.read() + if not audio_bytes: + raise HTTPException(400, "empty audio upload") + # Bound per-request memory + temp-file retention. Streaming the upload to + # disk with an incremental size check is a future improvement; this rejects + # outsized uploads after a single read. + if len(audio_bytes) > MAX_UPLOAD_BYTES: + raise HTTPException(413, f"upload exceeds {MAX_UPLOAD_BYTES} bytes") + + # WhisperX + faster-whisper read from a path (ffmpeg is invoked as a + # subprocess). Spool the upload to a NamedTemporaryFile so the ffmpeg + # child can read it. + suffix = Path(file.filename or "audio.wav").suffix or ".wav" + with tempfile.NamedTemporaryFile(suffix=suffix, delete=True) as tmp: + tmp.write(audio_bytes) + tmp.flush() + # _transcribe is fully synchronous and long-running (GPU inference plus + # an ffmpeg subprocess). Running it inline would block uvicorn's single + # event loop for the whole transcription, starving every other + # coroutine — including GET /ping — and risking SageMaker health-check + # failures. Offload it to a worker thread so the loop stays responsive. + # The limiter serializes the whole _transcribe (transcription + align + + # diarization): WhisperX's FasterWhisperPipeline mutates + # self.options/self.tokenizer mid-call and the shared diarization + # pipeline is not documented thread-safe, so concurrent inference on one + # instance is unsafe. MAX_CONCURRENT_REQUESTS (default 1) caps it; excess + # requests wait on the event loop (not in the worker pool) so /ping stays + # responsive. + result = await anyio.to_thread.run_sync( + functools.partial( + _transcribe, + audio_path=tmp.name, + language=language, + want_words=want_words or subtitle_formatting, + diarize=diarize, + min_speakers=min_speakers, + max_speakers=max_speakers, + ), + limiter=_INFERENCE_LIMITER, + ) + + return _format_response( + result, + response_format, + want_words, + diarize, + max_line_width, + max_line_count, + highlight_words, + ) + + +@app.post("/v1/audio/transcriptions") +async def transcribe( + file: UploadFile = File(...), + language: str | None = Form(None), + response_format: str = Form("json"), + timestamp_granularities: list[str] | None = Form(None, alias="timestamp_granularities[]"), + diarize: bool = Form(False), + min_speakers: int | None = Form(None), + max_speakers: int | None = Form(None), + max_line_width: int | None = Form(None), + max_line_count: int | None = Form(None), + highlight_words: bool = Form(False), +): + return await _handle_transcription( + file=file, + language=language, + response_format=response_format, + timestamp_granularities=timestamp_granularities, + diarize=diarize, + min_speakers=min_speakers, + max_speakers=max_speakers, + max_line_width=max_line_width, + max_line_count=max_line_count, + highlight_words=highlight_words, + ) + + +@app.post("/invocations") +async def invocations( + file: UploadFile = File(...), + language: str | None = Form(None), + response_format: str = Form("json"), + timestamp_granularities: list[str] | None = Form(None, alias="timestamp_granularities[]"), + diarize: bool = Form(False), + min_speakers: int | None = Form(None), + max_speakers: int | None = Form(None), + max_line_width: int | None = Form(None), + max_line_count: int | None = Form(None), + highlight_words: bool = Form(False), +): + return await _handle_transcription( + file=file, + language=language, + response_format=response_format, + timestamp_granularities=timestamp_granularities, + diarize=diarize, + min_speakers=min_speakers, + max_speakers=max_speakers, + max_line_width=max_line_width, + max_line_count=max_line_count, + highlight_words=highlight_words, + ) diff --git a/test/security/data/ecr_scan_allowlist/whisperx/framework_allowlist.json b/test/security/data/ecr_scan_allowlist/whisperx/framework_allowlist.json new file mode 100644 index 000000000000..66e3a0f8e2eb --- /dev/null +++ b/test/security/data/ecr_scan_allowlist/whisperx/framework_allowlist.json @@ -0,0 +1,52 @@ +[ + { + "vulnerability_id": "CVE-2026-4372", + "reason": "transformers is a transitive dependency pulled by whisperx 3.8.6, which constrains transformers<5; the fix requires transformers>=5.3.0 (a major bump) and cannot be taken without breaking the whisperx contract. Not on the container's inference path (the Trainer class is unused). Tracked for a whisperx version that supports transformers 5.x.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2026-1839", + "reason": "transformers is a transitive dependency pulled by whisperx 3.8.6, which constrains transformers<5; the fix requires transformers>=5.0.0 (a major bump) and cannot be taken without breaking the whisperx contract. The vulnerable Trainer._load_rng_state path is not exercised by inference serving. Tracked for a whisperx version that supports transformers 5.x.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2026-5241", + "reason": "transformers is a transitive dependency pulled by whisperx 3.8.6, which constrains transformers<5; the fix requires transformers>=5.5.0 (a major bump) and cannot be taken without breaking the whisperx contract. The vulnerable code path is not exercised by the ASR inference server. Tracked for a whisperx version that supports transformers 5.x.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2026-24747", + "reason": "torch is pinned to 2.8.0 by whisperx 3.8.6 (torch~=2.8.0, via torchvision~=0.23.0); the fix requires torch>=2.10.0 and cannot be taken without breaking the whisperx contract. torch.load defaults to weights_only=True on 2.8, and model checkpoints are baked or resolved from HuggingFace by default; an operator-supplied model directory is a deployment trust boundary. Tracked for a whisperx version that raises its torch cap.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2025-55552", + "reason": "torch is pinned to 2.8.0 by whisperx 3.8.6 (torch~=2.8.0); the fix requires torch>=2.9.0 and cannot be taken without breaking the whisperx contract. The reported torch.rot90/torch.randn_like interaction is not on the ASR inference path. Tracked for a whisperx version that raises its torch cap.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2025-55551", + "reason": "torch is pinned to 2.8.0 by whisperx 3.8.6 (torch~=2.8.0); the fix requires torch>=2.9.0 and cannot be taken without breaking the whisperx contract. The reported torch.linalg.lu slice denial-of-service is not on the ASR inference path (transcription/alignment/diarization do not call torch.linalg.lu). Tracked for a whisperx version that raises its torch cap.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2026-54283", + "reason": "starlette is a transitive dependency of fastapi 0.121.0, which pins starlette<0.50.0; the fix requires starlette>=1.3.1 and cannot be taken without a major fastapi bump that breaks the serving stack. Tracked for a fastapi release that allows starlette 1.x.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2026-48818", + "reason": "starlette is a transitive dependency of fastapi 0.121.0, which pins starlette<0.50.0; the fix requires starlette>=1.1.0 (major bump). The vulnerability is a Windows-only StaticFiles SSRF and is not reachable on this Linux container, which serves no StaticFiles. Tracked for a fastapi release that allows starlette 1.x.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2025-23339", + "reason": "cuda-toolkit-config-common comes from the nvidia/cuda:12.8.2-*-amzn2023 base image; the fix requires the CUDA 13 toolkit (>=13.0.48), which conflicts with the image's CUDA 12.8 contract (torch cu128 wheels). The vulnerable cuobjdump tool is a developer utility not invoked at inference time. Tracked against a future CUDA 13 base migration.", + "review_by": "2026-10-13" + }, + { + "vulnerability_id": "CVE-2025-23308", + "reason": "cuda-toolkit-config-common comes from the nvidia/cuda:12.8.2-*-amzn2023 base image; the fix requires the CUDA 13 toolkit (>=13.0.48), which conflicts with the image's CUDA 12.8 contract (torch cu128 wheels). The vulnerable nvdisasm tool is a developer utility not invoked at inference time. Tracked against a future CUDA 13 base migration.", + "review_by": "2026-10-13" + } +] diff --git a/test/whisperx/ec2/common.py b/test/whisperx/ec2/common.py new file mode 100644 index 000000000000..6144362d4b43 --- /dev/null +++ b/test/whisperx/ec2/common.py @@ -0,0 +1,265 @@ +# Copyright 2018-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""Shared constants, helpers, fixtures, and HTTP utilities for WhisperX EC2 tests. + +The GPU test module imports from here, setting only DEVICE and DOCKER_RUN_FLAGS. + +Like the Ray EC2 suite, these tests run the container locally via `docker run`, +hit the WhisperX FastAPI server on port 8000, and validate the OpenAI-compatible +transcription contract. Unlike Ray, no model tarball is mounted: the Whisper +default model and wav2vec2 aligners download lazily from HuggingFace at runtime +and the pyannote diarization models are baked into the image, so the container +is started with just the image URI and the entrypoint runs uvicorn itself. +""" + +import logging +import os +import subprocess +import time +from contextlib import contextmanager + +import pytest +import requests + +LOGGER = logging.getLogger(__name__) +LOGGER.setLevel(logging.INFO) + +# The WhisperX EC2 target EXPOSEs 8000 and its entrypoint runs +# `uvicorn server:app --host 0.0.0.0 --port 8000`. +WHISPERX_PORT = 8000 + +# First boot warm-loads the default Whisper model (downloaded from HF) and the +# baked pyannote diarization pipeline inside the FastAPI lifespan hook *before* +# uvicorn binds the socket, so /ping is only reachable once models are resident. +# Give it a generous budget. +HEALTH_TIMEOUT = 300 # seconds to wait for /ping to return 200 +HEALTH_INTERVAL = 5 # seconds between health checks +# A transcription that needs word timestamps or diarization triggers a lazy +# download of the wav2vec2 aligner for the detected language on first use, so +# the request timeout must tolerate a cold HF fetch. +REQUEST_TIMEOUT = 600 # seconds to wait for a transcription response + +# Audio fixtures staged in S3 (see test-fixtures/audio/ under the models bucket). +S3_BUCKET = "dlc-cicd-models" +S3_AUDIO_PREFIX = "test-fixtures/audio" +AUDIO_EN = "asr_en.wav" # English ~15s clip +AUDIO_ZH = "asr_zh.wav" # Chinese clip (non-English + language param) +AUDIO_DIARIZE_2SPK = "asr_diarize_2spk.wav" # 2 distinct speakers ~21s + + +# --------------------------------------------------------------------------- +# Container lifecycle helpers +# --------------------------------------------------------------------------- + + +def start_container(image_uri, device, docker_run_flags=None, env=None): + """Start a Docker container running the WhisperX FastAPI server. + + No model mount and no serve command: the image bakes/downloads its models + and the ENTRYPOINT runs uvicorn on port 8000. + + Args: + image_uri: Full ECR image URI. + device: "gpu" or "cpu" (informational; GPU access comes from flags). + docker_run_flags: Extra flags for docker run (e.g. ["--gpus", "all"]). + env: Optional dict of environment variables, emitted as `-e K=V` flags + (e.g. {"WHISPERX_DEFAULT_MODEL": "tiny"}). Lets a test pin the served + model, alias, or override behavior at launch. + + Returns: + (container_id, port) + """ + cmd = [ + "docker", + "run", + "-d", + "--shm-size=2g", + "-p", + f"{WHISPERX_PORT}:{WHISPERX_PORT}", + ] + for key, value in (env or {}).items(): + cmd.extend(["-e", f"{key}={value}"]) + if docker_run_flags: + cmd.extend(docker_run_flags) + cmd.append(image_uri) + + LOGGER.info(f"Starting {device} container: {' '.join(cmd)}") + result = subprocess.run(cmd, capture_output=True, text=True, check=True) + container_id = result.stdout.strip() + LOGGER.info(f"Container started: {container_id[:12]}") + return container_id, WHISPERX_PORT + + +def wait_for_health(port=WHISPERX_PORT, timeout=HEALTH_TIMEOUT, interval=HEALTH_INTERVAL): + """Poll GET /ping until it returns 200, or raise TimeoutError. + + The socket is refused until the lifespan warm-load finishes, so connection + errors are expected and simply retried. + """ + endpoint = f"http://localhost:{port}/ping" + deadline = time.time() + timeout + while time.time() < deadline: + try: + resp = requests.get(endpoint, timeout=5) + if resp.status_code == 200: + LOGGER.info("WhisperX server is healthy") + return True + except requests.RequestException: + pass + time.sleep(interval) + raise TimeoutError(f"WhisperX server did not become healthy within {timeout}s") + + +def get_container_logs(container_id): + """Return combined stdout+stderr from `docker logs` for debugging.""" + result = subprocess.run( + ["docker", "logs", container_id], + capture_output=True, + text=True, + ) + return result.stdout + result.stderr + + +def stop_container(container_id): + """Dump logs (best effort) then force-remove a Docker container.""" + LOGGER.info(f"Stopping container {container_id[:12]}") + subprocess.run(["docker", "rm", "-f", container_id], capture_output=True) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def make_container_fixture(device, docker_run_flags=None): + """Create the container fixture: start container, health-check, cleanup. + + Yields a dict with container_id and port for the test to use. On a health + timeout the container logs are dumped and the container is removed before + failing, so failures point directly at an image/server problem. + """ + + @pytest.fixture(scope="function") + def container(image_uri): + container_id, port = start_container(image_uri, device, docker_run_flags) + try: + wait_for_health(port=port) + except TimeoutError: + logs = get_container_logs(container_id) + LOGGER.error(f"Container logs:\n{logs}") + stop_container(container_id) + pytest.fail("WhisperX server health check timed out") + + yield {"container_id": container_id, "port": port} + + stop_container(container_id) + + return container + + +@contextmanager +def run_container_with_env(image_uri, env, device="gpu", flags=("--gpus", "all")): + """Start a WhisperX container with custom env vars; yield {container_id, port}. + + A context-manager sibling of make_container_fixture for tests that need a + per-test container launched with specific env (e.g. WHISPERX_DEFAULT_MODEL=tiny + to pin a fast-warming served model). It mirrors the fixture's start -> + health-check -> (log dump on timeout) -> yield -> teardown flow, but wraps the + yield in try/finally so `docker rm -f` always runs — a health timeout or a + failing test body must never leak the container (a leak blocks the runner). + + Usage: + with run_container_with_env(image_uri, {"WHISPERX_DEFAULT_MODEL": "tiny"}) as c: + ... requests to c["port"] ... + """ + container_id, port = start_container(image_uri, device, docker_run_flags=list(flags), env=env) + try: + try: + wait_for_health(port=port) + except TimeoutError: + logs = get_container_logs(container_id) + LOGGER.error(f"Container logs:\n{logs}") + pytest.fail("WhisperX server health check timed out") + + yield {"container_id": container_id, "port": port} + finally: + stop_container(container_id) + + +# --------------------------------------------------------------------------- +# S3 fixture download +# --------------------------------------------------------------------------- + + +def download_fixture(aws_session, key, dest): + """Download an audio fixture from S3 to a local path. + + Args: + aws_session: The shared AWSSessionManager (provides a boto3 s3 client). + key: Bare fixture filename under S3_AUDIO_PREFIX (e.g. "asr_en.wav"). + dest: Local path to write the downloaded file to. + + Returns: + dest (for convenient inline use). + """ + s3_key = f"{S3_AUDIO_PREFIX}/{key}" + LOGGER.info(f"Downloading s3://{S3_BUCKET}/{s3_key} -> {dest}") + aws_session.s3.download_file(S3_BUCKET, s3_key, dest) + return dest + + +# --------------------------------------------------------------------------- +# HTTP helpers +# --------------------------------------------------------------------------- + + +def post_transcription(port, audio_path, timeout=REQUEST_TIMEOUT, **form_fields): + """POST an audio file to /v1/audio/transcriptions as multipart/form-data. + + The audio is sent as the required `file` part. Remaining transcription + parameters are passed as keyword args and encoded as form fields: + + - None values are dropped (the server applies its own defaults). + - bool values become "true"/"false" (what Pydantic bool parsing expects). + - list/tuple values are emitted as a repeated form field, which is exactly + how the server reads the `timestamp_granularities[]` list, e.g. + `post_transcription(port, path, response_format="verbose_json", + **{"timestamp_granularities[]": ["word"]})`. + + Returns the raw requests.Response so callers can assert on status code, + headers/content-type, and body (json or text). + """ + with open(audio_path, "rb") as fh: + audio_bytes = fh.read() + files = {"file": (os.path.basename(audio_path), audio_bytes, "audio/wav")} + + # A list of (key, value) tuples lets us repeat a key for list-valued fields + # while `files` is present (requests encodes both into one multipart body). + data = [] + for key, value in form_fields.items(): + if value is None: + continue + if isinstance(value, bool): + data.append((key, "true" if value else "false")) + elif isinstance(value, (list, tuple)): + for item in value: + data.append((key, str(item))) + else: + data.append((key, str(value))) + + return requests.post( + f"http://localhost:{port}/v1/audio/transcriptions", + files=files, + data=data, + timeout=timeout, + ) diff --git a/test/whisperx/ec2/requirements.txt b/test/whisperx/ec2/requirements.txt new file mode 100644 index 000000000000..f2293605cf1b --- /dev/null +++ b/test/whisperx/ec2/requirements.txt @@ -0,0 +1 @@ +requests diff --git a/test/whisperx/ec2/test_ec2_gpu.py b/test/whisperx/ec2/test_ec2_gpu.py new file mode 100644 index 000000000000..055edcaf89f9 --- /dev/null +++ b/test/whisperx/ec2/test_ec2_gpu.py @@ -0,0 +1,257 @@ +# Copyright 2018-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""EC2 container integration tests for the WhisperX ASR DLC (GPU). + +Runs the container locally via `docker run --gpus all`, waits for the FastAPI +server on port 8000 to warm-load its models, and validates the OpenAI-compatible +transcription contract against staged audio fixtures. + +Assertions are deliberately loose about ASR *content* (ASR output is +nondeterministic): we assert the response contract — status codes, response +shapes, content types, and timestamp/speaker structure — not exact transcript +strings. Container lifecycle lives in common.py; this file only sets device +config and the test cases. +""" + +import requests +from whisperx.ec2.common import ( + AUDIO_DIARIZE_2SPK, + AUDIO_EN, + AUDIO_ZH, + download_fixture, + make_container_fixture, + post_transcription, +) + +DEVICE = "gpu" +DOCKER_RUN_FLAGS = ["--gpus", "all"] + +# Register the container fixture for this module (function-scoped: fresh +# container per test, mirroring the Ray EC2 suite). +container = make_container_fixture(DEVICE, docker_run_flags=DOCKER_RUN_FLAGS) + + +def test_ping_and_models(container): + """/ping is healthy and /v1/models advertises at least one served model id.""" + port = container["port"] + + resp = requests.get(f"http://localhost:{port}/ping", timeout=10) + assert resp.status_code == 200 + assert resp.json() == {"status": "ok"} + + resp = requests.get(f"http://localhost:{port}/v1/models", timeout=10) + assert resp.status_code == 200 + body = resp.json() + assert body.get("object") == "list" + data = body.get("data") + assert isinstance(data, list) and len(data) >= 1, f"expected >=1 model, got {body}" + assert data[0].get("id"), "expected a non-empty model id" + + +def test_basic_transcription(container, aws_session, tmp_path): + """Default json transcription of English audio returns non-empty text.""" + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + + resp = post_transcription(container["port"], audio, response_format="json") + assert resp.status_code == 200, resp.text + + body = resp.json() + # json response shape is exactly {"text": "..."}. + assert body.get("text", "").strip(), "expected non-empty transcription text" + + +def test_word_timestamps(container, aws_session, tmp_path): + """verbose_json + timestamp_granularities[]=word yields coherent word timings.""" + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + + resp = post_transcription( + container["port"], + audio, + response_format="verbose_json", + # The server reads the list off the literal "timestamp_granularities[]" key. + **{"timestamp_granularities[]": ["word"]}, + ) + assert resp.status_code == 200, resp.text + + body = resp.json() + words = body.get("words") + assert isinstance(words, list) and words, f"expected non-empty words list, got {body.keys()}" + for w in words: + assert {"word", "start", "end"} <= set(w), f"word entry missing keys: {w}" + + # WhisperX alignment can leave a few tokens (e.g. digits) without timings; + # assert coherence on the words that ARE timed and require at least one. + timed = [ + (w["start"], w["end"]) for w in words if w["start"] is not None and w["end"] is not None + ] + assert timed, "expected at least one word with start/end timestamps" + for start, end in timed: + assert isinstance(start, (int, float)) and isinstance(end, (int, float)) + assert start <= end, f"word start {start} after end {end}" + + +def test_language_non_english(container, aws_session, tmp_path): + """Forcing language=zh transcribes Chinese audio and echoes the language.""" + audio = download_fixture(aws_session, AUDIO_ZH, str(tmp_path / AUDIO_ZH)) + + resp = post_transcription( + container["port"], + audio, + language="zh", + response_format="verbose_json", + ) + assert resp.status_code == 200, resp.text + + body = resp.json() + assert body.get("language") == "zh", f"expected language 'zh', got {body.get('language')!r}" + assert body.get("text", "").strip(), "expected non-empty transcription text" + + +def test_diarization(container, aws_session, tmp_path): + """diarize=true on a 2-speaker clip returns >=2 speakers and per-segment labels.""" + audio = download_fixture(aws_session, AUDIO_DIARIZE_2SPK, str(tmp_path / AUDIO_DIARIZE_2SPK)) + + resp = post_transcription( + container["port"], + audio, + diarize=True, + response_format="verbose_json", + ) + assert resp.status_code == 200, resp.text + + body = resp.json() + speakers = body.get("speakers") + assert isinstance(speakers, list), f"expected speakers list, got {body.keys()}" + assert len(set(speakers)) >= 2, f"expected >=2 distinct speakers, got {speakers}" + + segments = body.get("segments", []) + seg_speakers = [s.get("speaker") for s in segments if s.get("speaker")] + assert seg_speakers, "expected at least one segment to carry a speaker label" + + +def test_response_formats(container, aws_session, tmp_path): + """text / srt / vtt each return the right content-type and structure.""" + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + port = container["port"] + + resp = post_transcription(port, audio, response_format="text") + assert resp.status_code == 200, resp.text + assert "text/plain" in resp.headers.get("content-type", "") + assert resp.text.strip(), "expected non-empty plain text" + + resp = post_transcription(port, audio, response_format="srt") + assert resp.status_code == 200, resp.text + assert "application/x-subrip" in resp.headers.get("content-type", "") + assert "-->" in resp.text, "expected SRT cue timing arrow" + + resp = post_transcription(port, audio, response_format="vtt") + assert resp.status_code == 200, resp.text + assert "text/vtt" in resp.headers.get("content-type", "") + assert resp.text.startswith("WEBVTT"), "expected VTT header" + + +def test_errors(container, tmp_path): + """Contract errors: bad response_format -> 400, empty upload -> 400. + + The server validates response_format before reading the file, so the format + case needs only a present (unread) file part. + """ + port = container["port"] + + # A present-but-never-read file part (validation short-circuits before read). + dummy = tmp_path / "dummy.wav" + dummy.write_bytes(b"RIFF0000WAVE") + + resp = post_transcription(port, str(dummy), response_format="xml") + assert resp.status_code == 400, resp.text + + empty = tmp_path / "empty.wav" + empty.write_bytes(b"") + resp = post_transcription(port, str(empty), response_format="json") + assert resp.status_code == 400, resp.text + + +def test_verbose_json_fields(container, aws_session, tmp_path): + """verbose_json carries the full envelope: task, language, duration, text, segments.""" + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + + resp = post_transcription(container["port"], audio, response_format="verbose_json") + assert resp.status_code == 200, resp.text + + body = resp.json() + assert body.get("task") == "transcribe", f"expected task 'transcribe', got {body.get('task')!r}" + assert isinstance(body.get("language"), str) and body["language"], "expected a language code" + assert isinstance(body.get("duration"), (int, float)), "expected numeric duration" + assert body.get("text", "").strip(), "expected non-empty text" + assert isinstance(body.get("segments"), list), f"expected segments list, got {body.keys()}" + + +def test_json_minimal(container, aws_session, tmp_path): + """Default json is the minimal OpenAI shape: only `text`, no verbose extras.""" + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + + resp = post_transcription(container["port"], audio, response_format="json") + assert resp.status_code == 200, resp.text + + body = resp.json() + assert body.get("text", "").strip(), "expected non-empty text" + # json omits the verbose_json envelope entirely (OpenAI-compatible minimal body). + for extra in ("segments", "words", "speakers"): + assert extra not in body, f"json response should not carry {extra!r}: {body.keys()}" + + +def test_diarize_false_no_speakers(container, aws_session, tmp_path): + """diarize omitted (default false) -> no `speakers` key and no per-segment speaker. + + DESIGN Dim 5 / OpenAI byte-compat: without diarization the verbose_json body + must be indistinguishable from OpenAI's — no top-level speakers list and no + `speaker` field leaking onto any segment. + """ + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + + resp = post_transcription(container["port"], audio, response_format="verbose_json") + assert resp.status_code == 200, resp.text + + body = resp.json() + assert "speakers" not in body, f"non-diarized verbose_json must omit 'speakers': {body.keys()}" + for seg in body.get("segments", []): + assert "speaker" not in seg, f"non-diarized segment must omit 'speaker': {seg}" + + +def test_max_speakers_one(container, aws_session, tmp_path): + """diarize=true with max_speakers=1 caps the diarization at a single speaker. + + max_speakers is a hard upper bound on pyannote's clustering, so on a + 2-speaker clip it should collapse to one label. The `speakers >= 1` floor is + guaranteed for any diarized clip; the `<= 1` bound is the assertion that the + cap was actually honored and is the most backend-sensitive check in the + suite — if a future pyannote pipeline treats max_speakers as advisory this is + the line that would need revisiting. + """ + audio = download_fixture(aws_session, AUDIO_DIARIZE_2SPK, str(tmp_path / AUDIO_DIARIZE_2SPK)) + + resp = post_transcription( + container["port"], + audio, + diarize=True, + max_speakers=1, + response_format="verbose_json", + ) + assert resp.status_code == 200, resp.text + + body = resp.json() + speakers = body.get("speakers") + assert isinstance(speakers, list) and speakers, ( + f"expected a non-empty speakers list: {body.keys()}" + ) + assert len(set(speakers)) <= 1, f"max_speakers=1 should cap to a single speaker, got {speakers}" diff --git a/test/whisperx/ec2/test_ec2_model_config.py b/test/whisperx/ec2/test_ec2_model_config.py new file mode 100644 index 000000000000..2d9bddafab97 --- /dev/null +++ b/test/whisperx/ec2/test_ec2_model_config.py @@ -0,0 +1,113 @@ +# Copyright 2018-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"). You +# may not use this file except in compliance with the License. A copy of +# the License is located at +# +# http://aws.amazon.com/apache2.0/ +# +# or in the "license" file accompanying this file. This file is +# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF +# ANY KIND, either express or implied. See the License for the specific +# language governing permissions and limitations under the License. +"""EC2 model-launch-config tests for the WhisperX ASR DLC (GPU). + +Group A: these tests validate how the server's launch-time env vars shape the +served-model contract — WHISPERX_DEFAULT_MODEL (the one model per container) and +WHISPERX_SERVED_MODEL_NAME (a client-facing alias advertised by /v1/models). +Each case needs a container launched with different env, so unlike test_ec2_gpu.py +they do NOT use the shared default `container` fixture; instead each spins up (and +tears down) a custom-env container via common.run_container_with_env. + +The Whisper model is one-per-container: the request `model` field is accepted but +ignored (an OpenAI-compat no-op), so every transcription runs the launched model +regardless of the `model` value sent. + +WHISPERX_DEFAULT_MODEL=tiny is used throughout so the startup warm-load finishes +in seconds — tiny is a valid faster-whisper model that downloads far faster than +the large-v2 default. + +Assertions target the served contract only (advertised model id, and that any +`model` value transcribes 200) — never exact transcript text, which is +nondeterministic. +""" + +import requests +from whisperx.ec2.common import ( + AUDIO_EN, + download_fixture, + post_transcription, + run_container_with_env, +) + +DEVICE = "gpu" +DOCKER_RUN_FLAGS = ["--gpus", "all"] +# Small, fast-downloading Whisper model so the startup warm-load takes seconds, +# not the minutes large-v2 needs. tiny is a valid faster-whisper id. +FAST_MODEL = "tiny" + + +def test_custom_default_model(image_uri, aws_session, tmp_path): + """WHISPERX_DEFAULT_MODEL sets the one served model; the `model` field is ignored. + + Launches with WHISPERX_DEFAULT_MODEL=tiny and asserts /v1/models advertises + `tiny`, that omitting `model` and naming `tiny` both transcribe (200), and + that naming a *different* id (`large-v2`) also transcribes 200 — the field is + an ignored no-op, so the launched model runs regardless. + """ + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + with run_container_with_env( + image_uri, {"WHISPERX_DEFAULT_MODEL": FAST_MODEL}, device=DEVICE, flags=DOCKER_RUN_FLAGS + ) as c: + port = c["port"] + + resp = requests.get(f"http://localhost:{port}/v1/models", timeout=10) + assert resp.status_code == 200, resp.text + data = resp.json()["data"] + assert data[0]["id"] == FAST_MODEL, f"expected served id {FAST_MODEL!r}, got {data}" + + # Omitted model -> serves the launched model -> 200 with real text. + resp = post_transcription(port, audio, response_format="json") + assert resp.status_code == 200, resp.text + assert resp.json().get("text", "").strip(), "expected non-empty text with default model" + + # Naming the launched model explicitly -> 200. + resp = post_transcription(port, audio, model=FAST_MODEL, response_format="json") + assert resp.status_code == 200, resp.text + + # A different id -> still 200: the `model` field is ignored, so the + # launched model serves the request instead of being rejected. + resp = post_transcription(port, audio, model="large-v2", response_format="json") + assert resp.status_code == 200, resp.text + assert resp.json().get("text", "").strip(), "expected launched model to serve any model id" + + +def test_served_model_alias(image_uri, aws_session, tmp_path): + """WHISPERX_SERVED_MODEL_NAME sets the id /v1/models advertises for OpenAI clients. + + Launches tiny behind the alias `whisper-1` and asserts /v1/models reports + `whisper-1`, and that the alias, the real backing id `tiny`, and an unrelated + id all transcribe (200) — the `model` field is an ignored no-op, so every + request runs the launched model. + """ + audio = download_fixture(aws_session, AUDIO_EN, str(tmp_path / AUDIO_EN)) + env = {"WHISPERX_DEFAULT_MODEL": FAST_MODEL, "WHISPERX_SERVED_MODEL_NAME": "whisper-1"} + with run_container_with_env(image_uri, env, device=DEVICE, flags=DOCKER_RUN_FLAGS) as c: + port = c["port"] + + resp = requests.get(f"http://localhost:{port}/v1/models", timeout=10) + assert resp.status_code == 200, resp.text + data = resp.json()["data"] + assert data[0]["id"] == "whisper-1", f"expected alias id 'whisper-1', got {data}" + + # The advertised alias transcribes -> 200. + resp = post_transcription(port, audio, model="whisper-1", response_format="json") + assert resp.status_code == 200, resp.text + + # The real backing model id transcribes too -> 200. + resp = post_transcription(port, audio, model=FAST_MODEL, response_format="json") + assert resp.status_code == 200, resp.text + + # An unrelated id -> still 200: the `model` field is ignored. + resp = post_transcription(port, audio, model="whisper-large", response_format="json") + assert resp.status_code == 200, resp.text diff --git a/test/whisperx/sagemaker/requirements.txt b/test/whisperx/sagemaker/requirements.txt new file mode 100644 index 000000000000..f60dca57bf56 --- /dev/null +++ b/test/whisperx/sagemaker/requirements.txt @@ -0,0 +1,3 @@ +boto3 +pytest +sagemaker>=3.0.0 diff --git a/test/whisperx/sagemaker/test_async_endpoint.py b/test/whisperx/sagemaker/test_async_endpoint.py new file mode 100644 index 000000000000..e577a7ff3b50 --- /dev/null +++ b/test/whisperx/sagemaker/test_async_endpoint.py @@ -0,0 +1,261 @@ +"""WhisperX ASR SageMaker async-inference integration test — SageMaker SDK v3. + +Async inference removes SageMaker's 60s real-time invoke cap, which is the +recommended path for long audio. The request payload is uploaded to S3, the +endpoint is invoked with `invoke_endpoint_async`, and the JSON response is polled +back from an S3 output location. + +Config source: this suite uses the shared `image_uri` pytest fixture + `aws_session` +(the newer pattern already proven for async by test/vllm-omni), NOT openfold3's +TEST_IMAGE_URI/SM_ROLE_ARN env vars — so the CI workflow can pass `--image-uri` +uniformly across the EC2/sync/async suites. The openfold3 async *mechanics* +(S3 preflight + writability probe, stale-resource sweep, failure-location handling, +S3 polling) are mirrored here. + +S3 buckets: the SageMaker execution role (AmazonSageMakerFullAccess) can only +read/write buckets whose name contains "sagemaker", so both the async input payload +and the async output live in the account default `sagemaker--` +bucket. The audio fixture is read from `dlc-cicd-models` with the test runner's own +credentials, then re-uploaded (wrapped in a multipart body) to the sagemaker bucket. + +AMI pin (DESIGN decision 9): CUDA 12.8 image -> must use INFERENCE_AMI_VERSION_CU12 +(AL2, driver 550). The default CU13 AMI causes a zero-log CannotStartContainerError. + +Validation is smoke-level and deterministic: a successful async invocation returns +JSON with non-empty transcription text. WhisperX output is not byte-stable. +""" + +import json +import logging +import time +import uuid + +import pytest +from sagemaker.core.resources import Endpoint, EndpointConfig, Model +from sagemaker.core.shapes import ( + AsyncInferenceClientConfig, + AsyncInferenceConfig, + AsyncInferenceOutputConfig, + ContainerDefinition, + ProductionVariant, +) +from test_utils import random_suffix_name +from test_utils.constants import INFERENCE_AMI_VERSION_CU12, SAGEMAKER_ROLE + +LOGGER = logging.getLogger(__name__) +LOGGER.setLevel(logging.INFO) + +INSTANCE_TYPE = "ml.g6.xlarge" +MAX_CONCURRENT_INVOCATIONS = 1 +# Generous: startup warms the Whisper model + diarization pipeline before /ping. +STARTUP_HEALTH_CHECK_TIMEOUT = 1200 +DEPLOY_TIMEOUT = 1800 + +# Source audio fixture (read with the test runner's credentials). +AUDIO_BUCKET = "dlc-cicd-models" +AUDIO_PREFIX = "test-fixtures/audio" +AUDIO_LONG = "asr_long_60s.wav" # English 60s — long-audio async path + +# S3 layout on the sagemaker-named bucket (writable by the SageMaker role). +INPUT_PREFIX = "whisperx-async-input" +OUTPUT_PREFIX = "whisperx-async-output" + +# Bounded so a stuck request fails fast instead of hanging. +POLL_TIMEOUT = 900 +POLL_INTERVAL = 5 + +RESOURCE_NAME_PREFIX = "whisperx-async" + + +def _split(uri: str) -> tuple[str, str]: + """Parse s3://bucket/key into (bucket, key).""" + return uri.split("/")[2], "/".join(uri.split("/")[3:]) + + +def _build_multipart(audio_bytes: bytes, filename: str, fields: dict) -> tuple[bytes, str]: + """Build a multipart/form-data body with the audio `file` part + string fields.""" + boundary = uuid.uuid4().hex + parts = [ + f'--{boundary}\r\nContent-Disposition: form-data; name="file"; ' + f'filename="{filename}"\r\nContent-Type: audio/wav\r\n\r\n'.encode(), + audio_bytes, + b"\r\n", + ] + for key, value in fields.items(): + parts.append( + f'--{boundary}\r\nContent-Disposition: form-data; name="{key}"\r\n\r\n' + f"{value}\r\n".encode() + ) + parts.append(f"--{boundary}--\r\n".encode()) + return b"".join(parts), f"multipart/form-data; boundary={boundary}" + + +@pytest.fixture(scope="module") +def async_endpoint(aws_session, image_uri): + """Deploy one async WhisperX endpoint for the module; sweep stale first, clean up after. + + Yields (endpoint_name, io_bucket) — the test uploads its own input payload and + polls the output, mirroring openfold3's boto3 async flow. + """ + s3 = aws_session.s3 + account_id = aws_session.sts.get_caller_identity()["Account"] + io_bucket = f"sagemaker-{aws_session.region}-{account_id}" + + _preflight_s3(s3, io_bucket) + _sweep_stale(aws_session.sagemaker) + + endpoint_name = random_suffix_name(RESOURCE_NAME_PREFIX, 50) + model_name = endpoint_name + role_arn = aws_session.resolve_role_arn(SAGEMAKER_ROLE) + s3_output = f"s3://{io_bucket}/{OUTPUT_PREFIX}/" + + LOGGER.info(f"Using image: {image_uri}") + + model = endpoint_config = endpoint = None + try: + LOGGER.info(f"Creating model: {model_name}") + model = Model.create( + model_name=model_name, + primary_container=ContainerDefinition(image=image_uri), + execution_role_arn=role_arn, + ) + + LOGGER.info(f"Creating async endpoint config: {endpoint_name}") + endpoint_config = EndpointConfig.create( + endpoint_config_name=endpoint_name, + production_variants=[ + ProductionVariant( + variant_name="AllTraffic", + model_name=model_name, + initial_instance_count=1, + instance_type=INSTANCE_TYPE, + # CUDA 12.8 image -> must use the CU12 (driver 550) AMI. + inference_ami_version=INFERENCE_AMI_VERSION_CU12, + container_startup_health_check_timeout_in_seconds=STARTUP_HEALTH_CHECK_TIMEOUT, + ), + ], + async_inference_config=AsyncInferenceConfig( + output_config=AsyncInferenceOutputConfig(s3_output_path=s3_output), + client_config=AsyncInferenceClientConfig( + max_concurrent_invocations_per_instance=MAX_CONCURRENT_INVOCATIONS, + ), + ), + ) + + LOGGER.info(f"Deploying async endpoint {endpoint_name} (~10-15 min)...") + endpoint = Endpoint.create(endpoint_name=endpoint_name, endpoint_config_name=endpoint_name) + endpoint.wait_for_status("InService", timeout=DEPLOY_TIMEOUT) + LOGGER.info("Async endpoint InService") + + yield endpoint_name, io_bucket + finally: + for resource in (endpoint, endpoint_config, model): + if resource is None: + continue + try: + resource.delete() + except Exception as e: + LOGGER.warning(f"Cleanup {type(resource).__name__} failed: {e}") + + +def _preflight_s3(s3, io_bucket: str): + """Fail before the ~10-min deploy if the fixture is missing or output isn't writable.""" + try: + s3.head_object(Bucket=AUDIO_BUCKET, Key=f"{AUDIO_PREFIX}/{AUDIO_LONG}") + except s3.exceptions.ClientError as e: + raise AssertionError( + f"Audio fixture s3://{AUDIO_BUCKET}/{AUDIO_PREFIX}/{AUDIO_LONG} not found ({e})." + ) from e + probe = f"{OUTPUT_PREFIX}/.preflight-{uuid.uuid4()}" + try: + s3.put_object(Bucket=io_bucket, Key=probe, Body=b"ok") + s3.delete_object(Bucket=io_bucket, Key=probe) + except s3.exceptions.ClientError as e: + raise AssertionError( + f"Cannot write to s3://{io_bucket}/{OUTPUT_PREFIX}/ ({e}). The async endpoint " + f"output will not be writable — the bucket name must contain 'sagemaker' so the " + f"SageMaker execution role can write to it." + ) from e + + +def _sweep_stale(sm): + """Delete leftover whisperx-async-* SageMaker resources from a prior canceled run.""" + try: + for ep in sm.list_endpoints(NameContains=RESOURCE_NAME_PREFIX).get("Endpoints", []): + LOGGER.warning(f"[sweep] deleting stale endpoint {ep['EndpointName']}") + try: + sm.delete_endpoint(EndpointName=ep["EndpointName"]) + except Exception as e: + LOGGER.warning(f"[sweep] {e}") + for c in sm.list_endpoint_configs(NameContains=RESOURCE_NAME_PREFIX).get( + "EndpointConfigs", [] + ): + try: + sm.delete_endpoint_config(EndpointConfigName=c["EndpointConfigName"]) + except Exception as e: + LOGGER.warning(f"[sweep] {e}") + for m in sm.list_models(NameContains=RESOURCE_NAME_PREFIX).get("Models", []): + try: + sm.delete_model(ModelName=m["ModelName"]) + except Exception as e: + LOGGER.warning(f"[sweep] {e}") + except Exception as e: + LOGGER.warning(f"[sweep] skipped: {e}") + + +def test_async_long_audio(async_endpoint, aws_session): + """Long (60s) audio transcribes through the async S3-in/S3-out path -> non-empty text.""" + endpoint_name, io_bucket = async_endpoint + s3 = aws_session.s3 + smr = aws_session.session.client("sagemaker-runtime") + + # Fetch the fixture with the test runner's creds, wrap it in a multipart body, and + # stage it on the sagemaker bucket (the only bucket the SageMaker role can read). + audio_bytes = s3.get_object(Bucket=AUDIO_BUCKET, Key=f"{AUDIO_PREFIX}/{AUDIO_LONG}")[ + "Body" + ].read() + body, content_type = _build_multipart(audio_bytes, AUDIO_LONG, {"language": "en"}) + input_key = f"{INPUT_PREFIX}/{endpoint_name}-input.bin" + s3.put_object(Bucket=io_bucket, Key=input_key, Body=body, ContentType=content_type) + input_location = f"s3://{io_bucket}/{input_key}" + + LOGGER.info(f"Submitting async invocation for {input_location}") + resp = smr.invoke_endpoint_async( + EndpointName=endpoint_name, + InputLocation=input_location, + ContentType=content_type, + ) + output_bucket, output_key = _split(resp["OutputLocation"]) + failure_uri = resp.get("FailureLocation", "") + failure_bucket, failure_key = _split(failure_uri) if failure_uri else (None, None) + + # Poll for the output object; surface a written failure object immediately. + deadline = time.time() + POLL_TIMEOUT + while time.time() < deadline: + try: + obj = s3.get_object(Bucket=output_bucket, Key=output_key) + result = json.loads(obj["Body"].read()) + text = result.get("text", "") + assert isinstance(text, str) and text.strip(), ( + f"Async transcription text is empty: {result!r}" + ) + LOGGER.info(f"Async transcription text (len={len(text)}): {text[:120]!r}") + return + except s3.exceptions.ClientError: + pass + if failure_bucket: + try: + s3.head_object(Bucket=failure_bucket, Key=failure_key) + err = ( + s3.get_object(Bucket=failure_bucket, Key=failure_key)["Body"] + .read() + .decode()[:2000] + ) + raise AssertionError(f"SageMaker async failure object: {err}") + except s3.exceptions.ClientError: + pass + time.sleep(POLL_INTERVAL) + + raise TimeoutError( + f"No async output at s3://{output_bucket}/{output_key} within {POLL_TIMEOUT}s" + ) diff --git a/test/whisperx/sagemaker/test_sm_endpoint.py b/test/whisperx/sagemaker/test_sm_endpoint.py new file mode 100644 index 000000000000..81b57315d9ae --- /dev/null +++ b/test/whisperx/sagemaker/test_sm_endpoint.py @@ -0,0 +1,223 @@ +"""Integration tests for the WhisperX ASR SageMaker endpoint (SYNC) — SageMaker SDK v3. + +Deploys one real-time inference endpoint for the module and drives it through the +SageMaker `/invocations` alias (identical handler to `/v1/audio/transcriptions`). +The container is a FastAPI WhisperX server that expects `multipart/form-data` with a +required `file` part plus optional form fields, so each request builds a multipart +body locally and sends the raw bytes via `endpoint.invoke(body=..., content_type=...)`. + +AMI pin (DESIGN decision 9): WhisperX is a CUDA 12.8 GPU image, so the endpoint MUST +run on INFERENCE_AMI_VERSION_CU12 (AL2, driver 550). The default test AMI is CU13 +(driver 580); using it makes the container fail to start with a zero-log +CannotStartContainerError. Pin CU12 explicitly. + +Validation is smoke-level and deterministic: transcription returns non-empty text and +diarization surfaces >=2 speakers. WhisperX output is not byte-stable, so no +exact-transcript assertions. +""" + +import json +import logging +import os +import time +import uuid + +import pytest +from sagemaker.core.resources import Endpoint, EndpointConfig, Model +from sagemaker.core.shapes import ContainerDefinition, ProductionVariant +from test_utils import random_suffix_name +from test_utils.constants import INFERENCE_AMI_VERSION_CU12, SAGEMAKER_ROLE + +LOGGER = logging.getLogger(__name__) +LOGGER.setLevel(logging.INFO) + +# Single L4 GPU; matches WhisperX VRAM needs (large-v2 + align + diarization). +INSTANCE_TYPE = "ml.g6.xlarge" +# Startup warms the default Whisper model + diarization pipeline in the lifespan +# hook before /ping returns 200, so InService already implies a warm model. +STARTUP_HEALTH_CHECK_TIMEOUT = 900 +# Deploy can take 10-15 min for a GPU endpoint. +DEPLOY_TIMEOUT = 1800 + +# Audio fixtures pre-staged in S3 (read with the test runner's credentials). +AUDIO_BUCKET = "dlc-cicd-models" +AUDIO_PREFIX = "test-fixtures/audio" +AUDIO_EN = "asr_en.wav" # English ~15s +AUDIO_DIARIZE = "asr_diarize_2spk.wav" # 2 speakers ~21s + + +def _cleanup(resources): + """Best-effort delete for a list of v3 resource objects (None-safe). + + Leaked SageMaker endpoints keep billing, so cleanup must run even when a + test fails — callers invoke this from a `finally` block. + """ + for resource in resources: + if resource is None: + continue + try: + resource.delete() + except Exception as e: + LOGGER.warning(f"Cleanup {type(resource).__name__} failed: {e}") + + +def _build_multipart(audio_path: str, fields: dict) -> tuple[bytes, str]: + """Build a multipart/form-data body with the audio `file` part + string fields. + + SageMaker InvokeEndpoint forwards the ContentType header (boundary included) + straight through to the model server, so the client owns the encoding. + """ + boundary = uuid.uuid4().hex + filename = os.path.basename(audio_path) + with open(audio_path, "rb") as f: + audio = f.read() + + parts = [ + f'--{boundary}\r\nContent-Disposition: form-data; name="file"; ' + f'filename="{filename}"\r\nContent-Type: audio/wav\r\n\r\n'.encode(), + audio, + b"\r\n", + ] + for key, value in fields.items(): + parts.append( + f'--{boundary}\r\nContent-Disposition: form-data; name="{key}"\r\n\r\n' + f"{value}\r\n".encode() + ) + parts.append(f"--{boundary}--\r\n".encode()) + + body = b"".join(parts) + content_type = f"multipart/form-data; boundary={boundary}" + return body, content_type + + +def _invoke_transcription(endpoint, audio_path: str, retries: int = 3, **fields) -> dict: + """Invoke /invocations with a multipart audio request and return the parsed JSON. + + The first invocation of a given code path may lazily load an alignment model, + which can nudge past SageMaker's 60s real-time invoke cap; the model is cached + server-side afterward, so a bounded retry rides out that cold start + deterministically rather than flaking. + """ + body, content_type = _build_multipart(audio_path, fields) + last_error = None + for attempt in range(retries): + try: + result = endpoint.invoke(body=body, content_type=content_type) + return json.loads(result.body.read()) + except Exception as e: + last_error = e + LOGGER.warning(f"Invoke attempt {attempt + 1}/{retries} failed: {e}") + if attempt < retries - 1: + time.sleep(30) + raise last_error + + +@pytest.fixture(scope="session") +def audio_cache(aws_session, tmp_path_factory): + """Return a downloader that fetches (and memoizes) an audio fixture locally.""" + cache_dir = str(tmp_path_factory.mktemp("whisperx-audio")) + downloaded: dict[str, str] = {} + + def _get(name: str) -> str: + if name not in downloaded: + dest = os.path.join(cache_dir, name) + LOGGER.info(f"Downloading s3://{AUDIO_BUCKET}/{AUDIO_PREFIX}/{name}") + aws_session.s3.download_file(AUDIO_BUCKET, f"{AUDIO_PREFIX}/{name}", dest) + downloaded[name] = dest + return downloaded[name] + + return _get + + +@pytest.fixture(scope="module") +def model_endpoint(aws_session, image_uri): + """Deploy one WhisperX real-time endpoint for the module; clean up after. + + Module-scoped so all sync tests share a single (expensive) GPU deploy. + """ + endpoint_name = random_suffix_name("whisperx", 50) + model_name = endpoint_name + role_arn = aws_session.resolve_role_arn(SAGEMAKER_ROLE) + + LOGGER.info(f"Using image: {image_uri}") + + model = endpoint_config = endpoint = None + try: + LOGGER.info(f"Creating model: {model_name}") + model = Model.create( + model_name=model_name, + primary_container=ContainerDefinition(image=image_uri), + execution_role_arn=role_arn, + ) + + LOGGER.info(f"Creating endpoint config: {endpoint_name}") + endpoint_config = EndpointConfig.create( + endpoint_config_name=endpoint_name, + production_variants=[ + ProductionVariant( + variant_name="AllTraffic", + model_name=model_name, + initial_instance_count=1, + instance_type=INSTANCE_TYPE, + # CUDA 12.8 image -> must use the CU12 (driver 550) AMI. + inference_ami_version=INFERENCE_AMI_VERSION_CU12, + container_startup_health_check_timeout_in_seconds=STARTUP_HEALTH_CHECK_TIMEOUT, + ), + ], + ) + + LOGGER.info(f"Deploying endpoint: {endpoint_name} (this may take 10-15 minutes)...") + endpoint = Endpoint.create( + endpoint_name=endpoint_name, + endpoint_config_name=endpoint_name, + ) + endpoint.wait_for_status("InService", timeout=DEPLOY_TIMEOUT) + LOGGER.info("Endpoint deployment completed successfully") + + yield endpoint + finally: + _cleanup([endpoint, endpoint_config, model]) + + +def test_endpoint_in_service(model_endpoint, audio_cache): + """Endpoint reaches InService and /invocations answers a minimal request with JSON. + + Reaching InService (the fixture) is the primary assertion; the minimal invoke + confirms the SageMaker /invocations route is wired and returns a parseable 200. + """ + endpoint = model_endpoint + audio = audio_cache(AUDIO_EN) + + body = _invoke_transcription(endpoint, audio, language="en") + assert isinstance(body, dict), f"Expected JSON object from /invocations, got: {body!r}" + LOGGER.info("Endpoint InService and /invocations returned a parseable response") + + +def test_invocations_transcription(model_endpoint, audio_cache): + """/invocations transcribes English audio and returns non-empty text (default json).""" + endpoint = model_endpoint + audio = audio_cache(AUDIO_EN) + + body = _invoke_transcription(endpoint, audio, language="en") + text = body.get("text", "") + assert isinstance(text, str) and text.strip(), f"Transcription text is empty: {body!r}" + LOGGER.info(f"Transcription text (len={len(text)}): {text[:120]!r}") + + +def test_invocations_diarization(model_endpoint, audio_cache): + """/invocations with diarize=true + verbose_json surfaces >=2 distinct speakers.""" + endpoint = model_endpoint + audio = audio_cache(AUDIO_DIARIZE) + + body = _invoke_transcription( + endpoint, + audio, + language="en", + response_format="verbose_json", + diarize="true", + ) + assert body.get("text", "").strip(), f"Diarization response has empty text: {body!r}" + speakers = body.get("speakers") + assert speakers, f"verbose_json missing non-empty 'speakers': {body!r}" + assert len(speakers) >= 2, f"Expected >=2 speakers, got {speakers!r}" + LOGGER.info(f"Diarization detected {len(speakers)} speakers: {speakers}") diff --git a/test/whisperx/sagemaker/test_sm_model_config.py b/test/whisperx/sagemaker/test_sm_model_config.py new file mode 100644 index 000000000000..3290c5984681 --- /dev/null +++ b/test/whisperx/sagemaker/test_sm_model_config.py @@ -0,0 +1,223 @@ +"""WhisperX model-config integration test — proves env config flows through SageMaker. + +Deploys a SECOND real-time endpoint whose ``ContainerDefinition.environment`` sets +``WHISPERX_DEFAULT_MODEL=tiny`` and verifies the container actually honors it. This +is the whole point of the suite: an env var set on the SageMaker Model must reach +the running container and change which Whisper model it pins and serves. + +We prove the pin took effect purely through the SageMaker ``/invocations`` route +(the only inference surface SageMaker exposes — there is no GET /v1/models over the +invoke API). Per-request model selection was removed, so the request `model` field +is an ignored no-op; the container always serves the one launched model. We prove +the env override reached the container via two positive invocations: + * a request with no `model` returns 200 non-empty text, which is only possible if + `tiny` was warm-loaded and is being served; and + * a request naming `large-v2` (the server's built-in default) is IGNORED and ALSO + returns 200 non-empty text — the launched `tiny` model serves regardless. + +Kept separate from ``test_sm_endpoint.py`` so the custom-env deploy never touches +the default endpoint. ``tiny`` also warm-loads fast, keeping deploy + health check +quick. AMI pin (DESIGN decision 9): CUDA 12.8 image -> INFERENCE_AMI_VERSION_CU12 +(AL2, driver 550); the default CU13 AMI causes a zero-log CannotStartContainerError. + +Validation is deterministic: assert served-model behavior and non-empty text, never +an exact transcript (WhisperX output is not byte-stable). +""" + +import json +import logging +import os +import time +import uuid + +import pytest +from sagemaker.core.resources import Endpoint, EndpointConfig, Model +from sagemaker.core.shapes import ContainerDefinition, ProductionVariant +from test_utils import random_suffix_name +from test_utils.constants import INFERENCE_AMI_VERSION_CU12, SAGEMAKER_ROLE + +LOGGER = logging.getLogger(__name__) +LOGGER.setLevel(logging.INFO) + +# Single L4 GPU; matches WhisperX VRAM needs. +INSTANCE_TYPE = "ml.g6.xlarge" +# Startup warms the pinned Whisper model + diarization pipeline before /ping +# returns 200, so InService already implies the (tiny) model is resident. +STARTUP_HEALTH_CHECK_TIMEOUT = 900 +DEPLOY_TIMEOUT = 1800 + +# Audio fixture pre-staged in S3 (read with the test runner's credentials). +AUDIO_BUCKET = "dlc-cicd-models" +AUDIO_PREFIX = "test-fixtures/audio" +AUDIO_EN = "asr_en.wav" # English ~15s + +# Env override under test: pin the container to the smallest Whisper model. +CUSTOM_MODEL = "tiny" +# The server's built-in default when WHISPERX_DEFAULT_MODEL is unset. A different id; +# naming it must be IGNORED (model selection was removed) and still served by `tiny`. +NORMAL_DEFAULT_MODEL = "large-v2" + + +def _cleanup(resources): + """Best-effort delete for a list of v3 resource objects (None-safe). + + Leaked SageMaker endpoints keep billing, so cleanup must run even when a + test fails — callers invoke this from a `finally` block. + """ + for resource in resources: + if resource is None: + continue + try: + resource.delete() + except Exception as e: + LOGGER.warning(f"Cleanup {type(resource).__name__} failed: {e}") + + +def _build_multipart(audio_path: str, fields: dict) -> tuple[bytes, str]: + """Build a multipart/form-data body with the audio `file` part + string fields. + + SageMaker InvokeEndpoint forwards the ContentType header (boundary included) + straight through to the model server, so the client owns the encoding. + """ + boundary = uuid.uuid4().hex + filename = os.path.basename(audio_path) + with open(audio_path, "rb") as f: + audio = f.read() + + parts = [ + f'--{boundary}\r\nContent-Disposition: form-data; name="file"; ' + f'filename="{filename}"\r\nContent-Type: audio/wav\r\n\r\n'.encode(), + audio, + b"\r\n", + ] + for key, value in fields.items(): + parts.append( + f'--{boundary}\r\nContent-Disposition: form-data; name="{key}"\r\n\r\n' + f"{value}\r\n".encode() + ) + parts.append(f"--{boundary}--\r\n".encode()) + + body = b"".join(parts) + content_type = f"multipart/form-data; boundary={boundary}" + return body, content_type + + +def _invoke_transcription(endpoint, audio_path: str, retries: int = 3, **fields) -> dict: + """Invoke /invocations with a multipart audio request and return the parsed JSON. + + A bounded retry rides out SageMaker's 60s real-time invoke cap on a cold first + request deterministically rather than flaking. + """ + body, content_type = _build_multipart(audio_path, fields) + last_error = None + for attempt in range(retries): + try: + result = endpoint.invoke(body=body, content_type=content_type) + return json.loads(result.body.read()) + except Exception as e: + last_error = e + LOGGER.warning(f"Invoke attempt {attempt + 1}/{retries} failed: {e}") + if attempt < retries - 1: + time.sleep(30) + raise last_error + + +@pytest.fixture(scope="module") +def audio_en(aws_session, tmp_path_factory): + """Download the English audio fixture once for the module.""" + dest = os.path.join(str(tmp_path_factory.mktemp("whisperx-audio")), AUDIO_EN) + LOGGER.info(f"Downloading s3://{AUDIO_BUCKET}/{AUDIO_PREFIX}/{AUDIO_EN}") + aws_session.s3.download_file(AUDIO_BUCKET, f"{AUDIO_PREFIX}/{AUDIO_EN}", dest) + return dest + + +@pytest.fixture(scope="module") +def custom_model_endpoint(aws_session, image_uri): + """Deploy one WhisperX endpoint pinned to `tiny` via env; clean up after. + + Modeled on ``test_sm_endpoint.py``'s ``model_endpoint`` fixture, but the + ContainerDefinition carries ``environment={"WHISPERX_DEFAULT_MODEL": "tiny"}`` + so the running container serves `tiny` instead of the built-in `large-v2`. + Module-scoped so the (expensive) GPU deploy is shared across this file's tests. + """ + endpoint_name = random_suffix_name("whisperx-cfg", 50) + model_name = endpoint_name + role_arn = aws_session.resolve_role_arn(SAGEMAKER_ROLE) + + LOGGER.info(f"Using image: {image_uri}") + + model = endpoint_config = endpoint = None + try: + LOGGER.info(f"Creating model: {model_name} (WHISPERX_DEFAULT_MODEL={CUSTOM_MODEL})") + model = Model.create( + model_name=model_name, + primary_container=ContainerDefinition( + image=image_uri, + environment={"WHISPERX_DEFAULT_MODEL": CUSTOM_MODEL}, + ), + execution_role_arn=role_arn, + ) + + LOGGER.info(f"Creating endpoint config: {endpoint_name}") + endpoint_config = EndpointConfig.create( + endpoint_config_name=endpoint_name, + production_variants=[ + ProductionVariant( + variant_name="AllTraffic", + model_name=model_name, + initial_instance_count=1, + instance_type=INSTANCE_TYPE, + # CUDA 12.8 image -> must use the CU12 (driver 550) AMI. + inference_ami_version=INFERENCE_AMI_VERSION_CU12, + container_startup_health_check_timeout_in_seconds=STARTUP_HEALTH_CHECK_TIMEOUT, + ), + ], + ) + + LOGGER.info(f"Deploying endpoint: {endpoint_name} (this may take 10-15 minutes)...") + endpoint = Endpoint.create( + endpoint_name=endpoint_name, + endpoint_config_name=endpoint_name, + ) + endpoint.wait_for_status("InService", timeout=DEPLOY_TIMEOUT) + LOGGER.info("Endpoint deployment completed successfully") + + yield endpoint + finally: + _cleanup([endpoint, endpoint_config, model]) + + +def test_custom_default_model_served(custom_model_endpoint, audio_en): + """WHISPERX_DEFAULT_MODEL=tiny propagates through SageMaker to the container. + + Per-request model selection was removed, so the served model can't be probed via + a rejected id anymore. Two positive invocations over the SageMaker /invocations + route prove the env override reached the running container: + + 1. A request with no `model` returns 200 with non-empty text — only possible if + `tiny` warm-loaded and is being served. + 2. A request naming `large-v2` (the server's built-in default) is IGNORED and + ALSO returns 200 with non-empty text — the launched `tiny` model serves + regardless, demonstrating the request `model` field is a no-op. + """ + endpoint = custom_model_endpoint + + # (1) Served-model check: no `model` field -> the launched (tiny) model serves. + body = _invoke_transcription(endpoint, audio_en, language="en") + text = body.get("text", "") + assert isinstance(text, str) and text.strip(), ( + f"Default (tiny) model returned empty text; is tiny actually served? {body!r}" + ) + LOGGER.info(f"tiny model served; transcription text (len={len(text)}): {text[:120]!r}") + + # (2) Ignored-field check: naming `large-v2` is a no-op — the launched `tiny` + # model still serves and returns 200 non-empty text. With model selection + # removed the served id can't be probed via a rejected id, so InService plus + # this non-empty transcription is the proof the env var propagated. + body = _invoke_transcription(endpoint, audio_en, model=NORMAL_DEFAULT_MODEL, language="en") + text = body.get("text", "") + assert isinstance(text, str) and text.strip(), ( + f"Ignored `{NORMAL_DEFAULT_MODEL}` request returned empty text; " + f"is the launched tiny model serving regardless? {body!r}" + ) + LOGGER.info(f"`{NORMAL_DEFAULT_MODEL}` ignored; tiny still served (len={len(text)})") diff --git a/test/whisperx/test_server_options.py b/test/whisperx/test_server_options.py new file mode 100644 index 000000000000..407939707289 --- /dev/null +++ b/test/whisperx/test_server_options.py @@ -0,0 +1,335 @@ +"""Unit tests for the WhisperX server's launch-time option builders. + +These run on CPU with no container and no GPU: torch / whisperx / fastapi are +stubbed in ``sys.modules`` before ``server.py`` is imported, so importing the +server never pulls in CUDA, the whisperX wheel, or a real ASGI app. We then +exercise the pure ``env -> dict`` helpers (``_build_asr_options`` / +``_build_vad_options`` / ``_env_bool``) that translate WHISPERX_* launch env +into ``whisperx.load_model`` kwargs, plus the module-level constants they feed. +A few tests go further and assert those constants are actually forwarded into +the ``whisperx`` calls: ``_get_align`` pins the aligner via ``model_name`` and +``_transcribe`` passes ``task`` (and no stray decoding kwargs) to transcribe(). +""" + +import importlib.util +import sys +import types +from pathlib import Path + +import pytest + +# server.py lives in the image-build tree, a sibling of test/; load it by path. +SERVER_PATH = Path(__file__).resolve().parents[2] / "scripts" / "docker" / "whisperx" / "server.py" + +# Every WHISPERX_* var the option builders read. Cleared before each test so a +# stray value in the runner's environment can't skew the "defaults" assertions. +_WHISPERX_ENV = ( + "WHISPERX_TEMPERATURE", + "WHISPERX_TEMPERATURE_INCREMENT_ON_FALLBACK", + "WHISPERX_BEAM_SIZE", + "WHISPERX_BEST_OF", + "WHISPERX_PATIENCE", + "WHISPERX_LENGTH_PENALTY", + "WHISPERX_COMPRESSION_RATIO_THRESHOLD", + "WHISPERX_LOGPROB_THRESHOLD", + "WHISPERX_NO_SPEECH_THRESHOLD", + "WHISPERX_CONDITION_ON_PREVIOUS_TEXT", + "WHISPERX_INITIAL_PROMPT", + "WHISPERX_HOTWORDS", + "WHISPERX_SUPPRESS_TOKENS", + "WHISPERX_SUPPRESS_NUMERALS", + "WHISPERX_CHUNK_SIZE", + "WHISPERX_VAD_ONSET", + "WHISPERX_VAD_OFFSET", + "WHISPERX_VAD_METHOD", + "WHISPERX_TASK", + "WHISPERX_ALIGN_MODEL", +) + + +def _install_stubs() -> None: + """Register lightweight fakes for the heavy imports server.py performs. + + server.py runs ``torch.cuda.is_available()`` and builds a FastAPI app at + import time; none of that is relevant to the option builders, so we swap the + modules for minimal stand-ins that satisfy the import and the decorators. + """ + + class _FakeApp: + def __init__(self, *args, **kwargs): + pass + + def get(self, *args, **kwargs): + return lambda fn: fn + + def post(self, *args, **kwargs): + return lambda fn: fn + + def _form_marker(default=None, **kwargs): + return default + + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(is_available=lambda: False) + + fastapi = types.ModuleType("fastapi") + fastapi.FastAPI = _FakeApp + fastapi.File = _form_marker + fastapi.Form = _form_marker + fastapi.UploadFile = type("UploadFile", (), {}) + fastapi.HTTPException = type("HTTPException", (Exception,), {}) + + fastapi_responses = types.ModuleType("fastapi.responses") + fastapi_responses.JSONResponse = type("JSONResponse", (), {}) + fastapi_responses.PlainTextResponse = type("PlainTextResponse", (), {}) + + whisperx = types.ModuleType("whisperx") + whisperx.load_model = lambda *a, **k: None + whisperx.load_audio = lambda *a, **k: None + whisperx.load_align_model = lambda *a, **k: (None, {}) + whisperx.align = lambda *a, **k: {} + whisperx.assign_word_speakers = lambda *a, **k: {} + + whisperx_diarize = types.ModuleType("whisperx.diarize") + whisperx_diarize.DiarizationPipeline = type("DiarizationPipeline", (), {}) + whisperx.diarize = whisperx_diarize + + # server.py imports WriteSRT/WriteVTT from whisperx.utils at module scope; + # stub the submodule so the import resolves (these tests never render subtitles). + whisperx_utils = types.ModuleType("whisperx.utils") + whisperx_utils.WriteSRT = type("WriteSRT", (), {}) + whisperx_utils.WriteVTT = type("WriteVTT", (), {}) + whisperx.utils = whisperx_utils + + sys.modules.update( + { + "torch": torch, + "fastapi": fastapi, + "fastapi.responses": fastapi_responses, + "whisperx": whisperx, + "whisperx.diarize": whisperx_diarize, + "whisperx.utils": whisperx_utils, + } + ) + + +def _load_server(): + """Import server.py fresh (stubs installed) so it re-reads WHISPERX_* env.""" + _install_stubs() + spec = importlib.util.spec_from_file_location("whisperx_server_under_test", SERVER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +server = _load_server() + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + """Start every test from a WHISPERX_*-free environment.""" + for name in _WHISPERX_ENV: + monkeypatch.delenv(name, raising=False) + + +def test_asr_options_defaults(): + assert server._build_asr_options() == { + "beam_size": 5, + "best_of": 5, + "patience": 1.0, + "length_penalty": 1.0, + "temperatures": (0.0, 0.2, 0.4, 0.6, 0.8, 1.0), + "compression_ratio_threshold": 2.4, + "log_prob_threshold": -1.0, + "no_speech_threshold": 0.6, + "condition_on_previous_text": False, + "initial_prompt": None, + "hotwords": None, + "suppress_tokens": [-1], + "suppress_numerals": False, + } + + +def test_vad_options_defaults(): + assert server._build_vad_options() == { + "chunk_size": 30, + "vad_onset": 0.5, + "vad_offset": 0.363, + } + + +def test_beam_size_override(monkeypatch): + monkeypatch.setenv("WHISPERX_BEAM_SIZE", "8") + assert server._build_asr_options()["beam_size"] == 8 + + +def test_vad_onset_override(monkeypatch): + monkeypatch.setenv("WHISPERX_VAD_ONSET", "0.7") + assert server._build_vad_options()["vad_onset"] == 0.7 + + +def test_suppress_tokens_csv(monkeypatch): + monkeypatch.setenv("WHISPERX_SUPPRESS_TOKENS", "1,2,3") + assert server._build_asr_options()["suppress_tokens"] == [1, 2, 3] + + +def test_suppress_tokens_ignores_blanks(monkeypatch): + monkeypatch.setenv("WHISPERX_SUPPRESS_TOKENS", "1, ,3,") + assert server._build_asr_options()["suppress_tokens"] == [1, 3] + + +def test_suppress_numerals_truthy(monkeypatch): + monkeypatch.setenv("WHISPERX_SUPPRESS_NUMERALS", "true") + assert server._build_asr_options()["suppress_numerals"] is True + + +def test_temperature_fallback_schedule(monkeypatch): + monkeypatch.setenv("WHISPERX_TEMPERATURE", "0.4") + # WHISPERX_TEMPERATURE_INCREMENT_ON_FALLBACK defaults to 0.2. + assert server._build_asr_options()["temperatures"] == (0.4, 0.6, 0.8, 1.0) + + +def test_temperature_increment_disabled(monkeypatch): + monkeypatch.setenv("WHISPERX_TEMPERATURE", "0.4") + monkeypatch.setenv("WHISPERX_TEMPERATURE_INCREMENT_ON_FALLBACK", "0") + assert server._build_asr_options()["temperatures"] == (0.4,) + + +def test_initial_prompt_set(monkeypatch): + monkeypatch.setenv("WHISPERX_INITIAL_PROMPT", "medical dictation") + assert server._build_asr_options()["initial_prompt"] == "medical dictation" + + +def test_initial_prompt_empty_is_none(monkeypatch): + monkeypatch.setenv("WHISPERX_INITIAL_PROMPT", "") + assert server._build_asr_options()["initial_prompt"] is None + + +def test_condition_on_previous_text_toggle(monkeypatch): + monkeypatch.setenv("WHISPERX_CONDITION_ON_PREVIOUS_TEXT", "yes") + assert server._build_asr_options()["condition_on_previous_text"] is True + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + ("1", True), + ("true", True), + ("YES", True), + ("on", True), + ("0", False), + ("false", False), + ("", False), + ("nonsense", False), + ], +) +def test_env_bool(monkeypatch, value, expected): + monkeypatch.setenv("WHISPERX_SUPPRESS_NUMERALS", value) + assert server._env_bool("WHISPERX_SUPPRESS_NUMERALS") is expected + + +def test_env_bool_default_when_unset(): + assert server._env_bool("WHISPERX_DOES_NOT_EXIST") is False + assert server._env_bool("WHISPERX_DOES_NOT_EXIST", default=True) is True + + +def test_module_constants_reflect_env(monkeypatch): + """The module-level ASR/VAD constants are built from env at import time.""" + monkeypatch.setenv("WHISPERX_BEAM_SIZE", "9") + monkeypatch.setenv("WHISPERX_CHUNK_SIZE", "12") + reloaded = _load_server() + assert reloaded.ASR_OPTIONS["beam_size"] == 9 + assert reloaded.VAD_OPTIONS["chunk_size"] == 12 + + +def test_pipeline_constants_override(monkeypatch): + monkeypatch.setenv("WHISPERX_VAD_METHOD", "silero") + monkeypatch.setenv("WHISPERX_TASK", "translate") + monkeypatch.setenv("WHISPERX_ALIGN_MODEL", "WAV2VEC2_ASR_LARGE_LV60K_960H") + reloaded = _load_server() + assert reloaded.VAD_METHOD == "silero" + assert reloaded.TASK == "translate" + assert reloaded.ALIGN_MODEL == "WAV2VEC2_ASR_LARGE_LV60K_960H" + + +def test_pipeline_constants_defaults(monkeypatch): + # Reload under a cleaned env instead of trusting the collection-time module: + # a runner with WHISPERX_TASK/VAD_METHOD/ALIGN_MODEL exported must not break + # the defaults assertions. + for name in ("WHISPERX_VAD_METHOD", "WHISPERX_TASK", "WHISPERX_ALIGN_MODEL"): + monkeypatch.delenv(name, raising=False) + reloaded = _load_server() + assert reloaded.VAD_METHOD == "pyannote" + assert reloaded.TASK == "transcribe" + assert reloaded.ALIGN_MODEL is None + + +def test_get_align_forwards_default_model_name(monkeypatch): + """Unset WHISPERX_ALIGN_MODEL => _get_align passes model_name=None (default aligner).""" + monkeypatch.delenv("WHISPERX_ALIGN_MODEL", raising=False) + reloaded = _load_server() + assert reloaded.ALIGN_MODEL is None + + recorded: dict = {} + + def _fake_load_align_model(*args, **kwargs): + recorded.update(kwargs) + return object(), {} + + reloaded.whisperx.load_align_model = _fake_load_align_model + reloaded._get_align("en") + assert recorded["model_name"] is None + + +def test_get_align_forwards_pinned_model_name(monkeypatch): + """WHISPERX_ALIGN_MODEL set => _get_align forwards it as model_name to pin the aligner.""" + monkeypatch.setenv("WHISPERX_ALIGN_MODEL", "WAV2VEC2_ASR_LARGE_LV60K_960H") + reloaded = _load_server() + assert reloaded.ALIGN_MODEL == "WAV2VEC2_ASR_LARGE_LV60K_960H" + + recorded: dict = {} + + def _fake_load_align_model(*args, **kwargs): + recorded.update(kwargs) + return object(), {} + + reloaded.whisperx.load_align_model = _fake_load_align_model + # Fresh reload => empty _ALIGN_LRU, so the stub is actually invoked. + reloaded._get_align("de") + assert recorded["model_name"] == "WAV2VEC2_ASR_LARGE_LV60K_960H" + + +def test_transcribe_forwards_task_and_no_decoding_kwargs(monkeypatch): + """Regression for the inert-knob bug: transcribe() receives only pipeline kwargs. + + ``task`` must reach ``FasterWhisperPipeline.transcribe`` (load_model discards + it without a pinned language), and decoding kwargs baked into ASR_OPTIONS + (initial_prompt/temperature/...) must NOT leak into the per-call kwargs. + """ + monkeypatch.setenv("WHISPERX_TASK", "translate") + reloaded = _load_server() + + recorded: dict = {} + + class _FakeModel: + def transcribe(self, audio, **kwargs): + recorded.update(kwargs) + return {"segments": [{"text": "hi", "start": 0.0, "end": 1.0}], "language": "en"} + + reloaded.whisperx.load_audio = lambda path: [0.0] * 16000 + reloaded._get_whisper = lambda name: _FakeModel() + + result = reloaded._transcribe( + audio_path="/tmp/does-not-exist.wav", + language=None, + want_words=False, + diarize=False, + min_speakers=None, + max_speakers=None, + ) + + assert set(recorded).issubset({"batch_size", "chunk_size", "language", "task"}) + assert "initial_prompt" not in recorded + assert "temperature" not in recorded + assert recorded["task"] == "translate" + assert recorded["task"] == reloaded.TASK + assert result["segments"] # sanity: the pipeline still produced output diff --git a/test/whisperx/test_subtitle_format.py b/test/whisperx/test_subtitle_format.py new file mode 100644 index 000000000000..1a98e1fa1e70 --- /dev/null +++ b/test/whisperx/test_subtitle_format.py @@ -0,0 +1,372 @@ +"""Unit tests for the WhisperX server's subtitle line-formatting support. + +Like ``test_transcribe_behavior.py`` these run on CPU with no container and no +GPU: torch / whisperx / fastapi are stubbed in ``sys.modules`` before +``server.py`` is imported. server.py now does ``from whisperx.utils import +WriteSRT, WriteVTT``, so the stub adds a ``whisperx.utils`` module whose +WriteSRT / WriteVTT are *recording* fakes: they capture the ``(result, +options)`` handed to ``write_result`` and emit a sentinel string. + +The actual line-wrapping / word-highlighting algorithm is WhisperX's own +``SubtitlesWriter`` — not ours — so we do NOT reimplement or assert on it (that +would only test a fake). Instead we test the logic we wrote: + + * ``_render_subtitle`` builds the correct writer class for the format and + passes an options dict carrying ALL THREE keys with the exact values. + * ``_format_response`` routes srt/vtt through ``_render_subtitle`` only when a + subtitle knob is set, and through the legacy ``_to_srt`` / ``_to_vtt`` + otherwise (so default srt/vtt output is unchanged), and json/text ignore the + knobs entirely (no error). + * ``_handle_transcription`` forces word-level alignment (want_words=True) when + a subtitle knob is set on srt/vtt, and leaves want_words as derived + otherwise (OR, never overwrite). + +The real WhisperX writers emit outputs that satisfy the existing EC2 format +assertions — WriteVTT prints ``"WEBVTT\\n"`` first (whisperx/utils.py:377) and +WriteSRT prints ``-->`` (whisperx/utils.py:391), verified against whisperx +3.8.6 source. ``_format_response`` returns that output verbatim (proven here via +the sentinel passthrough), so the knob path keeps satisfying those assertions. +""" + +import asyncio +import importlib.util +import sys +import types +from pathlib import Path + +# server.py lives in the image-build tree, a sibling of test/; load it by path. +SERVER_PATH = Path(__file__).resolve().parents[2] / "scripts" / "docker" / "whisperx" / "server.py" + +_SRT_SENTINEL = "<>" +_VTT_SENTINEL = "<>" + + +def _make_recording_writer(name: str, sentinel: str): + """Build a fake SubtitlesWriter that records calls and writes ``sentinel``. + + Records construction ``output_dir`` and every ``write_result(result, + options)`` on class-level lists. Fresh classes are created per stub install, + so the lists start empty for each server import. + """ + + class _RecordingWriter: + writer_name = name + init_output_dirs: list = [] + write_calls: list = [] + + def __init__(self, output_dir=None, **kwargs): + type(self).init_output_dirs.append(output_dir) + + def write_result(self, result, file, options): + type(self).write_calls.append({"result": result, "options": options}) + file.write(sentinel) + + return _RecordingWriter + + +def _install_stubs() -> None: + """Register lightweight fakes for the heavy imports server.py performs. + + Same shape as test_transcribe_behavior.py, plus a ``whisperx.utils`` module + whose WriteSRT / WriteVTT are recording fakes. The response stubs also + capture ``media_type`` so the srt/vtt content-type routing is checkable. + """ + + class _FakeApp: + def __init__(self, *args, **kwargs): + pass + + def get(self, *args, **kwargs): + return lambda fn: fn + + def post(self, *args, **kwargs): + return lambda fn: fn + + def _form_marker(default=None, **kwargs): + return default + + class _HTTPException(Exception): + def __init__(self, status_code=None, detail=None, **kwargs): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + class _JSONResponse: + def __init__(self, content=None, **kwargs): + self.content = content + self.media_type = kwargs.get("media_type") + + class _PlainTextResponse: + def __init__(self, content=None, **kwargs): + self.content = content + self.media_type = kwargs.get("media_type") + + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(is_available=lambda: False) + + fastapi = types.ModuleType("fastapi") + fastapi.FastAPI = _FakeApp + fastapi.File = _form_marker + fastapi.Form = _form_marker + fastapi.UploadFile = type("UploadFile", (), {}) + fastapi.HTTPException = _HTTPException + + fastapi_responses = types.ModuleType("fastapi.responses") + fastapi_responses.JSONResponse = _JSONResponse + fastapi_responses.PlainTextResponse = _PlainTextResponse + + whisperx = types.ModuleType("whisperx") + whisperx.load_model = lambda *a, **k: None + whisperx.load_audio = lambda *a, **k: None + whisperx.load_align_model = lambda *a, **k: (None, {}) + whisperx.align = lambda *a, **k: {} + whisperx.assign_word_speakers = lambda *a, **k: {} + + whisperx_diarize = types.ModuleType("whisperx.diarize") + whisperx_diarize.DiarizationPipeline = type("DiarizationPipeline", (), {}) + whisperx.diarize = whisperx_diarize + + whisperx_utils = types.ModuleType("whisperx.utils") + whisperx_utils.WriteSRT = _make_recording_writer("WriteSRT", _SRT_SENTINEL) + whisperx_utils.WriteVTT = _make_recording_writer("WriteVTT", _VTT_SENTINEL) + whisperx.utils = whisperx_utils + + sys.modules.update( + { + "torch": torch, + "fastapi": fastapi, + "fastapi.responses": fastapi_responses, + "whisperx": whisperx, + "whisperx.diarize": whisperx_diarize, + "whisperx.utils": whisperx_utils, + } + ) + + +def _load_server(): + """Import server.py fresh (stubs installed) so each test gets empty LRUs.""" + _install_stubs() + spec = importlib.util.spec_from_file_location("whisperx_server_under_test", SERVER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _result() -> dict: + """A minimal _transcribe-shaped result carrying language + segments + words.""" + return { + "task": "transcribe", + "language": "en", + "duration": 1.0, + "text": "hello world", + "segments": [ + { + "text": "hello world", + "start": 0.0, + "end": 1.0, + "words": [ + {"word": "hello", "start": 0.0, "end": 0.5}, + {"word": "world", "start": 0.5, "end": 1.0}, + ], + } + ], + "words": [ + {"word": "hello", "start": 0.0, "end": 0.5}, + {"word": "world", "start": 0.5, "end": 1.0}, + ], + "speakers": None, + } + + +# --------------------------------------------------------------------------- +# _render_subtitle — writer selection + options assembly +# --------------------------------------------------------------------------- +def test_render_subtitle_srt_passes_all_three_options(): + """srt picks WriteSRT and forwards all three option keys with exact values.""" + server = _load_server() + result = _result() + out = server._render_subtitle(result, "srt", 42, 2, True) + + assert out == _SRT_SENTINEL # verbatim writer output flows back + assert server.WriteVTT.write_calls == [] # wrong writer never used + assert len(server.WriteSRT.write_calls) == 1 + call = server.WriteSRT.write_calls[0] + assert call["options"] == { + "max_line_width": 42, + "max_line_count": 2, + "highlight_words": True, + } + assert call["result"] is result # full result dict (has language + segments) + assert server.WriteSRT.init_output_dirs == [""] # output_dir unused by write_result + + +def test_render_subtitle_vtt_selects_vtt_writer_with_none_options(): + """vtt picks WriteVTT; None/None/False knobs still yield all three keys.""" + server = _load_server() + out = server._render_subtitle(_result(), "vtt", None, None, False) + + assert out == _VTT_SENTINEL + assert server.WriteSRT.write_calls == [] + assert len(server.WriteVTT.write_calls) == 1 + assert server.WriteVTT.write_calls[0]["options"] == { + "max_line_width": None, + "max_line_count": None, + "highlight_words": False, + } + + +# --------------------------------------------------------------------------- +# _format_response — routing between the WhisperX writer and legacy formatters +# --------------------------------------------------------------------------- +def test_format_response_srt_with_knob_uses_render_subtitle(): + """A knob on srt routes through _render_subtitle (WhisperX writer).""" + server = _load_server() + resp = server._format_response( + _result(), "srt", want_words=True, diarize=False, max_line_width=42 + ) + assert len(server.WriteSRT.write_calls) == 1 + assert resp.content == _SRT_SENTINEL + assert resp.media_type == "application/x-subrip" + + +def test_format_response_vtt_with_knob_uses_render_subtitle(): + """A knob on vtt routes through _render_subtitle with the vtt media type.""" + server = _load_server() + resp = server._format_response( + _result(), "vtt", want_words=True, diarize=False, highlight_words=True + ) + assert len(server.WriteVTT.write_calls) == 1 + assert resp.content == _VTT_SENTINEL + assert resp.media_type == "text/vtt" + + +def test_format_response_srt_without_knob_uses_legacy(): + """No knob on srt keeps the legacy _to_srt output (WhisperX writer untouched).""" + server = _load_server() + resp = server._format_response(_result(), "srt", want_words=False, diarize=False) + assert server.WriteSRT.write_calls == [] # WhisperX writer not invoked + assert "-->" in resp.content # existing EC2 assertion still holds + assert resp.media_type == "application/x-subrip" + + +def test_format_response_vtt_without_knob_uses_legacy(): + """No knob on vtt keeps the legacy _to_vtt output starting with WEBVTT.""" + server = _load_server() + resp = server._format_response(_result(), "vtt", want_words=False, diarize=False) + assert server.WriteVTT.write_calls == [] + assert resp.content.startswith("WEBVTT") # existing EC2 assertion still holds + assert resp.media_type == "text/vtt" + + +def test_format_response_json_ignores_knobs(): + """json + knobs returns normally; the WhisperX writer is never invoked.""" + server = _load_server() + resp = server._format_response( + _result(), + "json", + want_words=False, + diarize=False, + max_line_width=42, + max_line_count=2, + highlight_words=True, + ) + assert server.WriteSRT.write_calls == [] + assert server.WriteVTT.write_calls == [] + assert resp.content == {"text": "hello world"} + + +def test_format_response_text_ignores_knobs(): + """text + knobs returns the plain transcript; no WhisperX writer call.""" + server = _load_server() + resp = server._format_response( + _result(), "text", want_words=False, diarize=False, highlight_words=True + ) + assert server.WriteSRT.write_calls == [] + assert server.WriteVTT.write_calls == [] + assert resp.content == "hello world" + + +# --------------------------------------------------------------------------- +# _handle_transcription — subtitle knob forces alignment (want_words) on srt/vtt +# --------------------------------------------------------------------------- +class _FakeUpload: + filename = "audio.wav" + + async def read(self): + return b"RIFFfake-audio-bytes" + + +def _run_handle(server, *, response_format, timestamp_granularities=None, **knobs): + """Drive _handle_transcription with _transcribe stubbed to record want_words.""" + recorded: dict = {} + + def _recorder(**kwargs): + recorded["want_words"] = kwargs["want_words"] + return _result() + + server._transcribe = _recorder + resp = asyncio.run( + server._handle_transcription( + file=_FakeUpload(), + language=None, + response_format=response_format, + timestamp_granularities=timestamp_granularities, + diarize=False, + min_speakers=None, + max_speakers=None, + **knobs, + ) + ) + return recorded, resp + + +def test_handle_transcription_srt_knob_forces_alignment(): + """srt + max_line_width forces want_words=True even without granularities[]=word.""" + server = _load_server() + recorded, resp = _run_handle(server, response_format="srt", max_line_width=42) + assert recorded["want_words"] is True + assert resp.content == _SRT_SENTINEL # routed through the WhisperX writer + + +def test_handle_transcription_vtt_highlight_forces_alignment(): + """vtt + highlight_words forces alignment too.""" + server = _load_server() + recorded, _ = _run_handle(server, response_format="vtt", highlight_words=True) + assert recorded["want_words"] is True + + +def test_handle_transcription_srt_no_knob_does_not_force_alignment(): + """srt with no knob leaves want_words as derived (False here).""" + server = _load_server() + recorded, resp = _run_handle(server, response_format="srt") + assert recorded["want_words"] is False + assert server.WriteSRT.write_calls == [] # legacy formatter path + + +def test_handle_transcription_json_knob_does_not_force_alignment(): + """json + max_line_width: knob ignored, no alignment forced, no error.""" + server = _load_server() + recorded, resp = _run_handle(server, response_format="json", max_line_width=42) + assert recorded["want_words"] is False + assert resp.content == {"text": "hello world"} # returns normally + assert server.WriteSRT.write_calls == [] + assert server.WriteVTT.write_calls == [] + + +def test_handle_transcription_word_granularity_preserved_with_knob(): + """subtitle_formatting ORs into want_words; explicit word granularity is kept.""" + server = _load_server() + recorded, _ = _run_handle( + server, + response_format="srt", + timestamp_granularities=["word"], + max_line_width=42, + ) + assert recorded["want_words"] is True + + +def test_handle_transcription_word_granularity_without_knob(): + """granularities[]=word still forces want_words even when no subtitle knob is set.""" + server = _load_server() + recorded, _ = _run_handle(server, response_format="json", timestamp_granularities=["word"]) + assert recorded["want_words"] is True diff --git a/test/whisperx/test_transcribe_behavior.py b/test/whisperx/test_transcribe_behavior.py new file mode 100644 index 000000000000..739006bbcc40 --- /dev/null +++ b/test/whisperx/test_transcribe_behavior.py @@ -0,0 +1,562 @@ +"""Unit tests for the WhisperX server's transcription behavior. + +Like ``test_server_options.py`` these run on CPU with no container and no GPU: +torch / whisperx / fastapi are stubbed in ``sys.modules`` before ``server.py`` +is imported (``anyio`` and ``functools`` are the real thing — ``anyio`` ships +with fastapi/starlette). We then exercise two behaviors that are pure control +flow and need no model: + + * FIX B — the alignment fallback only degrades on the exceptions whisperX + raises for a genuinely unsupported language (ValueError / NotImplementedError + / KeyError). Any other error (e.g. a CUDA-OOM RuntimeError) propagates as a + real failure instead of a silent HTTP 200, and an explicit ``diarize=true`` + that degrades becomes a 422 rather than a speaker-less 200. + * FIX A — the blocking ``_transcribe`` call is offloaded off the event loop + via ``anyio.to_thread.run_sync`` so long transcriptions don't starve /ping. + +Later review findings covered here (all GPU-free control flow): + + * FINDING #1 — GPU inference is serialized via ``_INFERENCE_LIMITER`` + (``WHISPERX_MAX_CONCURRENT_REQUESTS``, default 1), so two concurrent + ``_handle_transcription`` calls never run ``_transcribe`` at the same time. + * FINDING #3 — the returned dict's ``task`` reflects ``WHISPERX_TASK``. + * FINDING #4 — admission control: a full inference queue is shed with 503 + before the body is read, and an oversized upload is rejected with 413. + * FINDING #5 — passive readiness: ``/ping`` reports 503 once the inference + path observes a fatal CUDA fault; transient OOM leaves it healthy. + * FINDING #6 — extension params are validated at the boundary (422). +""" + +import asyncio +import importlib.util +import inspect +import sys +import threading +import time +import types +from pathlib import Path + +import pytest + +# server.py lives in the image-build tree, a sibling of test/; load it by path. +SERVER_PATH = Path(__file__).resolve().parents[2] / "scripts" / "docker" / "whisperx" / "server.py" + + +def _install_stubs() -> None: + """Register lightweight fakes for the heavy imports server.py performs. + + Same shape as test_server_options.py, but the fastapi fakes carry real + behavior we assert on: HTTPException stores ``status_code``/``detail`` (so + the 422 path is checkable) and the response classes accept their payload. + """ + + class _FakeApp: + def __init__(self, *args, **kwargs): + pass + + def get(self, *args, **kwargs): + return lambda fn: fn + + def post(self, *args, **kwargs): + return lambda fn: fn + + def _form_marker(default=None, **kwargs): + return default + + class _HTTPException(Exception): + def __init__(self, status_code=None, detail=None, **kwargs): + super().__init__(detail) + self.status_code = status_code + self.detail = detail + + class _JSONResponse: + def __init__(self, content=None, **kwargs): + self.content = content + self.status_code = kwargs.get("status_code", 200) + self.media_type = kwargs.get("media_type") + + class _PlainTextResponse: + def __init__(self, content=None, **kwargs): + self.content = content + self.media_type = kwargs.get("media_type") + + torch = types.ModuleType("torch") + torch.cuda = types.SimpleNamespace(is_available=lambda: False) + + fastapi = types.ModuleType("fastapi") + fastapi.FastAPI = _FakeApp + fastapi.File = _form_marker + fastapi.Form = _form_marker + fastapi.UploadFile = type("UploadFile", (), {}) + fastapi.HTTPException = _HTTPException + + fastapi_responses = types.ModuleType("fastapi.responses") + fastapi_responses.JSONResponse = _JSONResponse + fastapi_responses.PlainTextResponse = _PlainTextResponse + + whisperx = types.ModuleType("whisperx") + whisperx.load_model = lambda *a, **k: None + whisperx.load_audio = lambda *a, **k: None + whisperx.load_align_model = lambda *a, **k: (None, {}) + whisperx.align = lambda *a, **k: {} + whisperx.assign_word_speakers = lambda *a, **k: {} + + whisperx_diarize = types.ModuleType("whisperx.diarize") + whisperx_diarize.DiarizationPipeline = type("DiarizationPipeline", (), {}) + whisperx.diarize = whisperx_diarize + + # server.py imports WriteSRT/WriteVTT from whisperx.utils at module scope; + # stub the submodule so the import resolves (these tests never render subtitles). + whisperx_utils = types.ModuleType("whisperx.utils") + whisperx_utils.WriteSRT = type("WriteSRT", (), {}) + whisperx_utils.WriteVTT = type("WriteVTT", (), {}) + whisperx.utils = whisperx_utils + + sys.modules.update( + { + "torch": torch, + "fastapi": fastapi, + "fastapi.responses": fastapi_responses, + "whisperx": whisperx, + "whisperx.diarize": whisperx_diarize, + "whisperx.utils": whisperx_utils, + } + ) + + +def _load_server(): + """Import server.py fresh (stubs installed) so each test gets empty LRUs.""" + _install_stubs() + spec = importlib.util.spec_from_file_location("whisperx_server_under_test", SERVER_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class _FakeModel: + """Stand-in Whisper model: transcribe() yields one Spanish segment.""" + + def transcribe(self, audio, **kwargs): + return {"segments": [{"text": "hola", "start": 0.0, "end": 1.0}], "language": "es"} + + +def _server_with_align_error(exc: Exception): + """Load server with load_audio/_get_whisper stubbed and whisperx.align raising ``exc``.""" + server = _load_server() + server.whisperx.load_audio = lambda path: [0.0] * 16000 + server._get_whisper = lambda name: _FakeModel() + + def _raise(*args, **kwargs): + raise exc + + server.whisperx.align = _raise + return server + + +def _transcribe(server, *, want_words: bool, diarize: bool): + return server._transcribe( + audio_path="/tmp/does-not-exist.wav", + language=None, + want_words=want_words, + diarize=diarize, + min_speakers=None, + max_speakers=None, + ) + + +# --------------------------------------------------------------------------- +# FIX B — narrowed alignment except +# --------------------------------------------------------------------------- +def test_alignment_runtime_error_propagates(): + """A non-degradation error (CUDA OOM) must bubble up, not become a 200.""" + server = _server_with_align_error(RuntimeError("CUDA OOM")) + with pytest.raises(RuntimeError, match="CUDA OOM"): + _transcribe(server, want_words=True, diarize=False) + + +def test_unsupported_language_degrades_for_want_words(): + """A genuine no-aligner ValueError degrades a want_words request to segments.""" + server = _server_with_align_error(ValueError("No default align-model for language: xx")) + result = _transcribe(server, want_words=True, diarize=False) + assert isinstance(result, dict) + assert result["segments"] == [{"text": "hola", "start": 0.0, "end": 1.0}] + assert result["words"] is None # degraded: no word timings, but no exception + assert result["speakers"] is None + assert result["language"] == "es" + + +def test_diarize_requested_but_alignment_unavailable_raises_422(): + """An explicit diarize=true that can't be aligned is a 422, not a silent 200.""" + server = _server_with_align_error(ValueError("No default align-model for language: xx")) + with pytest.raises(server.HTTPException) as exc_info: + _transcribe(server, want_words=False, diarize=True) + assert exc_info.value.status_code == 422 + assert "diarization was requested" in exc_info.value.detail + + +def test_notimplementederror_also_degrades(): + """NotImplementedError (unsupported align model type) is a degrade case too.""" + server = _server_with_align_error(NotImplementedError("Align model of type ... not supported")) + result = _transcribe(server, want_words=True, diarize=False) + assert result["words"] is None + assert result["segments"] + + +# --------------------------------------------------------------------------- +# FIX A — offload blocking work off the event loop +# --------------------------------------------------------------------------- +class _FakeUpload: + filename = "audio.wav" + + async def read(self): + return b"RIFFfake-audio-bytes" + + +def test_handle_transcription_is_coroutine_and_anyio_wired(): + """Smoke check that the offload path is present: async handler + anyio import.""" + server = _load_server() + assert inspect.iscoroutinefunction(server._handle_transcription) + assert hasattr(server, "anyio") + assert "anyio" in sys.modules + + +def test_transcribe_runs_in_worker_thread(): + """The blocking _transcribe must run on a different thread than the loop.""" + server = _load_server() + main_thread_id = threading.get_ident() + recorded: dict = {} + + def _recorder(**kwargs): + recorded["thread_id"] = threading.get_ident() + recorded["kwargs"] = kwargs + return {"text": "hi", "segments": [{"text": "hi", "start": 0.0, "end": 1.0}]} + + server._transcribe = _recorder + + asyncio.run( + server._handle_transcription( + file=_FakeUpload(), + language=None, + response_format="json", + timestamp_granularities=None, + diarize=False, + min_speakers=None, + max_speakers=None, + ) + ) + + assert recorded["thread_id"] != main_thread_id + assert recorded["kwargs"]["audio_path"] # tempfile path was forwarded through + + +# --------------------------------------------------------------------------- +# FIX C — cache getters dedupe concurrent loads under the worker-thread pool +# --------------------------------------------------------------------------- +def _run_concurrently(target, n=8): + """Fire ``n`` threads at ``target`` at once; return once all have joined.""" + barrier = threading.Barrier(n) + + def _worker(): + barrier.wait() # release all threads together to widen the race window + target() + + threads = [threading.Thread(target=_worker) for _ in range(n)] + for t in threads: + t.start() + for t in threads: + t.join() + + +def test_get_whisper_dedupes_concurrent_loads(): + """Concurrent cold misses must trigger exactly one model load (the lock). + + One model per container: _get_whisper caches a single resident model. Without + _WHISPER_LOCK, N worker threads racing the first request would all see + _WHISPER_MODEL is None before any store completes and each call load_model — + N redundant ~30 s loads. The lock serializes check-load-store so only the + first loads (defense-in-depth if MAX_CONCURRENT_REQUESTS>1). + """ + server = _load_server() + calls = {"count": 0} + lock = threading.Lock() + sentinel = object() + + def _fake_load_model(*args, **kwargs): + with lock: + calls["count"] += 1 + time.sleep(0.05) # widen the check-then-act window + return sentinel + + server.whisperx.load_model = _fake_load_model + + results = [] + results_lock = threading.Lock() + + def _call(): + model = server._get_whisper("large-v2") + with results_lock: + results.append(model) + + _run_concurrently(_call, n=8) + + assert calls["count"] == 1 # lock deduped 8 concurrent misses to one load + assert len(results) == 8 + assert all(m is sentinel for m in results) + + +def _handle(server, upload=None, **overrides): + """Drive _handle_transcription with sane defaults; ``overrides`` win.""" + kwargs = { + "file": upload or _FakeUpload(), + "language": None, + "response_format": "json", + "timestamp_granularities": None, + "diarize": False, + "min_speakers": None, + "max_speakers": None, + } + kwargs.update(overrides) + return asyncio.run(server._handle_transcription(**kwargs)) + + +# --------------------------------------------------------------------------- +# FINDING #2 — one model per container: the `model` field left the request path +# --------------------------------------------------------------------------- +def test_model_field_removed_from_request_path(): + """`model` is gone from every request-path signature, and so is override.""" + server = _load_server() + for fn in (server._handle_transcription, server.transcribe, server.invocations): + assert "model" not in inspect.signature(fn).parameters + assert "model_name" not in inspect.signature(server._transcribe).parameters + # The override machinery is fully deleted. + assert not hasattr(server, "_resolve_model") + assert not hasattr(server, "ALLOW_MODEL_OVERRIDE") + + +# --------------------------------------------------------------------------- +# FINDING #3 — the returned dict's task reflects WHISPERX_TASK, not a literal +# --------------------------------------------------------------------------- +def test_transcribe_task_reflects_env(monkeypatch): + """WHISPERX_TASK=translate => result['task'] == 'translate' (was hardcoded).""" + monkeypatch.setenv("WHISPERX_TASK", "translate") + server = _load_server() + server.whisperx.load_audio = lambda path: [0.0] * 16000 + server._get_whisper = lambda name: _FakeModel() + + result = _transcribe(server, want_words=False, diarize=False) + assert server.TASK == "translate" + assert result["task"] == "translate" + + +# --------------------------------------------------------------------------- +# FINDING #1 — GPU inference is serialized to one at a time (default capacity 1) +# --------------------------------------------------------------------------- +def test_inference_serialized_to_capacity_one(): + """Two concurrent handlers must never run _transcribe at the same time. + + The stubbed _transcribe tracks a live-count with a lock and sleeps; if the + CapacityLimiter (default 1) serializes correctly the max observed overlap is + exactly 1. Without it the two worker threads would both be inside at once. + """ + server = _load_server() + live = {"cur": 0, "max": 0} + lock = threading.Lock() + + def _blocking(**kwargs): + with lock: + live["cur"] += 1 + live["max"] = max(live["max"], live["cur"]) + time.sleep(0.05) # hold the slot so an overlap would be observable + with lock: + live["cur"] -= 1 + return {"text": "hi", "segments": [{"text": "hi", "start": 0.0, "end": 1.0}]} + + server._transcribe = _blocking + + async def _drive(): + await asyncio.gather(_run_one(), _run_one()) + + async def _run_one(): + await server._handle_transcription( + file=_FakeUpload(), + language=None, + response_format="json", + timestamp_granularities=None, + diarize=False, + min_speakers=None, + max_speakers=None, + ) + + asyncio.run(_drive()) + assert live["max"] == 1 # the limiter prevented any overlap + + +# --------------------------------------------------------------------------- +# FINDING #4 — admission control: shed a full queue (503) and cap upload (413) +# --------------------------------------------------------------------------- +def test_full_queue_returns_503_before_reading_body(): + """tasks_waiting >= MAX_QUEUE => 503 raised BEFORE the upload is read.""" + server = _load_server() + + class _FakeLimiter: + def statistics(self): + return types.SimpleNamespace(tasks_waiting=99) + + server._INFERENCE_LIMITER = _FakeLimiter() + + read_called = {"v": False} + + class _WatchedUpload: + filename = "audio.wav" + + async def read(self): + read_called["v"] = True + return b"RIFFfake-audio-bytes" + + with pytest.raises(server.HTTPException) as exc_info: + _handle(server, upload=_WatchedUpload()) + assert exc_info.value.status_code == 503 + assert read_called["v"] is False # body never read under overload + + +def test_oversized_upload_returns_413(monkeypatch): + """A body larger than WHISPERX_MAX_UPLOAD_BYTES is rejected 413.""" + monkeypatch.setenv("WHISPERX_MAX_UPLOAD_BYTES", "8") + server = _load_server() + assert server.MAX_UPLOAD_BYTES == 8 + + class _BigUpload: + filename = "audio.wav" + + async def read(self): + return b"x" * 64 # exceeds the 8-byte cap + + with pytest.raises(server.HTTPException) as exc_info: + _handle(server, upload=_BigUpload()) + assert exc_info.value.status_code == 413 + + +# --------------------------------------------------------------------------- +# FINDING #5 — passive readiness health reflected by /ping +# --------------------------------------------------------------------------- +def test_ping_healthy_returns_200(): + server = _load_server() + assert server._HEALTHY is True + resp = server.ping() + assert resp.status_code == 200 + assert resp.content == {"status": "ok"} + + +def test_ping_unhealthy_returns_503(): + server = _load_server() + server._HEALTHY = False + resp = server.ping() + assert resp.status_code == 503 + assert resp.content == {"status": "unavailable"} + + +def test_is_fatal_cuda_error_classification(): + server = _load_server() + assert server._is_fatal_cuda_error(RuntimeError("CUDA error: device-side assert triggered")) + assert server._is_fatal_cuda_error(RuntimeError("out of memory")) is False + assert server._is_fatal_cuda_error(ValueError("x")) is False + + +def test_fatal_cuda_error_flips_health_and_reraises(): + """A fatal CUDA fault in inference flips _HEALTHY False and still errors.""" + server = _load_server() + server.whisperx.load_audio = lambda path: [0.0] * 16000 + + class _BoomModel: + def transcribe(self, audio, **kwargs): + raise RuntimeError("CUDA error: illegal memory access was encountered") + + server._get_whisper = lambda name: _BoomModel() + assert server._HEALTHY is True + with pytest.raises(RuntimeError, match="illegal memory access"): + _transcribe(server, want_words=False, diarize=False) + assert server._HEALTHY is False + + +def test_transient_oom_does_not_flip_health(): + """A plain CUDA OOM is transient: it propagates but leaves readiness intact.""" + server = _load_server() + server.whisperx.load_audio = lambda path: [0.0] * 16000 + + class _OomModel: + def transcribe(self, audio, **kwargs): + raise RuntimeError("CUDA out of memory") + + server._get_whisper = lambda name: _OomModel() + with pytest.raises(RuntimeError, match="out of memory"): + _transcribe(server, want_words=False, diarize=False) + assert server._HEALTHY is True + + +# --------------------------------------------------------------------------- +# FINDING #6 — parameter validation at the boundary (422) +# --------------------------------------------------------------------------- +@pytest.mark.parametrize( + "overrides", + [ + {"min_speakers": 0}, + {"max_speakers": 0}, + {"min_speakers": 3, "max_speakers": 2}, + {"max_line_width": -1}, + {"max_line_count": -1}, + ], +) +def test_invalid_params_return_422(overrides): + server = _load_server() + + def _must_not_run(**kwargs): + raise AssertionError("_transcribe ran despite invalid params") + + server._transcribe = _must_not_run + with pytest.raises(server.HTTPException) as exc_info: + _handle(server, **overrides) + assert exc_info.value.status_code == 422 + + +def test_valid_params_reach_transcribe(): + """A valid speaker/line combo passes validation and reaches _transcribe.""" + server = _load_server() + reached = {"v": False} + + def _recorder(**kwargs): + reached["v"] = True + return {"text": "hi", "segments": [{"text": "hi", "start": 0.0, "end": 1.0}]} + + server._transcribe = _recorder + _handle(server, min_speakers=1, max_speakers=2, max_line_width=40, max_line_count=2) + assert reached["v"] is True + + +def test_get_align_dedupes_concurrent_loads(): + """Same dedupe guarantee for the align cache, which is keyed by language.""" + server = _load_server() + calls = {"count": 0} + lock = threading.Lock() + sentinel = object() + + def _fake_load_align_model(*args, **kwargs): + with lock: + calls["count"] += 1 + time.sleep(0.05) # widen the check-then-act window + return (sentinel, {}) + + server.whisperx.load_align_model = _fake_load_align_model + + results = [] + results_lock = threading.Lock() + + def _call(): + model, metadata = server._get_align("es") # one language across threads + with results_lock: + results.append(model) + + _run_concurrently(_call, n=8) + + assert calls["count"] == 1 # lock deduped 8 concurrent misses to one load + assert len(results) == 8 + assert all(m is sentinel for m in results)