tools: Add ability to build a python only package

Adds the ability to generate the binary build matrix with a singular
version of python to account for packages that only need to build a
package that looks similar to:

`torchtune-0.0.1-py3-none-any.whl`

Signed-off-by: Eli Uriegas <[email protected]>

.github/workflows/generate_binary_build_matrix.yml

@@ -39,6 +39,10 @@ on:
         description: "Use only download.pytorch.org when generating wheel install command?"
         default: "false"
         type: string
+      build-python-only:
+        description: "Generate binary build matrix for a python only package (i.e. only one python version)"
+        default: "disable"
+        type: string
     outputs:
       matrix:
         description: "Generated build matrix"
@@ -74,6 +78,7 @@ jobs:
           # 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 }}
+          BUILD_PYTHON_ONLY: ${{ inputs.build-python-only }}
         run: |
           set -eou pipefail
           MATRIX_BLOB="$(python3 tools/scripts/generate_binary_build_matrix.py)"

.github/workflows/test_build_wheels_linux_without_cuda.yml

@@ -33,6 +33,7 @@ jobs:
         include:
           - repository: pytorch/torchtune
             package-name: torchtune
+            build-python-only: enable
     uses: ./.github/workflows/build_wheels_linux.yml
     name: ${{ matrix.repository }}
     with:
@@ -44,3 +45,4 @@ jobs:
       package-name: ${{ matrix.package-name }}
       trigger-event: "${{ github.event_name }}"
       build-platform: 'python-build-package'
+      build-python-only: ${{ matrix.build-python-only }}

tools/scripts/generate_binary_build_matrix.py

@@ -124,15 +124,19 @@ def validation_runner(arch_type: str, os: str) -> str:
         return LINUX_CPU_RUNNER
 
 
-def initialize_globals(channel: str):
+def initialize_globals(channel: str, build_python_only: bool):
     if channel == TEST:
         mod.CURRENT_VERSION = CURRENT_CANDIDATE_VERSION
     else:
         mod.CURRENT_VERSION = CURRENT_STABLE_VERSION
 
     mod.CUDA_ARCHES = CUDA_ARCHES_DICT[channel]
     mod.ROCM_ARCHES = ROCM_ARCHES_DICT[channel]
-    mod.PYTHON_ARCHES = PYTHON_ARCHES_DICT[channel]
+    if build_python_only:
+        # Only select the oldest version of python if building a python only package
+        mod.PYTHON_ARCHES = [PYTHON_ARCHES_DICT[channel][0]]
+    else:
+        mod.PYTHON_ARCHES = PYTHON_ARCHES_DICT[channel]
     mod.WHEEL_CONTAINER_IMAGES = {
         **{
             gpu_arch: f"pytorch/manylinux-builder:cuda{gpu_arch}"
@@ -542,6 +546,7 @@ def generate_build_matrix(
     with_cpu: str,
     limit_pr_builds: str,
     use_only_dl_pytorch_org: str,
+    build_python_only: str,
 ) -> Dict[str, List[Dict[str, str]]]:
     includes = []
 
@@ -553,7 +558,7 @@ def generate_build_matrix(
 
     for channel in channels:
         for package in package_types:
-            initialize_globals(channel)
+            initialize_globals(channel, build_python_only == ENABLE)
             includes.extend(
                 GENERATING_FUNCTIONS_BY_PACKAGE_TYPE[package](
                     operating_system,
@@ -630,6 +635,17 @@ def main(args) -> None:
         choices=["true", "false"],
         default=os.getenv("USE_ONLY_DL_PYTORCH_ORG", "false"),
     )
+    # Generates a single version python for building python packages only
+    # This basically makes it so that we only generate a matrix including the oldest
+    # version of python that we support
+    # For packages that look similar to torchtune-0.0.1-py3-none-any.whl
+    parser.add_argument(
+        "--build-python-only",
+        help="Build python only",
+        type=str,
+        choices=[ENABLE, DISABLE],
+        default=os.getenv("BUILD_PYTHON_ONLY", ENABLE),
+    )
 
     options = parser.parse_args(args)
 
@@ -646,6 +662,7 @@ def main(args) -> None:
         options.with_cpu,
         options.limit_pr_builds,
         options.use_only_dl_pytorch_org,
+        options.build_python_only,
     )
 
     print(json.dumps(build_matrix))