File size: 8,663 Bytes
9c6594c |
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 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 |
import json
import logging
import os
import shutil
import subprocess
from typing import Any, Dict, List, Optional, Tuple, Union
from wandb.docker import names
from wandb.errors import Error
class DockerError(Error):
"""Raised when attempting to execute a docker command."""
def __init__(
self,
command_launched: List[str],
return_code: int,
stdout: Optional[bytes] = None,
stderr: Optional[bytes] = None,
) -> None:
command_launched_str = " ".join(command_launched)
error_msg = (
f"The docker command executed was `{command_launched_str}`.\n"
f"It returned with code {return_code}\n"
)
if stdout is not None:
error_msg += f"The content of stdout is '{stdout.decode()}'\n"
else:
error_msg += (
"The content of stdout can be found above the "
"stacktrace (it wasn't captured).\n"
)
if stderr is not None:
error_msg += f"The content of stderr is '{stderr.decode()}'\n"
else:
error_msg += (
"The content of stderr can be found above the "
"stacktrace (it wasn't captured)."
)
super().__init__(error_msg)
entrypoint = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "wandb-entrypoint.sh"
)
log = logging.getLogger(__name__)
def shell(cmd: List[str]) -> Optional[str]:
"""Simple wrapper for calling docker,.
returning None on error and the output on success
"""
try:
return (
subprocess.check_output(["docker"] + cmd, stderr=subprocess.STDOUT)
.decode("utf8")
.strip()
)
except subprocess.CalledProcessError as e:
print(e) # noqa: T201
return None
_buildx_installed = None
def is_buildx_installed() -> bool:
"""Return `True` if docker buildx is installed and working."""
global _buildx_installed
if _buildx_installed is not None:
return _buildx_installed # type: ignore
if not shutil.which("docker"):
_buildx_installed = False
else:
help_output = shell(["buildx", "--help"])
_buildx_installed = help_output is not None and "buildx" in help_output
return _buildx_installed
def is_docker_installed() -> bool:
"""Return `True` if docker is installed and working, else `False`."""
try:
# Run the docker --version command
result = subprocess.run(
["docker", "--version"],
capture_output=True,
)
if result.returncode == 0:
return True
else:
return False
except FileNotFoundError:
# If docker command is not found
return False
def build(
tags: List[str], file: str, context_path: str, platform: Optional[str] = None
) -> str:
use_buildx = is_buildx_installed()
command = ["buildx", "build"] if use_buildx else ["build"]
command += ["--load"] if should_add_load_argument(platform) and use_buildx else []
if platform:
command += ["--platform", platform]
build_tags = []
for tag in tags:
build_tags += ["-t", tag]
args = ["docker"] + command + build_tags + ["-f", file, context_path]
stdout = run_command_live_output(
args,
)
return stdout
def should_add_load_argument(platform: Optional[str]) -> bool:
# the load option does not work when multiple platforms are specified:
# https://github.com/docker/buildx/issues/59
if platform is None or (platform and "," not in platform):
return True
return False
def run_command_live_output(args: List[Any]) -> str:
with subprocess.Popen(
args,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
universal_newlines=True,
bufsize=1,
) as process:
stdout = ""
while True:
chunk = os.read(process.stdout.fileno(), 4096) # type: ignore
if not chunk:
break
index = chunk.find(b"\r")
if index != -1:
print(chunk.decode(), end="") # noqa: T201
else:
stdout += chunk.decode()
print(chunk.decode(), end="\r") # noqa: T201
print(stdout) # noqa: T201
return_code = process.wait()
if return_code != 0:
raise DockerError(args, return_code, stdout.encode())
return stdout
def run(
args: List[Any],
capture_stdout: bool = True,
capture_stderr: bool = True,
input: Optional[bytes] = None,
return_stderr: bool = False,
env: Optional[Dict[str, str]] = None,
) -> Union[str, Tuple[str, str]]:
args = [str(x) for x in args]
subprocess_env = dict(os.environ)
subprocess_env.update(env or {})
if args[1] == "buildx":
subprocess_env["DOCKER_CLI_EXPERIMENTAL"] = "enabled"
stdout_dest: Optional[int] = subprocess.PIPE if capture_stdout else None
stderr_dest: Optional[int] = subprocess.PIPE if capture_stderr else None
completed_process = subprocess.run(
args, input=input, stdout=stdout_dest, stderr=stderr_dest, env=subprocess_env
)
if completed_process.returncode != 0:
raise DockerError(
args,
completed_process.returncode,
completed_process.stdout,
completed_process.stderr,
)
if return_stderr:
return (
_post_process_stream(completed_process.stdout),
_post_process_stream(completed_process.stderr),
)
else:
return _post_process_stream(completed_process.stdout)
def _post_process_stream(stream: Optional[bytes]) -> str:
if stream is None:
return ""
decoded_stream = stream.decode()
if len(decoded_stream) != 0 and decoded_stream[-1] == "\n":
decoded_stream = decoded_stream[:-1]
return decoded_stream
def default_image(gpu: bool = False) -> str:
tag = "all"
if not gpu:
tag += "-cpu"
return f"wandb/deepo:{tag}"
def parse_repository_tag(repo_name: str) -> Tuple[str, Optional[str]]:
parts = repo_name.rsplit("@", 1)
if len(parts) == 2:
return parts[0], parts[1]
parts = repo_name.rsplit(":", 1)
if len(parts) == 2 and "/" not in parts[1]:
return parts[0], parts[1]
return repo_name, None
def parse(image_name: str) -> Tuple[str, str, str]:
repository, tag = parse_repository_tag(image_name)
registry, repo_name = names.resolve_repository_name(repository)
if registry == "docker.io":
registry = "index.docker.io"
return registry, repo_name, (tag or "latest")
def image_id_from_registry(image_name: str) -> Optional[str]:
"""Query the image manifest to get its full ID including the digest.
Args:
image_name: The image name, such as "wandb/local".
Returns:
The image name followed by its digest, like "wandb/local@sha256:...".
"""
# https://docs.docker.com/reference/cli/docker/buildx/imagetools/inspect
inspect_cmd = ["buildx", "imagetools", "inspect", image_name]
format_args = ["--format", r"{{.Name}}@{{.Manifest.Digest}}"]
return shell([*inspect_cmd, *format_args])
def image_id(image_name: str) -> Optional[str]:
"""Retrieve the image id from the local docker daemon or remote registry."""
if "@sha256:" in image_name:
return image_name
else:
digests = shell(["inspect", image_name, "--format", "{{json .RepoDigests}}"])
if digests is None:
return image_id_from_registry(image_name)
try:
return json.loads(digests)[0]
except (ValueError, IndexError):
return image_id_from_registry(image_name)
def get_image_uid(image_name: str) -> int:
"""Retrieve the image default uid through brute force."""
image_uid = shell(["run", image_name, "id", "-u"])
return int(image_uid) if image_uid else -1
def push(image: str, tag: str) -> Optional[str]:
"""Push an image to a remote registry."""
return shell(["push", f"{image}:{tag}"])
def login(username: str, password: str, registry: str) -> Optional[str]:
"""Login to a registry."""
return shell(["login", "--username", username, "--password", password, registry])
def tag(image_name: str, tag: str) -> Optional[str]:
"""Tag an image."""
return shell(["tag", image_name, tag])
__all__ = [
"shell",
"build",
"run",
"image_id",
"image_id_from_registry",
"is_docker_installed",
"parse",
"parse_repository_tag",
"default_image",
"get_image_uid",
"push",
"login",
"tag",
]
|