From 0b2370819bf5a7a3bddf5f956b178b6e93b585f5 Mon Sep 17 00:00:00 2001
From: Alistair Muldal <alimuldal@google.com>
Date: Mon, 12 Apr 2021 13:05:11 +0100
Subject: [PATCH] Fix failing Travis build for `tvt/run.sh`

* Restrict the script to a single episode to avoid timeouts
* Ensure that `wheel` is up to date (we may need to build our own wheels for some of the requirements)

PiperOrigin-RevId: 367980522
---
 .travis.yml | 11 ++++++-----
 tvt/run.sh  |  7 ++++---
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 20bd9f3..7d1dad8 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -8,18 +8,19 @@ python:
   - "3.6"
 env:
   matrix:
-    - PROJECT="tvt"
+    - PROJECT="avae"
     - PROJECT="cs_gan"
     - PROJECT="iodine"
-    - PROJECT="transporter"
-    - PROJECT="polygen"
-    - PROJECT="learning_to_simulate"
-    - PROJECT="avae"
     - PROJECT="kfac_ferminet_alpha"
+    - PROJECT="learning_to_simulate"
+    - PROJECT="polygen"
+    - PROJECT="transporter"
+    - PROJECT="tvt"
 before_script:
   - sudo apt-get update -qq
   - pip install --upgrade pip
   - pip install --upgrade virtualenv
+  - pip install --upgrade wheel
 script:
   - ./$PROJECT/run.sh
 
diff --git a/tvt/run.sh b/tvt/run.sh
index bcee93e..722ae79 100755
--- a/tvt/run.sh
+++ b/tvt/run.sh
@@ -13,8 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-python3 -m venv tvt_venv
-source tvt_venv/bin/activate
+python3 -m venv /tmp/tvt_venv
+source /tmp/tvt_venv/bin/activate
+pip install -U pip
 pip install -r tvt/requirements.txt
 
-python3 -m tvt.main
+python3 -m tvt.main --num_episodes=1