Fix conda nightly to use the right pytorch nightly version corresponding to each platform (meta-pytorch#662)

ejguan · facebook-github-bot · commit 87d6dc3d6b0d · 2022-07-18T10:11:46.000-07:00
Summary: Fix one TODO in the releng improvement: meta-pytorch#253 Nightly release for TorchData is failing for the last two days because TorchData uses the latest version of PyTorch across platforms including macos, linux and windows. In some cases, like the current [workflow](https://github.com/pytorch/data/actions/runs/2690032872), PyTorch has not been updated nightly for macos with Python 3.7 and 3.9. In this case, TorchData conda release couldn't find the latest version of PyTorch for these two platforms then fail. I am adding this PR to dynamically detect the versions of PyTorch nightly based on OS (`conda search` would find the packages of the runner's OS by default), CPU/CUDA and Python version. Pull Request resolved: meta-pytorch#662 Reviewed By: NivekT Differential Revision: D37926216 Pulled By: ejguan fbshipit-source-id: 45531a57674815c8ada81e77bee46d8d80e883d8
diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash
@@ -154,7 +154,12 @@ setup_conda_pytorch_constraint() {
   CONDA_CHANNEL_FLAGS=${CONDA_CHANNEL_FLAGS:-}
   if [[ -z "$PYTORCH_VERSION" ]]; then
     export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c pytorch-nightly"
-    export PYTORCH_VERSION="$(conda search --json 'pytorch[channel=pytorch-nightly]' | python -c "import sys, json, re; print(re.sub(r'\\+.*$', '', json.load(sys.stdin)['pytorch'][-1]['version']))")"
+    export PYTORCH_VERSION="$(conda search --json 'pytorch[channel=pytorch-nightly]' | \
+        python -c "import json, os, re, sys; cuver = os.environ.get('CU_VERSION'); \
+            pyver = os.environ.get('PYTHON_VERSION'); \
+            print(re.sub(r'\\+.*$', '',
+                [x['version'] for x in json.load(sys.stdin)['pytorch'] \
+                    if (x['platform'] == 'darwin' or cuver in x['fn']) and 'py' + pyver in x['fn']][-1]))")"
   else
     export CONDA_CHANNEL_FLAGS="${CONDA_CHANNEL_FLAGS} -c pytorch -c pytorch-${UPLOAD_CHANNEL}"
   fi
