From ff2f751bfbd23efd6ecb9067d9cb3ad2fa991798 Mon Sep 17 00:00:00 2001
From: Xuehai Pan <XuehaiPan@pku.edu.cn>
Date: Mon, 21 Oct 2024 17:37:03 +0800
Subject: [PATCH] [tools] fix nightly pull tool when the conda environment not
 exists (#138448)

Now, `conda env remove --name env` exits with errors if the given environment does not exist. This PR check the existance of the environment before trying to remove it.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/138448
Approved by: https://github.com/ezyang
---
 tools/nightly.py | 32 ++++++++++++++++++++++++++++----
 1 file changed, 28 insertions(+), 4 deletions(-)

diff --git a/tools/nightly.py b/tools/nightly.py
index 6d1dc486040..f09563a8c03 100755
--- a/tools/nightly.py
+++ b/tools/nightly.py
@@ -57,6 +57,7 @@ from typing import Any, Callable, cast, Generator, Iterable, Iterator, Sequence,
 REPO_ROOT = Path(__file__).absolute().parent.parent
 GITHUB_REMOTE_URL = "https://github.com/pytorch/pytorch.git"
 SPECS_TO_INSTALL = ("pytorch", "mypy", "pytest", "hypothesis", "ipython", "sphinx")
+DEFAULT_ENV_NAME = "pytorch-deps"
 
 LOGGER: logging.Logger | None = None
 URL_FORMAT = "{base_url}/{platform}/{dist_name}.tar.bz2"
@@ -212,6 +213,28 @@ def check_branch(subcommand: str, branch: str | None) -> str | None:
     return None
 
 
+def check_conda_env_exists(name: str | None = None, prefix: str | None = None) -> bool:
+    """Checks that the conda environment exists."""
+    if name is not None and prefix is not None:
+        raise ValueError("Cannot specify both --name and --prefix")
+    if name is None and prefix is None:
+        raise ValueError("Must specify either --name or --prefix")
+
+    try:
+        cmd = ["conda", "info", "--envs"]
+        output = subprocess.check_output(cmd, text=True, encoding="utf-8")
+    except subprocess.CalledProcessError:
+        logger = cast(logging.Logger, LOGGER)
+        logger.warning("Failed to list conda environments", exc_info=True)
+        return False
+
+    if name is not None:
+        return len(re.findall(rf"^{name}\s+", output, flags=re.MULTILINE)) > 0
+    assert prefix is not None
+    prefix = Path(prefix).absolute()
+    return len(re.findall(rf"\s+{prefix}$", output, flags=re.MULTILINE)) > 0
+
+
 @contextlib.contextmanager
 def timer(logger: logging.Logger, prefix: str) -> Iterator[None]:
     """Timed context manager"""
@@ -271,7 +294,7 @@ def conda_solve(
     else:
         # create new environment
         existing_env = False
-        env_opts = ["--name", "pytorch-deps"]
+        env_opts = ["--name", DEFAULT_ENV_NAME]
     # run solve
     if existing_env:
         cmd = [
@@ -280,8 +303,8 @@ def conda_solve(
             "--yes",
             "--dry-run",
             "--json",
+            *env_opts,
         ]
-        cmd.extend(env_opts)
     else:
         cmd = [
             "conda",
@@ -321,8 +344,9 @@ def deps_install(deps: list[str], existing_env: bool, env_opts: list[str]) -> No
     """Install dependencies to deps environment"""
     if not existing_env:
         # first remove previous pytorch-deps env
-        cmd = ["conda", "env", "remove", "--yes", *env_opts]
-        subprocess.check_call(cmd)
+        if check_conda_env_exists(name=DEFAULT_ENV_NAME):
+            cmd = ["conda", "env", "remove", "--yes", *env_opts]
+            subprocess.check_output(cmd)
     # install new deps
     install_command = "install" if existing_env else "create"
     cmd = ["conda", install_command, "--yes", "--no-deps", *env_opts, *deps]