feat: Add enable-local-kvindexer into MDC

.github/actions/docker-build/action.yml

@@ -16,9 +16,6 @@ inputs:
   image_tag:
     description: 'Custom image tag (optional, defaults to framework:latest)'
     required: false
-  ngc_ci_access_token:
-    description: 'NGC CI Access Token'
-    required: false
   ci_token:
     description: 'CI Token'
     required: false
@@ -49,6 +46,12 @@ inputs:
   torch_backend:
     description: 'Optional override for TORCH_BACKEND build-arg (e.g., cu129)'
     required: false
+  enable_kvbm:
+    description: 'Enable KVBM support (optional)'
+    required: false
+  dynamo_base_image:
+    description: 'Pre-built Dynamo base image to use instead of building from scratch'
+    required: false
 
 outputs:
   image_tag:
@@ -61,20 +64,9 @@ runs:
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@e468171a9de216ec08956ac3ada2f0791b6bd435 #v3.11.1
       with:
-        driver: docker
+        driver: docker-container
         # Enable BuildKit for enhanced metadata
         buildkitd-flags: --debug
-    - name: Login to ECR
-      shell: bash
-      env:
-        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
-      run: |
-        aws ecr get-login-password --region ${{ inputs.aws_default_region }} | docker login --username AWS --password-stdin ${ECR_HOSTNAME}
-    - name: Login to NGC
-      if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name == 'push'
-      shell: bash
-      run: |
-        echo "${{ inputs.ngc_ci_access_token }}" | docker login nvcr.io -u '$oauthtoken' --password-stdin
     - name: Cleanup
       if: always()
       shell: bash
@@ -90,9 +82,12 @@ runs:
         AWS_ACCESS_KEY_ID: ${{ inputs.aws_access_key_id }}
         AWS_SECRET_ACCESS_KEY: ${{ inputs.aws_secret_access_key }}
         PLATFORM: ${{ inputs.platform }}
+        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
         GITHUB_RUN_ID: ${{ github.run_id }}
         GITHUB_JOB: ${{ github.job }}
+        GITHUB_REF_NAME: ${{ github.ref_name }}
       run: |
+        set -x
         # Determine image tag
         if [ -n "${{ inputs.image_tag }}" ]; then
           IMAGE_TAG="${{ inputs.image_tag }}"
@@ -112,18 +107,34 @@ runs:
         echo "📝 Build log will be saved to: ${BUILD_LOG_FILE}"
 
         # Collect optional overrides provided by the workflow
+        # Set base cache args and set --cache-to if this is a main commit
         EXTRA_ARGS=""
+        EXTRA_ARGS="--cache-to type=inline "
+        EXTRA_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:${{ inputs.framework }}-${PLATFORM##*/}-cache "
+        EXTRA_ARGS+="--cache-from type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:main-${{ inputs.framework }}-${PLATFORM##*/} "
+        if [[ "$GITHUB_REF_NAME" == "main" ]]; then
+          EXTRA_ARGS+="--cache-to type=registry,ref=${ECR_HOSTNAME}/ai-dynamo/dynamo:${{ inputs.framework }}-${PLATFORM##*/}-cache,mode=max "
+        fi
+
+        echo "$EXTRA_ARGS"
+        # Collect optional overrides provided by the workflow
         if [ -n "${{ inputs.base_image_tag }}" ]; then
-          EXTRA_ARGS+=" --base-image-tag ${{ inputs.base_image_tag }}"
+          EXTRA_ARGS+="--base-image-tag ${{ inputs.base_image_tag }} "
         fi
         if [ -n "${{ inputs.runtime_image_tag }}" ]; then
-          EXTRA_ARGS+=" --build-arg RUNTIME_IMAGE_TAG=${{ inputs.runtime_image_tag }}"
+          EXTRA_ARGS+="--build-arg RUNTIME_IMAGE_TAG=${{ inputs.runtime_image_tag }} "
         fi
         if [ -n "${{ inputs.cuda_version }}" ]; then
-          EXTRA_ARGS+=" --build-arg CUDA_VERSION=${{ inputs.cuda_version }}"
+          EXTRA_ARGS+="--build-arg CUDA_VERSION=${{ inputs.cuda_version }} "
         fi
         if [ -n "${{ inputs.torch_backend }}" ]; then
-          EXTRA_ARGS+=" --build-arg TORCH_BACKEND=${{ inputs.torch_backend }}"
+          EXTRA_ARGS+="--build-arg TORCH_BACKEND=${{ inputs.torch_backend }} "
+        fi
+        if [ -n "${{ inputs.dynamo_base_image }}" ]; then
+          EXTRA_ARGS+=" --dynamo-base-image ${{ inputs.dynamo_base_image }}"
+        fi
+        if [ -n "${{ inputs.enable_kvbm }}" ]; then
+          EXTRA_ARGS+=" --build-arg ENABLE_KVBM=${{ inputs.enable_kvbm }}"
         fi
 
         # Execute build and capture output (show on console AND save to file)
@@ -144,6 +155,26 @@ runs:
         # Exit with the build's exit code
         exit ${BUILD_EXIT_CODE}
 
+    - name: Run Sanity Check on Runtime Image
+      if: inputs.target == 'runtime'
+      shell: bash
+      run: |
+        IMAGE_TAG="${{ steps.build.outputs.image_tag }}"
+        echo "Running sanity check on image: $IMAGE_TAG"
+
+        # Run the sanity check script inside the container
+        # The script is located in /workspace/deploy/sanity_check.py in runtime containers
+        set +e
+        docker run --rm "$IMAGE_TAG" python /workspace/deploy/sanity_check.py --runtime-check --no-gpu-check
+        SANITY_CHECK_EXIT_CODE=$?
+        set -e
+        if [ ${SANITY_CHECK_EXIT_CODE} -ne 0 ]; then
+          echo "ERROR: Sanity check failed - ai-dynamo packages not properly installed"
+          exit ${SANITY_CHECK_EXIT_CODE}
+        else
+          echo "✅ Sanity check passed"
+        fi
+
     - name: Capture Build Metrics
       id: metrics
       shell: bash
@@ -223,8 +254,7 @@ runs:
         chmod +x .github/scripts/parse_buildkit_output.py
 
         # Check for build logs and build stage arguments dynamically
-        BASE_BUILD_LOG="build-logs/base-image-build.log"
-        FRAMEWORK_BUILD_LOG="build-logs/framework-${FRAMEWORK_LOWER}-build.log"
+        BUILD_LOG="build-logs/single-stage-build.log"
 
         # Path to container metadata created in previous step
         CONTAINER_METADATA="build-metrics/metrics-${{ inputs.framework }}-${PLATFORM_ARCH}-${WORKFLOW_ID}-${JOB_ID}.json"
@@ -237,18 +267,11 @@ runs:
         # Build stage arguments dynamically based on which logs exist
         STAGE_ARGS=()
 
-        if [ -f "$BASE_BUILD_LOG" ]; then
-          echo "  ✓ Found base image log: ${BASE_BUILD_LOG}"
-          STAGE_ARGS+=("base:${BASE_BUILD_LOG}")
-        else
-          echo "  ℹ️  No base image log found"
-        fi
-
-        if [ -f "$FRAMEWORK_BUILD_LOG" ]; then
-          echo "  ✓ Found framework log: ${FRAMEWORK_BUILD_LOG}"
-          STAGE_ARGS+=("runtime:${FRAMEWORK_BUILD_LOG}")
+        if [ -f "$BUILD_LOG" ]; then
+          echo "  ✓ Found base image log: ${BUILD_LOG}"
+          STAGE_ARGS+=("runtime:${BUILD_LOG}")
         else
-          echo "  ℹ️  No framework log found"
+          echo "  ℹ️  No image log found"
         fi
 
         # Check for any additional stage logs (e.g., build-logs/stage3-*.log)
@@ -280,13 +303,6 @@ runs:
         if [ ${PARSER_EXIT_CODE} -eq 0 ] && [ -f "$COMPREHENSIVE_JSON" ]; then
           echo "✅ Comprehensive build metrics generated successfully"
           echo "📄 Output file: ${COMPREHENSIVE_JSON}"
-          echo ""
-          echo "=========================================="
-          echo "📋 FULL JSON OUTPUT (for debugging)"
-          echo "=========================================="
-          cat "$COMPREHENSIVE_JSON"
-          echo ""
-          echo "=========================================="
         else
           echo "⚠️  Metrics generation had issues but continuing..."
         fi
@@ -296,7 +312,7 @@ runs:
       uses: actions/upload-artifact@v4
       if: always()
       with:
-        name: build-metrics-${{ inputs.framework }}-${{ env.PLATFORM_ARCH }}-${{ github.run_id }}-${{ job.check_run_id }}
+        name: build-metrics-${{ inputs.framework }}-${{ inputs.target }}-${{ env.PLATFORM_ARCH }}-${{ github.run_id }}-${{ job.check_run_id }}
         path: build-metrics/build-${{ inputs.framework }}-${{ env.PLATFORM_ARCH }}-${{ github.run_id }}-${{ job.check_run_id }}.json
         retention-days: 7
 

.github/actions/docker-login/action.yml

@@ -0,0 +1,46 @@
+name: 'Docker Login'
+description: 'Login to multiple container registries (ECR, NGC, ACR)'
+
+inputs:
+  ngc_ci_access_token:
+    description: 'NGC CI Access Token'
+    required: false
+  aws_default_region:
+    description: 'AWS Default Region'
+    required: false
+  aws_account_id:
+    description: 'AWS Account ID'
+    required: false
+  azure_acr_hostname:
+    description: 'Azure ACR hostname'
+    required: false
+  azure_acr_user:
+    description: 'Azure ACR user'
+    required: false
+  azure_acr_password:
+    description: 'Azure ACR password'
+    required: false
+
+runs:
+  using: "composite"
+  steps:
+    - name: ECR Login
+      shell: bash
+      if: ${{ inputs.aws_default_region != '' && inputs.aws_account_id != '' }}
+      env:
+        ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
+      run: |
+        set -euo pipefail
+        aws ecr get-login-password --region ${{ inputs.aws_default_region }} | docker login --username AWS --password-stdin "${ECR_HOSTNAME}"
+    - name: NGC Login
+      if: ${{ inputs.ngc_ci_access_token != '' }}
+      shell: bash
+      run: |
+        set -euo pipefail
+        echo "${{ inputs.ngc_ci_access_token }}" | docker login nvcr.io -u '$oauthtoken' --password-stdin
+    - name: ACR Login
+      shell: bash
+      if: ${{ inputs.azure_acr_hostname != '' && inputs.azure_acr_user != '' && inputs.azure_acr_password != '' }}
+      run: |
+        set -euo pipefail
+        echo "${{ inputs.azure_acr_password }}" | docker login "${{ inputs.azure_acr_hostname }}" --username "${{ inputs.azure_acr_user }}" --password-stdin

.github/actions/docker-tag-push/action.yml

@@ -1,12 +1,18 @@
+name: 'Docker Tag and Push'
 description: 'Tag and Push Docker Images'
 
 inputs:
   local_image:
     description: 'Local Image Name:Tag'
     required: true
-  push_tag:
-    description: 'Target Name:Tag'
+  push_tags:
+    description: 'Target Name:Tag (newline-separated list for multiple tags)'
     required: true
+  # There isn't a clean way to have an additional tag that is conditional
+  # Adding this to handle this use-case (we want multiple tags for main builds)
+  conditional_tag:
+    description: 'Optional tag for conditionals'
+    required: false
   aws_push:
     description: 'Push to AWS Boolean'
     required: false
@@ -21,54 +27,56 @@ inputs:
   aws_default_region:
     description: 'AWS Default Region'
     required: false
-  aws_access_key_id:
-    description: 'AWS Access Key ID'
-    required: false
-  aws_secret_access_key:
-    description: 'AWS Secret Access Key'
-    required: false
   azure_acr_hostname:
     description: 'Azure ACR hostname'
     required: false
-  azure_acr_user:
-    description: 'Azure ACR user'
-    required: false
-  azure_acr_password:
-    description: 'Azure ACR password'
-    required: false
 
 outputs:
-  image_tag:
-    description: 'Image Tag'
-    value: ${{ inputs.push_tag }}
+  image_tags:
+    description: 'Image Tags'
+    value: ${{ inputs.push_tags }}
 
 runs:
   using: "composite"
   steps:
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
-    - name: ACR Login
-      shell: bash
-      if: ${{ inputs.azure_push == 'true' }}
-      run: |
-        echo "${{ inputs.azure_acr_password }}" | docker login ${{ inputs.azure_acr_hostname }} --username ${{ inputs.azure_acr_user }} --password-stdin
     - name: ECR Tag and Push
       shell: bash
       if: ${{ inputs.aws_push == 'true' }}
       env:
         LOCAL_IMAGE: ${{ inputs.local_image }}
-        PUSH_TAG: ${{ inputs.push_tag }}
+        PUSH_TAGS: ${{ inputs.push_tags }}
+        CONDITIONAL_TAG: ${{ inputs.conditional_tag }}
         ECR_HOSTNAME: ${{ inputs.aws_account_id }}.dkr.ecr.${{ inputs.aws_default_region }}.amazonaws.com
       run: |
-        docker tag ${LOCAL_IMAGE} ${ECR_HOSTNAME}/${PUSH_TAG}
-        docker push ${ECR_HOSTNAME}/${PUSH_TAG}
+        set -euo pipefail
+        if [[ ${CONDITIONAL_TAG} != '' ]]; then
+            docker tag ${LOCAL_IMAGE} ${ECR_HOSTNAME}/${CONDITIONAL_TAG}
+            docker push ${ECR_HOSTNAME}/${CONDITIONAL_TAG}
+        fi
+        while IFS= read -r TAG; do
+          if [ -z "$TAG" ]; then
+            continue
+          fi
+          echo "Tagging and pushing: ${ECR_HOSTNAME}/${TAG}"
+          docker tag "${LOCAL_IMAGE}" "${ECR_HOSTNAME}/${TAG}"
+          docker push "${ECR_HOSTNAME}/${TAG}"
+        done <<< "$PUSH_TAGS"
     - name: ACR Tag and Push
       shell: bash
       if: ${{ inputs.azure_push == 'true' }}
       env:
         LOCAL_IMAGE: ${{ inputs.local_image }}
-        PUSH_TAG: ${{ inputs.push_tag }}
+        PUSH_TAGS: ${{ inputs.push_tags }}
         AZURE_ACR_HOSTNAME: ${{ inputs.azure_acr_hostname }}
       run: |
-        docker tag ${LOCAL_IMAGE} ${AZURE_ACR_HOSTNAME}/${PUSH_TAG}
-        docker push ${AZURE_ACR_HOSTNAME}/${PUSH_TAG}
+        set -euo pipefail
+        while IFS= read -r TAG; do
+          if [ -z "$TAG" ]; then
+            continue
+          fi
+          echo "Tagging and pushing: ${AZURE_ACR_HOSTNAME}/${TAG}"
+          docker tag "${LOCAL_IMAGE}" "${AZURE_ACR_HOSTNAME}/${TAG}"
+          docker push "${AZURE_ACR_HOSTNAME}/${TAG}"
+        done <<< "$PUSH_TAGS"