From 5a0ad02aeddfc74d4ab551c3ed9ab2e049b67ab6 Mon Sep 17 00:00:00 2001
From: Dimitri Gnidash <dimitri@tenstorrent.com>
Date: Tue, 17 Sep 2024 09:49:07 -0400
Subject: [PATCH]  #12775: Cleanup docker run action    - Remove unnecessary
 arch parameter of the docker run action    - Add a parameter to installing
 the wheel    - Remove LD_LIBRARY env variable since it is not used in the
 wheel environments    - Remove Github username since it is accessible inside
 of the custom Action

---
 .github/actions/docker-run/action.yml               | 13 ++++++++++---
 .../fast-dispatch-build-and-unit-tests.yaml         |  7 +------
 .github/workflows/models-post-commit.yaml           |  7 +------
 .github/workflows/ttnn-post-commit.yaml             |  7 +------
 4 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/.github/actions/docker-run/action.yml b/.github/actions/docker-run/action.yml
index b2a6a28b5cd..1c3178b9feb 100644
--- a/.github/actions/docker-run/action.yml
+++ b/.github/actions/docker-run/action.yml
@@ -9,12 +9,10 @@ inputs:
     description: 'Docker image to run commands in'
     required: false
     default: tt-metalium/ubuntu-20.04-amd64
-  docker_image_arch:
-    description: 'Is this a wormhole.b0, grayskull, or blackhole image'
-    required: true
   docker_username:
     description: docker login username
     required: true
+    default: ${{ github.actor }}
   docker_password:
     description: docker login password
     required: true
@@ -28,6 +26,11 @@ inputs:
     default: |
       -v /dev/hugepages-1G:/dev/hugepages-1G
       --device /dev/tenstorrent
+  install_wheel:
+    description: "Install the wheel that contains all of the Python environment. The artifact needs to be present."
+    type: boolean
+    required: false
+    default: false
 runs:
   using: "composite"
   steps:
@@ -65,4 +68,8 @@ runs:
           cp -r /github_workspace/* /usr/app/
           cd /usr/app/
           rm -rf tt_metal tt_eager
+          if [ ${{ inputs.install_wheel }} ]; then
+            WHEEL_FILENAME=$(ls -1 *.whl)
+            pip3 install --user $WHEEL_FILENAME
+          fi
           ${{ inputs.run_args }}
diff --git a/.github/workflows/fast-dispatch-build-and-unit-tests.yaml b/.github/workflows/fast-dispatch-build-and-unit-tests.yaml
index 73cf9bcc512..b5a85ca92f5 100644
--- a/.github/workflows/fast-dispatch-build-and-unit-tests.yaml
+++ b/.github/workflows/fast-dispatch-build-and-unit-tests.yaml
@@ -57,8 +57,6 @@ jobs:
     name: ${{ matrix.test-group.name }} ${{ inputs.arch }} ${{ inputs.runner-label }}
     env:
       LOGURU_LEVEL: INFO
-      # may not need this
-      LD_LIBRARY_PATH: ${{ github.workspace }}/build/lib
     runs-on:
       - ${{ inputs.runner-label }}
       - cloud-virtual-machine
@@ -72,12 +70,9 @@ jobs:
         timeout-minutes: ${{ inputs.timeout }}
         uses: ./.github/actions/docker-run
         with:
-          docker_username: ${{ github.actor }}
+          install_wheel: true
           docker_password: ${{ secrets.GITHUB_TOKEN }}
-          docker_image_arch: ${{ inputs.arch }}
           run_args: |
-            WHEEL_FILENAME=$(ls -1 *.whl)
-            pip3 install --user $WHEEL_FILENAME
             ${{ matrix.test-group.cmd }}
       - uses: ./.github/actions/slack-report
         if: ${{ failure() }}
diff --git a/.github/workflows/models-post-commit.yaml b/.github/workflows/models-post-commit.yaml
index 938822af6ce..b91d8dd9039 100644
--- a/.github/workflows/models-post-commit.yaml
+++ b/.github/workflows/models-post-commit.yaml
@@ -49,8 +49,6 @@ jobs:
     name: ${{ matrix.test-group.name }} ${{ inputs.arch }} ${{ inputs.runner-label }}
     env:
       LOGURU_LEVEL: INFO
-      # may not need this
-      LD_LIBRARY_PATH: ${{ github.workspace }}/build/lib
     runs-on:
       - ${{ inputs.runner-label }}
       - in-service
@@ -64,12 +62,9 @@ jobs:
         timeout-minutes: ${{ inputs.timeout }}
         uses: ./.github/actions/docker-run
         with:
-          docker_username: ${{ github.actor }}
+          install_wheel: true
           docker_password: ${{ secrets.GITHUB_TOKEN }}
-          docker_image_arch: ${{ inputs.arch }}
           run_args: |
-            WHEEL_FILENAME=$(ls -1 *.whl)
-            pip3 install --user $WHEEL_FILENAME
             source tests/scripts/run_python_model_tests.sh && run_python_model_tests_${{ inputs.arch }}
       - uses: ./.github/actions/slack-report
         if: ${{ failure() }}
diff --git a/.github/workflows/ttnn-post-commit.yaml b/.github/workflows/ttnn-post-commit.yaml
index 5777ac2ae19..5fdf49198d9 100644
--- a/.github/workflows/ttnn-post-commit.yaml
+++ b/.github/workflows/ttnn-post-commit.yaml
@@ -64,8 +64,6 @@ jobs:
     name: ${{ matrix.test-group.name }} ${{ inputs.arch }} ${{ inputs.runner-label }}
     env:
       LOGURU_LEVEL: INFO
-      # may not need this
-      LD_LIBRARY_PATH: ${{ github.workspace }}/build/lib
     runs-on:
       - ${{ inputs.runner-label }}
       - "in-service"
@@ -82,12 +80,9 @@ jobs:
         timeout-minutes: ${{ inputs.timeout }}
         uses: ./.github/actions/docker-run
         with:
-          docker_username: ${{ github.actor }}
+          install_wheel: true
           docker_password: ${{ secrets.GITHUB_TOKEN }}
-          docker_image_arch: ${{ inputs.arch }}
           run_args: |
-            WHEEL_FILENAME=$(ls -1 *.whl)
-            pip3 install --user $WHEEL_FILENAME
             ${{ matrix.test-group.cmd }}
       - uses: ./.github/actions/slack-report
         if: ${{ failure() }}