Add use-only-dl-pytorch-org option to release matrix

.github/workflows/generate_binary_build_matrix.yml

@@ -35,6 +35,10 @@ on:
         description: "Build with CPU?"
         default: "enable"
         type: string
+      use-only-dl-pytorch-org:
+        description: "Use only download.pytorch.org when generating wheel install command?"
+        default: "false"
+        type: string
     outputs:
       matrix:
         description: "Generated build matrix"
@@ -67,6 +71,9 @@ jobs:
           # limit pull request builds to one version of python unless ciflow/binaries/all is applied to the workflow
           # should not affect builds that are from events that are not the pull_request event
           LIMIT_PR_BUILDS: ${{ github.event_name == 'pull_request' && !contains( github.event.pull_request.labels.*.name, 'ciflow/binaries/all') }}
+          # This is used when testing release binaries only from download.pytorch.org.
+          # In cases when pipy binaries are not published yet.
+          USE_ONLY_DL_PYTORCH_ORG: ${{ inputs.use-only-dl-pytorch-org }}
         run: |
           set -eou pipefail
           MATRIX_BLOB="$(python3 tools/scripts/generate_binary_build_matrix.py)"

tools/scripts/generate_binary_build_matrix.py

@@ -272,10 +272,11 @@ def get_wheel_install_command(
     gpu_arch_version: str,
     desired_cuda: str,
     python_version: str,
+    use_only_dl_pytorch_org: bool,
 ) -> str:
 
     index_url_option = "--index-url" if os != LINUX_AARCH64 else "--extra-index-url"
-    if channel == RELEASE and (
+    if  channel == RELEASE and (not use_only_dl_pytorch_org) and (
         (gpu_arch_version == "12.1" and os == LINUX)
         or (
             gpu_arch_type == CPU
@@ -300,6 +301,7 @@ def generate_conda_matrix(
     with_rocm: str,
     with_cpu: str,
     limit_pr_builds: bool,
+    use_only_dl_pytorch_org: bool,
 ) -> List[Dict[str, str]]:
     ret: List[Dict[str, str]] = []
     python_versions = list(mod.PYTHON_ARCHES)
@@ -353,7 +355,8 @@ def generate_libtorch_matrix(
     with_cuda: str,
     with_rocm: str,
     with_cpu: str,
-    limit_pr_builds: str,
+    limit_pr_builds: bool,
+    use_only_dl_pytorch_org: bool,
     abi_versions: Optional[List[str]] = None,
     arches: Optional[List[str]] = None,
     libtorch_variants: Optional[List[str]] = None,
@@ -446,6 +449,7 @@ def generate_wheels_matrix(
     with_rocm: str,
     with_cpu: str,
     limit_pr_builds: bool,
+    use_only_dl_pytorch_org: bool,
     arches: Optional[List[str]] = None,
     python_versions: Optional[List[str]] = None,
 ) -> List[Dict[str, str]]:
@@ -514,6 +518,7 @@ def generate_wheels_matrix(
                         gpu_arch_version,
                         desired_cuda,
                         python_version,
+                        use_only_dl_pytorch_org,
                     ),
                     "channel": channel,
                     "upload_to_base_bucket": upload_to_base_bucket,
@@ -538,6 +543,7 @@ def generate_build_matrix(
     with_rocm: str,
     with_cpu: str,
     limit_pr_builds: str,
+    use_only_dl_pytorch_org: str,
 ) -> Dict[str, List[Dict[str, str]]]:
     includes = []
 
@@ -558,6 +564,7 @@ def generate_build_matrix(
                     with_rocm,
                     with_cpu,
                     limit_pr_builds == "true",
+                    use_only_dl_pytorch_org == "true",
                 )
             )
 
@@ -616,6 +623,15 @@ def main(args) -> None:
         choices=["true", "false"],
         default=os.getenv("LIMIT_PR_BUILDS", "false"),
     )
+    # This is used when testing release builds to test release binaries
+    # only from download.pytorch.org. When pipy binaries are not released yet.
+    parser.add_argument(
+        "--use-only-dl-pytorch-org",
+        help="Use only download.pytorch.org when gen wheel install command?",
+        type=str,
+        choices=["true", "false"],
+        default=os.getenv("USE_ONLY_DL_PYTORCH_ORG", "false"),
+    )
 
     options = parser.parse_args(args)
 
@@ -631,6 +647,7 @@ def main(args) -> None:
         options.with_rocm,
         options.with_cpu,
         options.limit_pr_builds,
+        options.use_only_dl_pytorch_org,
     )
 
     print(json.dumps(build_matrix))