Support detached checkout in tools/nightly.py (#154314)

Prompt for Sonnet 3.7 in Claude Code: Only inspect tools/nightly.py, all
other files are irrelevant to your task. Do not use any shell commands.
Task: Add a --detach argument to this script which instead of making a
new branch just directly checks out the correct commit in detached mode.

With two interventions:
- Branch and detach are mutually exclusive. So you should consolidate
  them into a single argument. Why don't we take over the 'None' option?
- Do you know that nightly_version is guaranteed to be a commit hash? It
  seems it would be safer to explicitly pass --detach

I tested by running `python tools/nightly.py checkout` and observing
that my worktree was detached at this point.

Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/154314
Approved by: https://github.com/XuehaiPan, https://github.com/malfet

tools/nightly.py

@@ -9,6 +9,11 @@ You can use this script to check out a new nightly branch with the following::
     $ ./tools/nightly.py checkout -b my-nightly-branch
     $ source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
 
+Or if you would like to check out the nightly commit in detached HEAD mode::
+
+    $ ./tools/nightly.py checkout
+    $ source venv/bin/activate  # or `& .\venv\Scripts\Activate.ps1` on Windows
+
 Or if you would like to re-use an existing virtual environment, you can pass in
 the prefix argument (--prefix)::
 
@@ -613,19 +618,17 @@ def check_branch(subcommand: str, branch: str | None) -> str | None:
     """Checks that the branch name can be checked out."""
     if subcommand != "checkout":
         return None
-    # first make sure actual branch name was given
-    if branch is None:
-        return "Branch name to checkout must be supplied with '-b' option"
     # next check that the local repo is clean
     cmd = git("status", "--untracked-files=no", "--porcelain")
     stdout = subprocess.check_output(cmd, text=True, encoding="utf-8")
     if stdout.strip():
         return "Need to have clean working tree to checkout!\n\n" + stdout
-    # next check that the branch name doesn't already exist
-    cmd = git("show-ref", "--verify", "--quiet", f"refs/heads/{branch}")
-    p = subprocess.run(cmd, capture_output=True, check=False)  # type: ignore[assignment]
-    if not p.returncode:
-        return f"Branch {branch!r} already exists"
+    # next check that the branch name doesn't already exist (if a branch name is provided)
+    if branch is not None:
+        cmd = git("show-ref", "--verify", "--quiet", f"refs/heads/{branch}")
+        p = subprocess.run(cmd, capture_output=True, check=False)  # type: ignore[assignment]
+        if not p.returncode:
+            return f"Branch {branch!r} already exists"
     return None
 
 
@@ -680,10 +683,15 @@ def _nightly_version(site_dir: Path) -> str:
 
 
 @timed("Checking out nightly PyTorch")
-def checkout_nightly_version(branch: str, site_dir: Path) -> None:
+def checkout_nightly_version(branch: str | None, site_dir: Path) -> None:
     """Get's the nightly version and then checks it out."""
     nightly_version = _nightly_version(site_dir)
-    cmd = git("checkout", "-b", branch, nightly_version)
+    if branch is None:
+        # Detached mode - explicitly use --detach flag
+        cmd = git("checkout", "--detach", nightly_version)
+    else:
+        # Branch mode
+        cmd = git("checkout", "-b", branch, nightly_version)
     subprocess.check_call(cmd)
 
 
@@ -860,7 +868,7 @@ def install(
 
     with venv.extracted_wheel(torch_wheel) as wheel_site_dir:
         if subcommand == "checkout":
-            checkout_nightly_version(cast(str, branch), wheel_site_dir)
+            checkout_nightly_version(branch, wheel_site_dir)
         elif subcommand == "pull":
             pull_nightly_version(wheel_site_dir)
         else:
@@ -893,7 +901,7 @@ def make_parser() -> argparse.ArgumentParser:
     checkout.add_argument(
         "-b",
         "--branch",
-        help="Branch name to checkout",
+        help="Branch name to checkout (if omitted, checks out in detached HEAD mode)",
         dest="branch",
         default=None,
         metavar="NAME",