Reducing wait interval for BQ integration tests (#1827)

adchia · web-flow · commit 09e0955f709f · 2021-09-03T09:23:57.000-07:00
* For tests, reducing BQ's wait time. Some quick logging indicated that the uploading of the entity_df in integration tests was very slow

Signed-off-by: Danny Chiao &lt;danny@tecton.ai&gt;

* Moving environment variables into yml and MakeFile

Signed-off-by: Danny Chiao &lt;danny@tecton.ai&gt;

* Reducing prod retry cadence

Signed-off-by: Danny Chiao &lt;danny@tecton.ai&gt;
diff --git a/.github/workflows/integration_tests.yml b/.github/workflows/integration_tests.yml
@@ -52,7 +52,7 @@ jobs:
       - name: Install dependencies
         run: make install-python-ci-dependencies
       - name: Test python
-        run: FEAST_USAGE=False pytest -n 8 --cov=./ --cov-report=xml --verbose --color=yes sdk/python/tests --integration
+        run: FEAST_USAGE=False IS_TEST=True pytest -n 8 --cov=./ --cov-report=xml --verbose --color=yes sdk/python/tests --integration
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v1
         with:
diff --git a/.github/workflows/pr_integration_tests.yml b/.github/workflows/pr_integration_tests.yml
@@ -63,7 +63,7 @@ jobs:
       - name: Install dependencies
         run: make install-python-ci-dependencies
       - name: Test python
-        run: FEAST_USAGE=False pytest -n 8 --cov=./ --cov-report=xml --verbose --color=yes sdk/python/tests --integration
+        run: FEAST_USAGE=False IS_TEST=True pytest -n 8 --cov=./ --cov-report=xml --verbose --color=yes sdk/python/tests --integration
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v1
         with:
diff --git a/Makefile b/Makefile
@@ -56,7 +56,7 @@ test-python:
 	FEAST_USAGE=False pytest -n 8 sdk/python/tests
 
 test-python-integration:
-	FEAST_USAGE=False pytest -n 8 --integration sdk/python/tests
+	FEAST_USAGE=False IS_TEST=True pytest -n 8 --integration sdk/python/tests
 
 format-python:
 	# Sort
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -1,3 +1,4 @@
+import os
 import uuid
 from datetime import date, datetime, timedelta
 from typing import Dict, List, Optional, Union
@@ -239,7 +240,7 @@ def block_until_done(
     client: Client,
     bq_job: Union[bigquery.job.query.QueryJob, bigquery.job.load.LoadJob],
     timeout: int = 1800,
-    retry_cadence: int = 10,
+    retry_cadence: float = 1,
 ):
     """
     Waits for bq_job to finish running, up to a maximum amount of time specified by the timeout parameter (defaulting to 30 minutes).
@@ -255,6 +256,11 @@ def block_until_done(
         BigQueryJobCancelled exception to signify when that the job has been cancelled (i.e. from timeout or KeyboardInterrupt).
     """
 
+    # For test environments, retry more aggressively
+    is_test = os.getenv("IS_TEST", default="False") == "True"
+    if is_test:
+        retry_cadence = 0.1
+
     def _wait_until_done(job_id):
         if client.get_job(job_id).state in ["PENDING", "RUNNING"]:
             raise BigQueryJobStillRunning(job_id=job_id)
