# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

name: Inference Python Benchmarks Dataflow

on:
  schedule:
    - cron: '50 3 * * *'
  workflow_dispatch:

#Setting explicit permissions for the action to avoid the default permissions which are `write-all` in case of pull_request_target event
permissions:
  actions: write
  pull-requests: read
  checks: read
  contents: read
  deployments: read
  id-token: none
  issues: read
  discussions: read
  packages: read
  pages: read
  repository-projects: read
  security-events: read
  statuses: read

# This allows a subsequently queued workflow run to interrupt previous runs
concurrency:
  group: '${{ github.workflow }} @ ${{ github.event.issue.number || github.sha || github.head_ref || github.ref }}-${{ github.event.schedule || github.event.comment.id || github.event.sender.login }}'
  cancel-in-progress: true

env:
  DEVELOCITY_ACCESS_KEY: ${{ secrets.DEVELOCITY_ACCESS_KEY }}
  GRADLE_ENTERPRISE_CACHE_USERNAME: ${{ secrets.GE_CACHE_USERNAME }}
  GRADLE_ENTERPRISE_CACHE_PASSWORD: ${{ secrets.GE_CACHE_PASSWORD }}
  INFLUXDB_USER: ${{ secrets.INFLUXDB_USER }}
  INFLUXDB_USER_PASSWORD: ${{ secrets.INFLUXDB_USER_PASSWORD }}

jobs:
  beam_Inference_Python_Benchmarks_Dataflow:
    if: |
      github.event_name == 'workflow_dispatch' ||
      (github.event_name == 'schedule' && github.repository == 'apache/beam') ||
      github.event.comment.body == 'Run Inference Benchmarks'
    runs-on: [self-hosted, ubuntu-20.04, main]
    timeout-minutes: 1000
    name: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
    strategy:
      matrix:
        job_name: ["beam_Inference_Python_Benchmarks_Dataflow"]
        job_phrase: ["Run Inference Benchmarks"]
    steps:
      - uses: actions/checkout@v4
      - name: Setup repository
        uses: ./.github/actions/setup-action
        with:
          comment_phrase: ${{ matrix.job_phrase }}
          github_token: ${{ secrets.GITHUB_TOKEN }}
          github_job: ${{ matrix.job_name }} (${{ matrix.job_phrase }})
      - name: Setup Python environment
        uses: ./.github/actions/setup-environment-action
        with:
          java-version: default
          python-version: '3.10'
      - name: Package Python SDK using Gradle
        run: ./gradlew :sdks:python:sdist -PpythonVersion=3.10
      - name: Configure Docker for Artifact Registry
        run: gcloud auth configure-docker us-docker.pkg.dev
      - name: Prepare test arguments
        uses: ./.github/actions/test-arguments-action
        with:
          test-type: load
          test-language: python
          argument-file-paths: |
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Vision_Classification_Resnet_101.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Imagenet_Classification_Resnet_152.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Language_Modeling_Bert_Base_Uncased.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Language_Modeling_Bert_Large_Uncased.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Imagenet_Classification_Resnet_152_Tesla_T4_GPU.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Streaming_DistilBert_Base_Uncased.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_Pytorch_Sentiment_Batch_DistilBert_Base_Uncased.txt
            ${{ github.workspace }}/.github/workflows/load-tests-pipeline-options/beam_Inference_Python_Benchmarks_Dataflow_VLLM_Gemma_Batch.txt
      # The env variables are created and populated in the test-arguments-action as "<github.job>_test_arguments_<argument_file_paths_index>"
      - name: get current time
        run: echo "NOW_UTC=$(date '+%m%d%H%M%S' --utc)" >> $GITHUB_ENV
      - name: Build VLLM Development Image
        id: build_vllm_image
        uses: ./.github/actions/build-push-docker-action
        with:
          dockerfile_path: 'sdks/python/apache_beam/ml/inference/test_resources/vllm.dockerfile'
          image_name: 'us-docker.pkg.dev/apache-beam-testing/beam-temp/beam-vllm-gpu-base'
          image_tag: ${{ github.sha }}
      - name: Run VLLM Gemma Batch Test
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.vllm_gemma_benchmarks \
            -Prunner=DataflowRunner \
            -PsdkLocationOverride=false \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/vllm_tests_requirements.txt '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_8 }} --mode=batch --job_name=benchmark-tests-vllm-with-gemma-2b-it-batch-${{env.NOW_UTC}} --sdk_container_image=${{ steps.build_vllm_image.outputs.image_url }}'
      - name: run Pytorch Sentiment Streaming using Hugging Face distilbert-base-uncased model
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_sentiment_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_6 }} --mode=streaming --job_name=benchmark-tests-pytorch-sentiment-distilbert-base-uncased-streaming-${{env.NOW_UTC}} --output_table=apache-beam-testing.beam_run_inference.result_sentiment_distilbert_base_uncased' \
      - name: run Pytorch Sentiment Batch using Hugging Face distilbert-base-uncased model
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_sentiment_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_7 }} --mode=batch --job_name=benchmark-tests-pytorch-sentiment-distilbert-base-uncased-batch-${{env.NOW_UTC}} --output_table=apache-beam-testing.beam_run_inference.result_sentiment_distilbert_base_uncased' \
      - name: run Pytorch Vision Classification with Resnet 101
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_image_classification_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_1 }} --job_name=benchmark-tests-pytorch-imagenet-python-101-${{env.NOW_UTC}} --output=gs://temp-storage-for-end-to-end-tests/torch/result_resnet101-${{env.NOW_UTC}}.txt' \
      - name: run Pytorch Imagenet Classification with Resnet 152
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_image_classification_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_2 }} --job_name=benchmark-tests-pytorch-imagenet-python-152-${{env.NOW_UTC}} --output=gs://temp-storage-for-end-to-end-tests/torch/result_resnet152-${{env.NOW_UTC}}.txt' \
      - name: run Pytorch Language Modeling using Hugging Face bert-base-uncased model
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_language_modeling_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_3 }} --job_name=benchmark-tests-pytorch-language-modeling-bert-base-uncased-${{env.NOW_UTC}} --output=gs://temp-storage-for-end-to-end-tests/torch/result_bert_base_uncased-${{env.NOW_UTC}}.txt' \
      - name: run Pytorch Langauge Modeling using Hugging Face bert-large-uncased model
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_language_modeling_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_4 }} --job_name=benchmark-tests-pytorch-language-modeling-bert-large-uncased-${{env.NOW_UTC}} --output=gs://temp-storage-for-end-to-end-tests/torch/result_bert_large_uncased-${{env.NOW_UTC}}.txt' \
      - name: run Pytorch Imagenet Classification with Resnet 152 with Tesla T4 GPU
        uses: ./.github/actions/gradle-command-self-hosted-action
        timeout-minutes: 180
        with:
          gradle-command: :sdks:python:apache_beam:testing:load_tests:run
          arguments: |
            -PloadTest.mainClass=apache_beam.testing.benchmarks.inference.pytorch_image_classification_benchmarks \
            -Prunner=DataflowRunner \
            -PpythonVersion=3.10 \
            -PloadTest.requirementsTxtFile=apache_beam/ml/inference/torch_tests_requirements.txt \
            '-PloadTest.args=${{ env.beam_Inference_Python_Benchmarks_Dataflow_test_arguments_5 }} --job_name=benchmark-tests-pytorch-imagenet-python-gpu-${{env.NOW_UTC}} --output=gs://temp-storage-for-end-to-end-tests/torch/result_resnet152_gpu-${{env.NOW_UTC}}.txt'