From 3193102e13bbd228a8466e308db34cde8c974afe Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Tue, 1 Nov 2022 00:47:06 -0700
Subject: [PATCH] feat(cli): log conditional environment variables (#3156)

---
 src/bentoml_cli/env.py | 86 +++++++++++++++++++++++++++++++++---------
 1 file changed, 69 insertions(+), 17 deletions(-)

diff --git a/src/bentoml_cli/env.py b/src/bentoml_cli/env.py
index f9460e21c8..7a41248ebb 100644
--- a/src/bentoml_cli/env.py
+++ b/src/bentoml_cli/env.py
@@ -2,6 +2,7 @@
 
 import os
 import sys
+import shlex
 import shutil
 import typing as t
 import platform
@@ -12,12 +13,38 @@
 
 conda_packages_name = "conda_packages"
 
+_ENVVAR = [
+    "BENTOML_DEBUG",
+    "BENTOML_QUIET",
+    "BENTOML_BUNDLE_LOCAL_BUILD",
+    "BENTOML_DO_NOT_TRACK",
+    "BENTOML_CONFIG",
+    "BENTOML_CONFIG_OPTIONS",
+    "BENTOML_PORT",
+    "BENTOML_HOST",
+    "BENTOML_API_WORKERS",
+]
+_CONDITIONAL_ENVVAR = [
+    # Only print if set
+    "BENTOML_RETRY_RUNNER_REQUESTS",
+    "BENTOML_NUM_THREAD",
+    "OMP_NUM_THREADS",
+    "OPENBLAS_NUM_THREADS",
+    "MKL_NUM_THREADS",
+    "VECLIB_MAXIMUM_THREADS",
+    "NUMEXPR_NUM_THREADS",
+    "RAYON_RS_NUM_CPUS",
+    "TF_NUM_INTEROP_THREADS",
+    "TF_NUM_INTRAOP_THREADS",
+    "CUDA_VISIBLE_DEVICES",
+]
+
 
 def run_cmd(cmd: list[str]) -> list[str]:
     return subprocess.check_output(cmd).decode("utf-8").split("\n")[:-1]
 
 
-def format_dropdown(title: str, content: t.Iterable[str]) -> str:
+def _format_dropdown(title: str, content: t.Iterable[str]) -> str:
     processed = "\n".join(content)
     return f"""\
 <details><summary><code>{title}</code></summary>
@@ -32,23 +59,48 @@ def format_dropdown(title: str, content: t.Iterable[str]) -> str:
 """
 
 
-def format_keys(key: str, markdown: bool) -> str:
-    return f"`{key}`" if markdown else key
+def _format_env(env: list[str]) -> list[str]:
+    return [f"{e}={shlex.quote(os.environ.get(e, ''))}" for e in env]
 
 
-def pretty_format(
-    info_dict: dict[str, str | list[str]], output: t.Literal["md", "plain"]
-) -> str:
+def format_md(env: list[str], info_dict: dict[str, str | list[str]]) -> list[str]:
+    out: list[str] = []
+    out.append(
+        """\
+#### Environment variable
+
+```bash
+{env}
+```
+""".format(
+            env="\n".join(_format_env(env))
+        )
+    )
+    out.append("#### System information\n")
+    for key, value in info_dict.items():
+        if isinstance(value, list):
+            out.append(_format_dropdown(key, value))
+        else:
+            out.append(f"`{key}`: {value}")
+    return out
+
+
+def format_bash(env: list[str], info_dict: dict[str, str | list[str]]) -> list[str]:
     out: list[str] = []
+    out.extend(_format_env(env))
     for key, value in info_dict.items():
         if isinstance(value, list):
-            if output == "md":
-                out.append(format_dropdown(key, value))
-            else:
-                out.append(f"{key}:\n  " + "\n  ".join(value))
+            out.append(f"{key}=( " + " ".join(map(lambda s: f'"{s}"', value)) + " )")
         else:
-            out.append(f"{format_keys(key, markdown=output=='md')}: {value}")
-    return "\n".join(out)
+            out.append(f'{key}="{value}"')
+    return out
+
+
+def pretty_format(
+    info_dict: dict[str, str | list[str]], output: t.Literal["md", "bash"]
+) -> str:
+    env = _ENVVAR + list(filter(lambda e: e in os.environ, _CONDITIONAL_ENVVAR))
+    return "\n".join({"md": format_md, "bash": format_bash}[output](env, info_dict))
 
 
 def add_env_command(cli: click.Group) -> None:
@@ -61,14 +113,14 @@ def add_env_command(cli: click.Group) -> None:
     @click.option(
         "-o",
         "--output",
-        type=click.Choice(["md", "plain"]),
+        type=click.Choice(["md", "bash"]),
         default="md",
         show_default=True,
-        help="Output format. '-o plain' to display without format.",
+        help="Output format. '-o bash' to display without format.",
     )
     @click.pass_context
-    def env(ctx: click.Context, output: t.Literal["md", "plain"]) -> None:  # type: ignore (unused warning)
-        if output not in ["md", "plain"]:
+    def env(ctx: click.Context, output: t.Literal["md", "bash"]) -> None:  # type: ignore (unused warning)
+        if output not in ["md", "bash"]:
             raise CLIException(f"Unknown output format: {output}")
 
         is_windows = sys.platform == "win32"
@@ -86,7 +138,7 @@ def env(ctx: click.Context, output: t.Literal["md", "plain"]) -> None:  # type:
             is_admin: bool = windll.shell32.IsUserAnAdmin() != 0
             info_dict["is_window_admin"] = str(is_admin)
         else:
-            info_dict["uid:gid"] = f"{os.geteuid()}:{os.getegid()}"
+            info_dict["uid_gid"] = f"{os.geteuid()}:{os.getegid()}"
 
         if "CONDA_PREFIX" in os.environ:
             # conda packages