From ae1fbbe52aa8177ae3799a49bb8066729445a6fd Mon Sep 17 00:00:00 2001 From: Roman Kazantsev Date: Mon, 23 Dec 2024 21:21:17 +0400 Subject: [PATCH] [GHA][JAX FE] Separate JAX layer tests to special job and have dedicated req file (#28178) **Details:** Separate JAX layer tests to special job and have dedicated req file **Ticket:** TBD --------- Signed-off-by: Kazantsev, Roman --- .github/workflows/job_jax_layer_tests.yml | 133 ++++++++++++++++++ .github/workflows/job_jax_models_tests.yml | 2 +- .github/workflows/job_python_unit_tests.yml | 8 -- .github/workflows/linux_arm64.yml | 10 ++ .github/workflows/mac.yml | 9 ++ .github/workflows/mac_arm64.yml | 9 ++ .github/workflows/ubuntu_22.yml | 10 ++ .github/workflows/ubuntu_24.yml | 10 ++ .github/workflows/windows_vs2019_release.yml | 9 ++ tests/CMakeLists.txt | 2 +- .../test_tf_UnaryOpsAllRealDomain.py | 2 +- tests/model_hub_tests/jax/requirements.txt | 10 -- tests/requirements_jax | 13 ++ tests/requirements_tensorflow | 2 +- 14 files changed, 207 insertions(+), 22 deletions(-) create mode 100644 .github/workflows/job_jax_layer_tests.yml delete mode 100644 tests/model_hub_tests/jax/requirements.txt create mode 100644 tests/requirements_jax diff --git a/.github/workflows/job_jax_layer_tests.yml b/.github/workflows/job_jax_layer_tests.yml new file mode 100644 index 00000000000000..25f171060f43be --- /dev/null +++ b/.github/workflows/job_jax_layer_tests.yml @@ -0,0 +1,133 @@ +name: JAX Layer Tests + +on: + workflow_call: + inputs: + runner: + description: 'Machine on which the tests would run' + type: string + required: true + container: + description: 'JSON to be converted to the value of the "container" configuration for the job' + type: string + required: false + default: '{"image": null}' + affected-components: + description: 'Components that are affected by changes in the commit defined by the Smart CI Action' + type: string + required: true + python-version: + description: 'Python version to setup. E.g., "3.11"' + type: string + required: true + +permissions: read-all + +env: + PIP_CACHE_PATH_LINUX: /mount/caches/pip/linux + PIP_CACHE_PATH_WIN: "C:\\mount\\caches\\pip\\win" + +jobs: + JAX_Layer_Tests: + name: JAX Layer Tests + timeout-minutes: 40 + runs-on: ${{ inputs.runner }} + container: ${{ fromJSON(inputs.container) }} + defaults: + run: + shell: ${{ contains(inputs.runner, 'win') && 'pwsh' || 'bash' }} + env: + DEBIAN_FRONTEND: noninteractive # to prevent apt-get from waiting user input + OPENVINO_REPO: ${{ github.workspace }}/openvino + INSTALL_DIR: ${{ github.workspace }}/install + INSTALL_TEST_DIR: ${{ github.workspace }}/install/tests + INSTALL_WHEELS_DIR: ${{ github.workspace }}/install/wheels + LAYER_TESTS_INSTALL_DIR: ${{ github.workspace }}/install/tests/layer_tests + steps: + - name: Download OpenVINO artifacts (tarballs) + uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8 + with: + pattern: openvino_[tests]* + path: ${{ env.INSTALL_DIR }} + merge-multiple: true + + - name: Download OpenVINO artifacts (wheels) + uses: actions/download-artifact@fa0a91b85d4f404e444e00e005971372dc801d16 # v4.1.8 + with: + pattern: openvino_[wheels]* + path: ${{ env.INSTALL_WHEELS_DIR }} + merge-multiple: true + + # Needed as ${{ github.workspace }} is not working correctly when using Docker + - name: Setup Variables + if: runner.os != 'Windows' + run: | + echo "OPENVINO_REPO=$GITHUB_WORKSPACE/openvino" >> "$GITHUB_ENV" + echo "INSTALL_DIR=$GITHUB_WORKSPACE/install" >> "$GITHUB_ENV" + echo "INSTALL_TEST_DIR=$GITHUB_WORKSPACE/install/tests" >> "$GITHUB_ENV" + echo "INSTALL_WHEELS_DIR=$GITHUB_WORKSPACE/install/wheels" >> "$GITHUB_ENV" + echo "LAYER_TESTS_INSTALL_DIR=$GITHUB_WORKSPACE/install/tests/layer_tests" >> "$GITHUB_ENV" + + - name: Install OpenVINO dependencies (mac) + if: runner.os == 'macOS' + run: brew install pigz + + - name: Extract OpenVINO packages (Linux, macOS) + if: runner.os != 'Windows' + run: | + pigz -dc openvino_tests.tar.gz | tar -xf - -C ${INSTALL_DIR} + working-directory: ${{ env.INSTALL_DIR }} + + - name: Extract OpenVINO artifacts (Windows) + if: runner.os == 'Windows' + run: | + Expand-Archive openvino_tests.zip -DestinationPath ${{ env.INSTALL_DIR }} + working-directory: ${{ env.INSTALL_DIR }} + + - name: Fetch setup_python and install wheels actions + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + timeout-minutes: 15 + with: + sparse-checkout: | + .github/actions/setup_python/action.yml + .github/actions/install_ov_wheels/action.yml + sparse-checkout-cone-mode: false + path: 'openvino' + + - name: Setup Python ${{ inputs.python-version }} + uses: ./openvino/.github/actions/setup_python + with: + version: ${{ inputs.python-version }} + pip-cache-path: ${{ runner.os == 'Linux' && env.PIP_CACHE_PATH_LINUX || env.PIP_CACHE_PATH_WIN }} + should-setup-pip-paths: ${{ runner.os != 'macOS' }} + self-hosted-runner: ${{ runner.os != 'macOS' }} + + - name: Install OpenVINO Python wheels + uses: ./openvino/.github/actions/install_ov_wheels + with: + wheels-dir-path: ${{ env.INSTALL_WHEELS_DIR }} + wheels-to-install: 'openvino' + + - name: Install JAX Layer tests dependencies + run: | + # jax test requirements + python3 -m pip install -r ${{ env.INSTALL_TEST_DIR }}/requirements_jax + + - name: JAX Layer Tests + if: ${{ fromJSON(inputs.affected-components).JAX_FE.test && runner.arch != 'ARM64' }} # Ticket: 126287, 142196 + run: python3 -m pytest ${{ env.LAYER_TESTS_INSTALL_DIR }}/jax_tests ${PARALLEL} -m precommit_jax_fe --junitxml=${{ env.INSTALL_TEST_DIR }}/TEST-jax.xml + env: + TEST_DEVICE: CPU + TEST_PRECISION: FP16 + JAX_TRACE_MODE: JAXPR + PARALLEL: ${{ runner.os == 'Windows' && ' ' || '-n logical'}} + + - name: Upload Test Results + uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3 + if: ${{ !cancelled() }} + with: + name: test-results-python-jax-layers + path: | + ${{ env.INSTALL_TEST_DIR }}/TEST*.html + ${{ env.INSTALL_TEST_DIR }}/TEST*.xml + if-no-files-found: 'warn' diff --git a/.github/workflows/job_jax_models_tests.yml b/.github/workflows/job_jax_models_tests.yml index 07155db1016057..57eb07a83aa423 100644 --- a/.github/workflows/job_jax_models_tests.yml +++ b/.github/workflows/job_jax_models_tests.yml @@ -89,7 +89,7 @@ jobs: - name: Install JAX tests requirements for precommit run: | - python3 -m pip install -r ${MODEL_HUB_TESTS_INSTALL_DIR}/jax/requirements.txt + python3 -m pip install -r ${{ env.INSTALL_TEST_DIR }}/requirements_jax - name: JAX/Flax Models Tests from Hugging Face if: ${{ inputs.model_scope == 'precommit' || inputs.model_scope == 'nightly' }} diff --git a/.github/workflows/job_python_unit_tests.yml b/.github/workflows/job_python_unit_tests.yml index b04f719c8e296f..e1532d530ff2db 100644 --- a/.github/workflows/job_python_unit_tests.yml +++ b/.github/workflows/job_python_unit_tests.yml @@ -162,14 +162,6 @@ jobs: export LD_LIBRARY_PATH=${PIP_INSTALL_PATH}/openvino/libs:$LD_LIBRARY_PATH python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/py_frontend_tests --junitxml=${INSTALL_TEST_DIR}/TEST-test_py_fontend.xml - - name: JAX Layer Tests - JAX FE - if: ${{ fromJSON(inputs.affected-components).JAX_FE.test && runner.arch != 'ARM64' && runner.os != 'macOS' }} - run: python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/jax_tests/ -m precommit_jax_fe --junitxml=${INSTALL_TEST_DIR}/TEST-jax_fe.xml - env: - TEST_DEVICE: CPU - TEST_PRECISION: FP16 - JAX_TRACE_MODE: JAXPR - - name: TensorFlow Lite Layer Tests - TFL FE if: fromJSON(inputs.affected-components).TFL_FE.test run: python3 -m pytest ${LAYER_TESTS_INSTALL_DIR}/tensorflow_lite_tests/ -n logical --junitxml=${INSTALL_TEST_DIR}/TEST-tfl_fe.xml diff --git a/.github/workflows/linux_arm64.yml b/.github/workflows/linux_arm64.yml index 66e825e5d5e126..ca1ca6e056e23d 100644 --- a/.github/workflows/linux_arm64.yml +++ b/.github/workflows/linux_arm64.yml @@ -202,6 +202,16 @@ jobs: affected-components: ${{ needs.smart_ci.outputs.affected_components }} python-version: '3.11' + JAX_Layer_Tests: + name: JAX Layer Tests + needs: [ Build, Docker, Smart_CI ] + uses: ./.github/workflows/job_jax_layer_tests.yml + with: + runner: 'aks-linux-16-cores-32gb-arm' + container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_test.ubuntu_20_04_arm64 }}", "volumes": ["/mount:/mount"]}' + affected-components: ${{ needs.smart_ci.outputs.affected_components }} + python-version: '3.11' + CPU_Functional_Tests: name: CPU functional tests if: fromJSON(needs.smart_ci.outputs.affected_components).CPU.test diff --git a/.github/workflows/mac.yml b/.github/workflows/mac.yml index 5e4335b8151c02..0fbc20cf19594b 100644 --- a/.github/workflows/mac.yml +++ b/.github/workflows/mac.yml @@ -356,6 +356,15 @@ jobs: affected-components: ${{ needs.smart_ci.outputs.affected_components }} python-version: '3.11' + JAX_Layer_Tests: + name: JAX Layer Tests + needs: [ Build, Smart_CI ] + uses: ./.github/workflows/job_jax_layer_tests.yml + with: + runner: 'macos-13' + affected-components: ${{ needs.smart_ci.outputs.affected_components }} + python-version: '3.11' + CPU_Functional_Tests: name: CPU functional tests # if: fromJSON(needs.smart_ci.outputs.affected_components).CPU.test diff --git a/.github/workflows/mac_arm64.yml b/.github/workflows/mac_arm64.yml index 855d76973cc2e4..b60daefa442c83 100644 --- a/.github/workflows/mac_arm64.yml +++ b/.github/workflows/mac_arm64.yml @@ -355,6 +355,15 @@ jobs: affected-components: ${{ needs.smart_ci.outputs.affected_components }} python-version: '3.11' + JAX_Layer_Tests: + name: JAX Layer Tests + needs: [ Build, Smart_CI ] + uses: ./.github/workflows/job_jax_layer_tests.yml + with: + runner: 'macos-13-xlarge' + affected-components: ${{ needs.smart_ci.outputs.affected_components }} + python-version: '3.11' + CPU_Functional_Tests: name: CPU functional tests if: fromJSON(needs.smart_ci.outputs.affected_components).CPU.test diff --git a/.github/workflows/ubuntu_22.yml b/.github/workflows/ubuntu_22.yml index 5aed74bbb242b8..e5c7d25003de1e 100644 --- a/.github/workflows/ubuntu_22.yml +++ b/.github/workflows/ubuntu_22.yml @@ -334,6 +334,16 @@ jobs: affected-components: ${{ needs.smart_ci.outputs.affected_components }} python-version: '3.11' + JAX_Layer_Tests: + name: JAX Layer Tests + needs: [ Docker, Build, Smart_CI ] + uses: ./.github/workflows/job_jax_layer_tests.yml + with: + runner: 'aks-linux-4-cores-16gb' + container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_test.ubuntu_22_04_x64 }}", "volumes": ["/mount:/mount"]}' + affected-components: ${{ needs.smart_ci.outputs.affected_components }} + python-version: '3.11' + CPU_Functional_Tests: name: CPU functional tests if: fromJSON(needs.smart_ci.outputs.affected_components).CPU.test diff --git a/.github/workflows/ubuntu_24.yml b/.github/workflows/ubuntu_24.yml index 25be095e692d35..beac15bfbda97d 100644 --- a/.github/workflows/ubuntu_24.yml +++ b/.github/workflows/ubuntu_24.yml @@ -156,6 +156,16 @@ jobs: affected-components: ${{ needs.smart_ci.outputs.affected_components }} python-version: '3.12' + JAX_Layer_Tests: + name: JAX Layer Tests + needs: [ Docker, Build, Smart_CI ] + uses: ./.github/workflows/job_jax_layer_tests.yml + with: + runner: 'aks-linux-4-cores-16gb' + container: '{"image": "${{ fromJSON(needs.docker.outputs.images).ov_test.ubuntu_24_04_x64 }}", "volumes": ["/mount:/mount"]}' + affected-components: ${{ needs.smart_ci.outputs.affected_components }} + python-version: '3.12' + TensorFlow_Layer_Tests: name: TensorFlow Layer Tests needs: [ Docker, Build, Smart_CI, Openvino_tokenizers ] diff --git a/.github/workflows/windows_vs2019_release.yml b/.github/workflows/windows_vs2019_release.yml index f1fd0be596baa2..de33f2603d7430 100644 --- a/.github/workflows/windows_vs2019_release.yml +++ b/.github/workflows/windows_vs2019_release.yml @@ -499,6 +499,15 @@ jobs: affected-components: ${{ needs.smart_ci.outputs.affected_components }} python-version: '3.11' + JAX_Layer_Tests: + name: JAX Layer Tests + needs: [ Build, Smart_CI ] + uses: ./.github/workflows/job_jax_layer_tests.yml + with: + runner: 'aks-win-8-cores-16gb' + affected-components: ${{ needs.smart_ci.outputs.affected_components }} + python-version: '3.11' + CXX_Unit_Tests: name: C++ unit tests needs: [ Build, Smart_CI ] diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 08b4308479ef03..de3ad80280d603 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -7,5 +7,5 @@ add_subdirectory(model_hub_tests) add_subdirectory(samples_tests) add_subdirectory(e2e_tests) -install(FILES requirements_pytorch requirements_tensorflow requirements_onnx +install(FILES requirements_pytorch requirements_tensorflow requirements_onnx requirements_jax DESTINATION tests COMPONENT tests EXCLUDE_FROM_ALL) diff --git a/tests/layer_tests/tensorflow_tests/test_tf_UnaryOpsAllRealDomain.py b/tests/layer_tests/tensorflow_tests/test_tf_UnaryOpsAllRealDomain.py index 4ff4d589cbae32..5c1037e38cfc84 100644 --- a/tests/layer_tests/tensorflow_tests/test_tf_UnaryOpsAllRealDomain.py +++ b/tests/layer_tests/tensorflow_tests/test_tf_UnaryOpsAllRealDomain.py @@ -67,4 +67,4 @@ def test_unary_ops(self, input_shape, input_type, op_type, pytest.skip("159585: accuracy error on ARM") self._test(*self.create_unary_net(input_shape, input_type, op_type), ie_device, precision, ir_version, temp_dir=temp_dir, - use_legacy_frontend=use_legacy_frontend, custom_eps=1e-3) + use_legacy_frontend=use_legacy_frontend, custom_eps=3 * 1e-3) diff --git a/tests/model_hub_tests/jax/requirements.txt b/tests/model_hub_tests/jax/requirements.txt deleted file mode 100644 index 328084ac050ca6..00000000000000 --- a/tests/model_hub_tests/jax/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ --c ../../constraints.txt -numpy -pytest -pytest-html -transformers -requests -jax -jaxlib -flax -pillow \ No newline at end of file diff --git a/tests/requirements_jax b/tests/requirements_jax new file mode 100644 index 00000000000000..c392df4359bee3 --- /dev/null +++ b/tests/requirements_jax @@ -0,0 +1,13 @@ +numpy==1.26.4; python_version < "3.12" or platform_system == "Darwin" and platform_machine == "x86_64" +numpy==2.2.1; python_version >= "3.12" and (platform_system != "Darwin" or platform_machine != "x86_64") +pytest==7.0.1 +pytest-xdist[psutil]==3.6.1 +pytest-html==4.1.1 +jax==0.4.38; (platform_system != "Darwin" or platform_machine != "x86_64") and python_version > "3.9" +# tensorflow 2.16.2 depends on ml-dtypes~=0.3.1 and jax 0.4.35 depends on ml-dtypes>=0.4.0 +jax==0.4.33; (platform_system == "Darwin" and platform_machine == "x86_64") and python_version > "3.9" +jax==0.4.30; python_version <= "3.9" +flax==0.10.2 +transformers==4.47.1 +defusedxml +pillow diff --git a/tests/requirements_tensorflow b/tests/requirements_tensorflow index 5369b0135f7618..8e0d1141695ef9 100644 --- a/tests/requirements_tensorflow +++ b/tests/requirements_tensorflow @@ -17,7 +17,7 @@ wrapt==1.15.0; python_version >= "3.12" # tensorflow-text is not available for both Windows and ARM platforms tensorflow-text==2.18.0; python_version < "3.12" and platform_system == "Linux" and platform_machine == "x86_64" tensorflow-hub==0.16.1 -jax==0.4.35; (platform_system != "Darwin" or platform_machine != "x86_64") and python_version > "3.9" +jax==0.4.38; (platform_system != "Darwin" or platform_machine != "x86_64") and python_version > "3.9" # tensorflow 2.16.2 depends on ml-dtypes~=0.3.1 and jax 0.4.35 depends on ml-dtypes>=0.4.0 jax==0.4.33; (platform_system == "Darwin" and platform_machine == "x86_64") and python_version > "3.9" jax==0.4.30; python_version <= "3.9"