Skip to content

Commit

Permalink
feat(cli): log conditional environment variables (#3156)
Browse files Browse the repository at this point in the history
  • Loading branch information
aarnphm committed Nov 1, 2022
1 parent 8365375 commit 3193102
Showing 1 changed file with 69 additions and 17 deletions.
86 changes: 69 additions & 17 deletions src/bentoml_cli/env.py
Expand Up @@ -2,6 +2,7 @@

import os
import sys
import shlex
import shutil
import typing as t
import platform
Expand All @@ -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>
Expand All @@ -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:
Expand All @@ -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"
Expand All @@ -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
Expand Down

0 comments on commit 3193102

Please sign in to comment.