diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index d0aaea27..0f037457 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -75,10 +75,10 @@ jobs: export label_extension="-intel-cpu" export docker_devices="none" export docker_volume="/mnt/cache" - export runs_on="ubuntu-latest" - # export runs_on="aws-highmemory-32-plus-priv" + # export runs_on="ubuntu-latest" + export runs_on="aws-highmemory-32-plus-priv" export platform="cpu" - export extra_pytest="-k test_flash_llama_load" + export extra_pytest="-k test_flash_llama_simple" ;; esac echo $dockerfile diff --git a/integration-tests/models/test_flash_llama.py b/integration-tests/models/test_flash_llama.py index c69314ff..bf49dc0b 100644 --- a/integration-tests/models/test_flash_llama.py +++ b/integration-tests/models/test_flash_llama.py @@ -15,7 +15,7 @@ async def flash_llama(flash_llama_handle): @pytest.mark.asyncio @pytest.mark.private -async def test_flash_llama(flash_llama, response_snapshot): +async def test_flash_llama_simple(flash_llama, response_snapshot): response = await flash_llama.generate( "Test request", max_new_tokens=10, decoder_input_details=True )