/
env.py
169 lines (138 loc) · 4.89 KB
/
env.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
from __future__ import annotations
import os
import sys
import shlex
import shutil
import typing as t
import platform
import subprocess
from gettext import gettext
import click
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",
]
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:
processed = "\n".join(content)
return f"""\
<details><summary><code>{title}</code></summary>
<br>
```{ 'yaml' if title == conda_packages_name else '' }
{processed}
```
</details>
"""
def _format_env(env: list[str]) -> list[str]:
return [f"{e}={shlex.quote(os.environ.get(e, ''))}" for e in env]
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):
out.append(f"{key}=( " + " ".join(map(lambda s: f'"{s}"', value)) + " )")
else:
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:
from bentoml import __version__ as BENTOML_VERSION
from bentoml.exceptions import CLIException
from bentoml._internal.utils.pkg import get_pkg_version
from bentoml._internal.utils.pkg import PackageNotFoundError
@cli.command(help=gettext("Print environment info and exit"))
@click.option(
"-o",
"--output",
type=click.Choice(["md", "bash"]),
default="md",
show_default=True,
help="Output format. '-o bash' to display without format.",
)
@click.pass_context
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"
info_dict: dict[str, str | list[str]] = {
"bentoml": BENTOML_VERSION,
"python": platform.python_version(),
"platform": platform.platform(),
}
if is_windows:
from ctypes import windll
# https://stackoverflow.com/a/1026626
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()}"
if "CONDA_PREFIX" in os.environ:
# conda packages
conda_like = None
for possible_exec in ["conda", "mamba", "micromamba"]:
if shutil.which(possible_exec) is not None:
conda_like = possible_exec
break
assert (
conda_like is not None
), "couldn't find a conda-like executable, while CONDA_PREFIX is set."
conda_packages = run_cmd([conda_like, "env", "export"])
# user is currently in a conda environment,
# doing this is faster than invoking `conda --version`
try:
conda_version = get_pkg_version("conda")
except PackageNotFoundError:
conda_version = run_cmd([conda_like, "--version"])[0].split(" ")[-1]
info_dict[conda_like] = conda_version
info_dict["in_conda_env"] = str(True)
info_dict["conda_packages"] = conda_packages
# process info from `pip freeze`
pip_packages = run_cmd(["pip", "freeze"])
info_dict["pip_packages"] = pip_packages
click.echo(pretty_format(info_dict, output=output))
ctx.exit(0)