Add conda constraint

pytorch · Nov 16, 2022 · 77a3d8b · 77a3d8b
1 parent 0f40c96
commit 77a3d8b
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 17 deletions.
diff --git a/packaging/install_torchdata.sh b/packaging/install_torchdata.sh
@@ -21,32 +21,29 @@ if [ "$package_type" = "wheel" ]; then
   install_cmd="pip install"
   if [ "$channel" = "nightly" ]; then
     install_cmd="${install_cmd} --pre"
-    install_channel="--extra-index-url https://download.pytorch.org/whl/nightly/cpu"
-  else
-    install_channel="--extra-index-url https://download.pytorch.org/whl/test/cpu"
-  fi
-else
-  install_cmd="conda install"
-  if [ "$channel" = "nightly" ]; then
-    install_channel="-c pytorch-nightly"
-  else
-    install_channel="-c pytorch-test"
   fi
+  install_channel="--extra-index-url https://download.pytorch.org/whl/${channel}/cpu"
+  # Only install torchdata for wheels
+  $install_cmd torchdata $install_channel
 fi
 
-$install_cmd torchdata $install_channel
-
 if [ "$package_type" = "wheel" ]; then
   TORCHDATA_VERSION="$(pip show torchdata | grep ^Version: | sed 's/Version:  *//' | sed 's/+.\+//')"
 else
-  TORCHDATA_VERSION="$(conda search --json torchdata[channel=pytorch-${channel}] | \
+  TORCHDATA_VERSION="$(conda search --json 'torchdata[channel=pytorch-'"${channel}"] | \
 	  python -c "import json, os, re, sys; \
 	    cuver = 'cpu'; \
-	    pyver = os.environ.get('PYTHON_VERSION'); \
+	    pyver = os.environ.get('PYTHON_VERSION').replace('.', ''); \
 	    print(re.sub(r'\\+.*$', '',
 	      [x['version'] for x in json.load(sys.stdin)['torchdata'] \
 		      if 'py' + pyver in x['fn']][-1]))"
 	    )"
+  echo "export CONDA_TORCHDATA_CONSTRAINT='- torchdata==${TORCHDATA_VERSION}'"
+  echo "export CONDA_TORCHDATA_CONSTRAINT='- torchdata==${TORCHDATA_VERSION}'" >> "${BUILD_ENV_FILE}"
 fi
 
-echo "export TORCHDATA_VERSION=${TORCHDATA_VERSION}" >> ${BUILD_ENV_FILE}
+echo "export TORCHDATA_VERSION=${TORCHDATA_VERSION}"
+echo "export TORCHDATA_VERSION=${TORCHDATA_VERSION}" >> "${BUILD_ENV_FILE}"
+
+
+tail -n 5 "${BUILD_ENV_FILE}"
diff --git a/packaging/pkg_helpers.bash b/packaging/pkg_helpers.bash
@@ -198,7 +198,7 @@ setup_pip_pytorch_version() {
     pip_install "torchdata==$TORCHDATA_VERSION" \
       -f https://download.pytorch.org/whl/torch_stable.html \
       -f "https://download.pytorch.org/whl/${UPLOAD_CHANNEL}/torch_${UPLOAD_CHANNEL}.html"
-  fi 
+  fi
 }
 
 # Fill PYTORCH_VERSION with the latest conda nightly version, and
@@ -233,6 +233,10 @@ setup_conda_pytorch_constraint() {
       export CONDA_EXTRA_BUILD_CONSTRAINT="- mkl<=2021.2.0"
     fi
   fi
+  if [[ -z "$TORCHDATA_VERSION" ]]; then
+    export TORCHDATA_VERSION="$(conda search --json 'torchdata[channel=pytorch-nightly]' | ${PYTHON} -c "import sys, json, re; print(re.sub(r'\\+.*$', '', json.load(sys.stdin)['torchdata'][-1]['version']))")"
+  fi
+  export CONDA_TORCHDATA_CONSTRAINT="- torchdata==$TORCHDATA_VERSION"
 }
 
 # Translate CUDA_VERSION into CUDA_CUDATOOLKIT_CONSTRAINT

diff --git a/packaging/torchtext/meta.yaml b/packaging/torchtext/meta.yaml
@@ -23,8 +23,8 @@ requirements:
     - python
     - requests
     - tqdm
-    - torchdata
     {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
+    {{ environ.get('CONDA_TORCHDATA_CONSTRAINT') }}
 
 build:
   string: py{{py}}