name: 'Docker Test Image Build'
description: 'Build and push a Dynamo test image from an existing runtime image base'

# Builds container/Dockerfile.test on top of a pre-built runtime image.
# Unlike docker-remote-build this action has no sccache wiring, no build metrics,
# and always uses the "test-<framework>-..." cache namespace.

inputs:
  image_tag:
    description: 'Full image URI to tag the test image with'
    required: true
  base_image:
    description: 'Full image URI of the runtime image to use as BASE_IMAGE'
    required: true
  framework:
    description: 'Framework name (vllm, sglang, trtllm) — used to derive the cache tag'
    required: true
  platform:
    description: 'Target platform architecture (amd64 or arm64)'
    required: true
  cuda_version:
    description: 'CUDA major.minor version (e.g. 12.9) — major version used in cache tag'
    required: true
  aws_account_id:
    description: 'AWS account ID for ECR hostname'
    required: true
  aws_default_region:
    description: 'AWS region for ECR hostname'
    required: true
  push_image:
    description: 'Push the image to the registry'
    required: false
    default: 'false'
  no_load:
    description: 'Do not load the image into the local Docker daemon'
    required: false
    default: 'true'
  no_cache:
    description: 'Disable Docker layer cache'
    required: false
    default: 'false'

runs:
  using: "composite"
  steps:
    - name: Build and push test image
      shell: bash
      env:
        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
      run: |
        CUDA_MAJOR=${{ inputs.cuda_version }}
        CUDA_MAJOR=${CUDA_MAJOR%%.*}
        CACHE_TAG="test-${{ inputs.framework }}-cuda${CUDA_MAJOR}-${{ inputs.platform }}-cache"
        CACHE_ARGS="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${CACHE_TAG}"
        CACHE_ARGS+=" --cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:release-${CACHE_TAG}"
        if [[ "$GITHUB_REF_NAME" == "main" ]]; then
          CACHE_ARGS+=" --cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${CACHE_TAG},mode=max"
        elif [[ "$GITHUB_REF_NAME" =~ ^release ]]; then
          CACHE_ARGS+=" --cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:release-${CACHE_TAG},mode=max"
        fi

        PUSH_ARGS=""
        if [ "${{ inputs.push_image }}" == "true" ]; then
          PUSH_ARGS="--push"
        elif [ "${{ inputs.no_load }}" == "false" ]; then
          PUSH_ARGS="--load"
        fi

        NO_CACHE_ARG=""
        if [ "${{ inputs.no_cache }}" == "true" ]; then
          NO_CACHE_ARG="--no-cache"
        fi

        docker buildx build \
          --progress=plain \
          ${PUSH_ARGS} \
          ${NO_CACHE_ARG} \
          --platform "linux/${{ inputs.platform }}" \
          -f container/Dockerfile.test \
          --build-arg "BASE_IMAGE=${{ inputs.base_image }}" \
          ${CACHE_ARGS} \
          -t "${{ inputs.image_tag }}" .