From 13bd1e971647d12eb805a4817e137c1e2fcfde68 Mon Sep 17 00:00:00 2001
From: Tristan Rice <rice@fn.lc>
Date: Thu, 21 Oct 2021 10:46:59 -0700
Subject: [PATCH] switch to torch 1.10 (#295)

Summary: Pull Request resolved: https://github.com/pytorch/torchx/pull/295

Reviewed By: kiukchung, aivanou

Differential Revision: D31832289

Pulled By: d4l3k

fbshipit-source-id: 9b18dceed75d6d80a20727f4f88597a856566a89
---
 .github/workflows/python-unittests.yaml | 17 +++--------------
 dev-requirements.txt                    |  5 +++--
 torchx/runtime/container/Dockerfile     | 17 +----------------
 3 files changed, 7 insertions(+), 32 deletions(-)

diff --git a/.github/workflows/python-unittests.yaml b/.github/workflows/python-unittests.yaml
index 294adfcac..148f1439b 100644
--- a/.github/workflows/python-unittests.yaml
+++ b/.github/workflows/python-unittests.yaml
@@ -10,20 +10,11 @@ jobs:
   unittest:
     strategy:
       matrix:
+        python-version: [3.7, 3.8, 3.9]
+        platform: [ubuntu-18.04]
         include:
-# TODO: uncomment this when torch.1-10 is released
-#          - python-version: 3.9
-#            platform: macos-latest
-# TODO: remove this when torch.1-10 is released
-          - python-version: 3.7
-            platform: ubuntu-18.04
-            torch-whl: https://download.pytorch.org/whl/test/cu113/torch-1.10.0%2Bcu113-cp37-cp37m-linux_x86_64.whl
-          - python-version: 3.8
-            platform: ubuntu-18.04
-            torch-whl: https://download.pytorch.org/whl/test/cu113/torch-1.10.0%2Bcu113-cp38-cp38-linux_x86_64.whl
           - python-version: 3.9
-            platform: ubuntu-18.04
-            torch-whl: https://download.pytorch.org/whl/test/cu113/torch-1.10.0%2Bcu113-cp39-cp39-linux_x86_64.whl
+            platform: macos-latest
       fail-fast: false
     runs-on: ${{ matrix.platform }}
     steps:
@@ -39,8 +30,6 @@ jobs:
           set -eux
           pip install coverage codecov
           pip install -e .[dev]
-          pip uninstall -y torch
-          pip install ${{ matrix.torch-whl }}
       - name: Run tests
         run: coverage run -m unittest discover --verbose --start-directory . --pattern "*_test.py"
       - name: Coverage
diff --git a/dev-requirements.txt b/dev-requirements.txt
index 393623913..7bbb2b170 100644
--- a/dev-requirements.txt
+++ b/dev-requirements.txt
@@ -5,8 +5,9 @@ pyre-extensions>=0.0.21
 black>=21.5b1
 usort==0.6.4
 pytorch-lightning>=1.4.9
-torch>=1.9.0
-torchvision>=0.10.0
+torch>=1.10.0
+torchvision>=0.11.1
+torchtext>=0.11.0
 classy-vision>=0.6.0
 flake8==3.9.0
 ts>=0.5.1
diff --git a/torchx/runtime/container/Dockerfile b/torchx/runtime/container/Dockerfile
index 1684db2ae..51cd98bc8 100644
--- a/torchx/runtime/container/Dockerfile
+++ b/torchx/runtime/container/Dockerfile
@@ -2,27 +2,12 @@ FROM pytorch/pytorch:1.9.0-cuda10.2-cudnn7-runtime
 
 RUN pip install fsspec[s3]
 
-# Note: the pytorch image installs version of the torchvision that is not compatible
-# with CUDA. By using it directly there will be a segfault.
-RUN pip uninstall -y torchvision
-RUN pip install torchvision
-
 WORKDIR /app
 
 # copy requirements early so we don't have to redownload dependencies on code
 # changes
 COPY dev-requirements.txt /app
-RUN pip install -r dev-requirements.txt
-
-# TODO: Remove this and get dependency on torch 1.10 when it is released
-
-RUN pip uninstall -y torch
-
-RUN pip install https://download.pytorch.org/whl/test/cu113/torch-1.10.0%2Bcu113-cp37-cp37m-linux_x86_64.whl
-
-RUN pip uninstall -y torchtext
-
-RUN pip install https://download.pytorch.org/whl/test/torchtext-0.10.0-cp37-cp37m-linux_x86_64.whl
+RUN pip install --upgrade -r dev-requirements.txt
 
 COPY . /app