|
10 | 10 | import os
|
11 | 11 | import subprocess
|
12 | 12 | import sys
|
| 13 | +import shlex |
13 | 14 | import tempfile
|
14 | 15 | import urllib.parse
|
15 | 16 | from distutils import dir_util
|
16 | 17 | from subprocess import Popen, PIPE, STDOUT
|
17 | 18 | from typing import Union, List, Tuple, Dict
|
18 |
| -import shlex |
19 | 19 | import yaml
|
20 | 20 |
|
21 | 21 | import ads
|
|
33 | 33 | )
|
34 | 34 |
|
35 | 35 |
|
| 36 | +CONTAINER_NETWORK = "CONTAINER_NETWORK" |
| 37 | + |
| 38 | + |
36 | 39 | def get_service_pack_prefix() -> Dict:
|
37 | 40 | curr_dir = os.path.dirname(os.path.abspath(__file__))
|
38 | 41 | service_config_file = os.path.join(curr_dir, "conda", "config.yaml")
|
@@ -175,11 +178,13 @@ def build_image(
|
175 | 178 | command += ["--build-arg", f"http_proxy={os.environ['http_proxy']}"]
|
176 | 179 | if os.environ.get("https_proxy"):
|
177 | 180 | command += ["--build-arg", f"https_proxy={os.environ['https_proxy']}"]
|
| 181 | + if os.environ.get(CONTAINER_NETWORK): |
| 182 | + command += ["--network", os.environ[CONTAINER_NETWORK]] |
178 | 183 | command += [os.path.abspath(curr_dir)]
|
179 |
| - logger.info(f"Build image with command {command}") |
| 184 | + logger.info("Build image with command %s", command) |
180 | 185 | proc = run_command(command)
|
181 | 186 | if proc.returncode != 0:
|
182 |
| - raise RuntimeError(f"Docker build failed.") |
| 187 | + raise RuntimeError("Docker build failed.") |
183 | 188 |
|
184 | 189 |
|
185 | 190 | def _get_image_name_dockerfile_target(type: str, gpu: bool) -> str:
|
@@ -321,27 +326,51 @@ def run_container(
|
321 | 326 | logger.info(f"command: {command}")
|
322 | 327 | logger.info(f"entrypoint: {entrypoint}")
|
323 | 328 |
|
| 329 | + # Print out the equivalent docker run command for debugging purpose |
| 330 | + docker_run_cmd = [">>> docker run --rm"] |
| 331 | + if entrypoint: |
| 332 | + docker_run_cmd.append(f"--entrypoint {entrypoint}") |
| 333 | + if env_vars: |
| 334 | + docker_run_cmd.extend([f"-e {key}={val}" for key, val in env_vars.items()]) |
| 335 | + if bind_volumes: |
| 336 | + docker_run_cmd.extend( |
| 337 | + [f'-v {source}:{bind.get("bind")}' for source, bind in bind_volumes.items()] |
| 338 | + ) |
| 339 | + docker_run_cmd.append(image) |
| 340 | + if command: |
| 341 | + docker_run_cmd.append(command) |
| 342 | + logger.debug(" ".join(docker_run_cmd)) |
| 343 | + |
324 | 344 | client = get_docker_client()
|
325 | 345 | try:
|
326 | 346 | client.api.inspect_image(image)
|
327 | 347 | except docker.errors.ImageNotFound:
|
328 |
| - logger.warn(f"Image {image} not found. Try pulling it now....") |
| 348 | + logger.warning(f"Image {image} not found. Try pulling it now....") |
329 | 349 | run_command(["docker", "pull", f"{image}"], None)
|
330 |
| - container = client.containers.run( |
331 |
| - image=image, |
332 |
| - volumes=bind_volumes, |
333 |
| - command=shlex.split(command) if command else None, |
334 |
| - environment=env_vars, |
335 |
| - detach=True, |
336 |
| - entrypoint=entrypoint, |
337 |
| - user=0, |
338 |
| - # auto_remove=True, |
339 |
| - ) |
340 |
| - logger.info(f"Container ID: {container.id}") |
341 |
| - |
342 |
| - for line in container.logs(stream=True, follow=True): |
343 |
| - logger.info(line.decode("utf-8").strip()) |
344 |
| - |
345 |
| - result = container.wait() |
346 |
| - container.remove() |
347 |
| - return result.get("StatusCode", -1) |
| 350 | + try: |
| 351 | + kwargs = {} |
| 352 | + if CONTAINER_NETWORK in os.environ: |
| 353 | + kwargs["network_mode"] = os.environ[CONTAINER_NETWORK] |
| 354 | + container = client.containers.run( |
| 355 | + image=image, |
| 356 | + volumes=bind_volumes, |
| 357 | + command=shlex.split(command) if command else None, |
| 358 | + environment=env_vars, |
| 359 | + detach=True, |
| 360 | + entrypoint=entrypoint, |
| 361 | + user=0, |
| 362 | + **kwargs |
| 363 | + # auto_remove=True, |
| 364 | + ) |
| 365 | + logger.info("Container ID: %s", container.id) |
| 366 | + for line in container.logs(stream=True, follow=True): |
| 367 | + print(line.decode("utf-8"), end="") |
| 368 | + |
| 369 | + result = container.wait() |
| 370 | + return result.get("StatusCode", -1) |
| 371 | + except docker.errors.APIError as ex: |
| 372 | + logger.error(ex.explanation) |
| 373 | + return -1 |
| 374 | + finally: |
| 375 | + # Remove the container |
| 376 | + container.remove() |
0 commit comments