Add 3.11 and --pytorch-only options

Author: malfet

build_aarch64_wheel.py

@@ -475,6 +475,7 @@ def start_build(host: RemoteHost, *,
                 compiler="gcc-8",
                 use_conda=True,
                 python_version="3.8",
+                pytorch_only:bool=False,
                 shallow_clone=True,
                 enable_mkldnn=False) -> Tuple[str, str]:
     git_clone_flags = " --depth 1 --shallow-submodules" if shallow_clone else ""
@@ -543,6 +544,8 @@ def start_build(host: RemoteHost, *,
     print('Installing PyTorch wheel')
     host.run_cmd(f"pip3 install pytorch/dist/{pytorch_wheel_name}")
 
+    if pytorch_only:
+        return pytorch_wheel_name, None
     vision_wheel_name = build_torchvision(host, branch=branch, use_conda=use_conda, git_clone_flags=git_clone_flags)
     build_torchaudio(host, branch=branch, use_conda=use_conda, git_clone_flags=git_clone_flags)
     build_torchtext(host, branch=branch, use_conda=use_conda, git_clone_flags=git_clone_flags)
@@ -674,9 +677,10 @@ def parse_arguments():
     parser.add_argument("--build-only", action="store_true")
     parser.add_argument("--test-only", type=str)
     parser.add_argument("--os", type=str, choices=list(os_amis.keys()), default='ubuntu18_04')
-    parser.add_argument("--python-version", type=str, choices=['3.6', '3.7', '3.8', '3.9', '3.10'], default=None)
+    parser.add_argument("--python-version", type=str, choices=['3.6', '3.7', '3.8', '3.9', '3.10', '3.11'], default=None)
     parser.add_argument("--alloc-instance", action="store_true")
     parser.add_argument("--list-instances", action="store_true")
+    parser.add_argument("--pytorch-only", action="store_true")
     parser.add_argument("--keep-running", action="store_true")
     parser.add_argument("--terminate-instances", action="store_true")
     parser.add_argument("--instance-type", type=str, default="t4g.2xlarge")
@@ -754,6 +758,7 @@ def parse_arguments():
                     branch=args.branch,
                     compiler=args.compiler,
                     python_version=python_version,
+                    pytorch_only=args.pytorch_only,
                     enable_mkldnn=args.enable_mkldnn)
     if not args.keep_running:
         print(f'Waiting for instance {inst.id} to terminate')