Usability fixes to CI runner script (#9752)

* Usability fixes to CI runner script

* address comments

Co-authored-by: driazati <[email protected]>

.gitignore

@@ -247,3 +247,4 @@ conda/pkg
 _docs/
 jvm/target
 .config/configstore/
+.ci-py-scripts/

tests/scripts/ci.py

@@ -24,12 +24,15 @@
 import getpass
 import inspect
 import argparse
+import json
+import shutil
 import grp
 import subprocess
 from pathlib import Path
 from typing import List, Dict, Any, Optional
 
 REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+SCRIPT_DIR = REPO_ROOT / ".ci-py-scripts"
 NPROC = multiprocessing.cpu_count()
 
 
@@ -44,48 +47,121 @@ class col:
     UNDERLINE = "\033[4m"
 
 
-def print_color(color: str, msg: str, **kwargs: Any) -> None:
+def print_color(color: str, msg: str, bold: bool, **kwargs: Any) -> None:
     if hasattr(sys.stdout, "isatty") and sys.stdout.isatty():
-        print(col.BOLD + color + msg + col.RESET, **kwargs)
+        bold_code = col.BOLD if bold else ""
+        print(bold_code + color + msg + col.RESET, **kwargs)
     else:
         print(msg, **kwargs)
 
 
+warnings = []
+
+
 def clean_exit(msg: str) -> None:
-    print_color(col.RED, msg, file=sys.stderr)
+    print_color(col.RED, msg, bold=True, file=sys.stderr)
+
+    for warning in warnings:
+        print_color(col.YELLOW, warning, bold=False, file=sys.stderr)
+
     exit(1)
 
 
 def cmd(commands: List[Any], **kwargs: Any):
     commands = [str(s) for s in commands]
     command_str = " ".join(commands)
-    print_color(col.BLUE, command_str)
+    print_color(col.BLUE, command_str, bold=True)
     proc = subprocess.run(commands, **kwargs)
     if proc.returncode != 0:
         raise RuntimeError(f"Command failed: '{command_str}'")
+    return proc
 
 
-def docker(name: str, image: str, scripts: List[str], env: Dict[str, str]):
-    """
-    Invoke a set of bash scripts through docker/bash.sh
-    """
+def check_docker():
+    executable = shutil.which("docker")
+    if executable is None:
+        clean_exit("'docker' executable not found, install it first (e.g. 'apt install docker.io')")
+
     if sys.platform == "linux":
         # Check that the user is in the docker group before running
         try:
             group = grp.getgrnam("docker")
             if getpass.getuser() not in group.gr_mem:
-                print_color(
-                    col.YELLOW, f"Note: User '{getpass.getuser()}' is not in the 'docker' group"
+                warnings.append(
+                    f"Note: User '{getpass.getuser()}' is not in the 'docker' group, either:\n"
+                    " * run with 'sudo'\n"
+                    " * add user to 'docker': sudo usermod -aG docker $(whoami), then log out and back in",
                 )
-        except KeyError:
-            print_color(col.YELLOW, f"Note: 'docker' group does not exist")
+        except KeyError as e:
+            warnings.append(f"Note: 'docker' group does not exist")
+
+
+def check_gpu():
+    if not (sys.platform == "linux" and shutil.which("lshw")):
+        # Can't check GPU on non-Linux platforms
+        return
+
+    # See if we can check if a GPU is present in case of later failures,
+    # but don't block on execution since this isn't critical
+    try:
+        proc = cmd(
+            ["lshw", "-json", "-C", "display"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            encoding="utf-8",
+        )
+        stdout = proc.stdout.strip().strip(",")
+        stdout = json.loads(stdout)
+    except (subprocess.CalledProcessError, json.decoder.JSONDecodeError) as e:
+        # Do nothing if any step failed
+        return
+
+    if isinstance(stdout, dict):
+        # Sometimes lshw outputs a single item as a dict instead of a list of
+        # dicts, so wrap it up if necessary
+        stdout = [stdout]
+    if not isinstance(stdout, list):
+        return
+
+    products = [s.get("product", "").lower() for s in stdout]
+    if not any("nvidia" in product for product in products):
+        warnings.append("nvidia GPU not found in 'lshw', maybe use --cpu flag?")
+
+
+def check_build():
+    if (REPO_ROOT / "build").exists():
+        warnings.append(
+            "Existing build dir found may be interfering with the Docker "
+            "build (you may need to remove it)"
+        )
+
+
+def docker(name: str, image: str, scripts: List[str], env: Dict[str, str]):
+    """
+    Invoke a set of bash scripts through docker/bash.sh
+
+    name: container name
+    image: docker image name
+    scripts: list of bash commands to run
+    env: environment to set
+    """
+    check_docker()
 
     docker_bash = REPO_ROOT / "docker" / "bash.sh"
     command = [docker_bash, "--name", name]
     for key, value in env.items():
         command.append("--env")
         command.append(f"{key}={value}")
-    command += [image, "bash", "-c", " && ".join(scripts)]
+
+    SCRIPT_DIR.mkdir(exist_ok=True)
+
+    script_file = SCRIPT_DIR / f"{name}.sh"
+    with open(script_file, "w") as f:
+        f.write("set -eux\n\n")
+        f.write("\n".join(scripts))
+        f.write("\n")
+
+    command += [image, "bash", str(script_file.relative_to(REPO_ROOT))]
 
     try:
         cmd(command)
@@ -110,17 +186,50 @@ def docs(
     full -- Build all language docs, not just Python
     precheck -- Run Sphinx precheck script
     tutorial-pattern -- Regex for which tutorials to execute when building docs (can also be set via TVM_TUTORIAL_EXEC_PATTERN)
-    cpu -- Use CMake defaults for building TVM (useful for building docs on a CPU machine.)
+    cpu -- Run with the ci-cpu image and use CMake defaults for building TVM (if no GPUs are available)
     """
     config = "./tests/scripts/task_config_build_gpu.sh"
     if cpu and full:
         clean_exit("--full cannot be used with --cpu")
 
+    extra_setup = []
+    image = "ci_gpu"
     if cpu:
-        # The docs import tvm.micro, so it has to be enabled in the build
-        config = "cd build && cp ../cmake/config.cmake . && echo set\(USE_MICRO ON\) >> config.cmake && cd .."
+        image = "ci_cpu"
+        config = " && ".join(
+            [
+                "mkdir -p build",
+                "pushd build",
+                "cp ../cmake/config.cmake .",
+                # The docs import tvm.micro, so it has to be enabled in the build
+                "echo set\(USE_MICRO ON\) >> config.cmake",
+                "popd",
+            ]
+        )
+
+        # These are taken from the ci-gpu image via pip freeze, consult that
+        # if there are any changes: https://github.com/apache/tvm/tree/main/docs#native
+        requirements = [
+            "Sphinx==4.2.0",
+            "tlcpack-sphinx-addon==0.2.1",
+            "synr==0.5.0",
+            "image==1.5.33",
+            "sphinx-gallery==0.4.0",
+            "sphinx-rtd-theme==1.0.0",
+            "matplotlib==3.3.4",
+            "commonmark==0.9.1",
+            "Pillow==8.3.2",
+            "autodocsumm==0.2.7",
+            "docutils==0.16",
+        ]
+
+        extra_setup = [
+            "python3 -m pip install --user " + " ".join(requirements),
+        ]
+    else:
+        check_gpu()
 
-    scripts = [
+    scripts = extra_setup + [
         config,
         f"./tests/scripts/task_build.sh build -j{NPROC}",
         "./tests/scripts/task_ci_setup.sh",
@@ -137,7 +246,8 @@ def docs(
         "PYTHON_DOCS_ONLY": "0" if full else "1",
         "IS_LOCAL": "1",
     }
-    docker(name="ci-docs", image="ci_gpu", scripts=scripts, env=env)
+    check_build()
+    docker(name="ci-docs", image=image, scripts=scripts, env=env)
 
 
 def serve_docs(directory: str = "_docs") -> None:
@@ -221,6 +331,10 @@ def main():
         add_subparser(func, subparsers)
 
     args = parser.parse_args()
+    if args.command is None:
+        parser.print_help()
+        exit(1)
+
     func = subparser_functions[args.command]
 
     # Extract out the parsed args and invoke the relevant function