Fix Docker image generation (#88741)

Pass install channel when building nightly images
Pass `TRITON_VERSION` argument to install triton for nightly images

Fix `generate_pytorch_version.py` to work with unannotated tags and avoid failures like the following:
```
% git checkout nightly
% ./.github/scripts/generate_pytorch_version.py

fatal: No annotated tags can describe '93f15b1b54ca5fb4a7ca9c21a813b4b86ebaeafa'.
However, there were unannotated tags: try --tags.
Traceback (most recent call last):
  File "/Users/nshulga/git/pytorch/pytorch-release/./.github/scripts/generate_pytorch_version.py", line 120, in <module>
    main()
  File "/Users/nshulga/git/pytorch/pytorch-release/./.github/scripts/generate_pytorch_version.py", line 115, in main
    print(version_obj.get_release_version())
  File "/Users/nshulga/git/pytorch/pytorch-release/./.github/scripts/generate_pytorch_version.py", line 75, in get_release_version
    if not get_tag():
  File "/Users/nshulga/git/pytorch/pytorch-release/./.github/scripts/generate_pytorch_version.py", line 37, in get_tag
    dirty_tag = subprocess.check_output(
  File "/Users/nshulga/miniforge3/lib/python3.9/subprocess.py", line 424, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/Users/nshulga/miniforge3/lib/python3.9/subprocess.py", line 528, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['git', 'describe']' returned non-zero exit status 128.
```
After the change nightly is reported as(due to autolabelling issue,
should be fixed by ttps://github.com/pytorch/test-infra/pull/1047 ):
```
 % ./.github/scripts/generate_pytorch_version.py
ciflow/inductor/26921+cpu
```

Even for tagged release commits version generation was wrong:
```
% git checkout release/1.13
% ./.github/scripts/generate_pytorch_version.py
ciflow/periodic/79617-4848-g7c98e70d44+cpu
```
After the fix, it is as expected:
```
% ./.github/scripts/generate_pytorch_version.py
1.13.0+cpu
```
Pull Request resolved: https://github.com/pytorch/pytorch/pull/88741
Approved by: https://github.com/dagitses, https://github.com/msaroufim

.github/scripts/generate_pytorch_version.py

@@ -23,27 +23,22 @@ def get_pytorch_root() -> Path:
 
 def get_tag() -> str:
     root = get_pytorch_root()
-    # We're on a tag
-    am_on_tag = (
-        subprocess.run(
-            ['git', 'describe', '--tags', '--exact'],
-            cwd=root,
-            stdout=subprocess.DEVNULL,
-            stderr=subprocess.DEVNULL
-        ).returncode == 0
-    )
-    tag = ""
-    if am_on_tag:
+    try:
         dirty_tag = subprocess.check_output(
-            ['git', 'describe'],
+            ['git', 'describe', '--tags', '--exact'],
             cwd=root
         ).decode('ascii').strip()
-        # Strip leading v that we typically do when we tag branches
-        # ie: v1.7.1 -> 1.7.1
-        tag = re.sub(LEADING_V_PATTERN, "", dirty_tag)
-        # Strip trailing rc pattern
-        # ie: 1.7.1-rc1 -> 1.7.1
-        tag = re.sub(TRAILING_RC_PATTERN, "", tag)
+    except subprocess.CalledProcessError:
+        return ""
+    # Strip leading v that we typically do when we tag branches
+    # ie: v1.7.1 -> 1.7.1
+    tag = re.sub(LEADING_V_PATTERN, "", dirty_tag)
+    # Strip trailing rc pattern
+    # ie: 1.7.1-rc1 -> 1.7.1
+    tag = re.sub(TRAILING_RC_PATTERN, "", tag)
+    # Ignore ciflow tags
+    if tag.startswith("ciflow/"):
+        return ""
     return tag
 
 def get_base_version() -> str:

.github/workflows/docker-release.yml

@@ -80,10 +80,13 @@ jobs:
           # Generate PyTorch version to use
           echo "PYTORCH_VERSION=$(python3 .github/scripts/generate_pytorch_version.py)" >> "${GITHUB_ENV}"
       - name: Setup nightly specific variables
-        if: ${{ github.event.ref == 'refs/heads/nightly' }}
+        if: ${{ github.event.ref == 'refs/heads/nightly' || startsWith(github.event.ref, 'refs/tags/ciflow/nightly/') }}
         run: |
-          # Use nightly image if building for nightly
-          echo "DOCKER_IMAGE=pytorch-nightly" >> "${GITHUB_ENV}"
+          {
+            echo "DOCKER_IMAGE=pytorch-nightly";
+            echo "INSTALL_CHANNEL=pytorch-nightly";
+            echo "TRITON_VERSION=2.0.0+$(cut -c -10 .github/ci_commit_pins/triton.txt)";
+          } >> "${GITHUB_ENV}"
       - name: Run docker build / push
         # WITH_PUSH is used here to determine whether or not to add the --push flag
         run: |

Dockerfile

@@ -66,6 +66,7 @@ ARG INSTALL_CHANNEL=pytorch-nightly
 RUN /opt/conda/bin/conda update -y conda
 RUN /opt/conda/bin/conda install -c "${INSTALL_CHANNEL}" -y python=${PYTHON_VERSION}
 ARG TARGETPLATFORM
+ARG TRITON_VERSION
 
 # On arm64 we can only install wheel packages
 RUN case ${TARGETPLATFORM} in \
@@ -74,6 +75,7 @@ RUN case ${TARGETPLATFORM} in \
     esac && \
     /opt/conda/bin/conda clean -ya
 RUN /opt/conda/bin/pip install torchelastic
+RUN if test -n "${TRITON_VERSION}" -a "${TARGETPLATFORM}" != "linux/arm64"; then /opt/conda/bin/pip install "torchtriton==${TRITON_VERSION}" --extra-index-url https://download.pytorch.org/whl/nightly/cpu ; fi
 
 FROM ${BASE_IMAGE} as official
 ARG PYTORCH_VERSION

docker.Makefile

@@ -23,12 +23,15 @@ PYTORCH_VERSION          ?= $(shell git describe --tags --always)
 # Can be either official / dev
 BUILD_TYPE               ?= dev
 BUILD_PROGRESS           ?= auto
+# Intentionally left blank
+TRITON_VERSION           ?=
 BUILD_ARGS                = --build-arg BASE_IMAGE=$(BASE_IMAGE) \
 							--build-arg PYTHON_VERSION=$(PYTHON_VERSION) \
 							--build-arg CUDA_VERSION=$(CUDA_VERSION) \
 							--build-arg CUDA_CHANNEL=$(CUDA_CHANNEL) \
 							--build-arg PYTORCH_VERSION=$(PYTORCH_VERSION) \
-							--build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL)
+							--build-arg INSTALL_CHANNEL=$(INSTALL_CHANNEL) \
+							--build-arg TRITON_VERSION=$(TRITON_VERSION)
 EXTRA_DOCKER_BUILD_FLAGS ?=
 
 BUILD                    ?= build