Skip to content

Commit 54cb0ff

Browse files
committed
Update utils.py to set docker run/build network based on env var.
1 parent f30b798 commit 54cb0ff

File tree

1 file changed

+11
-2
lines changed

1 file changed

+11
-2
lines changed

ads/opctl/utils.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,9 @@
3333
)
3434

3535

36+
CONTAINER_NETWORK = "CONTAINER_NETWORK"
37+
38+
3639
def get_service_pack_prefix() -> Dict:
3740
curr_dir = os.path.dirname(os.path.abspath(__file__))
3841
service_config_file = os.path.join(curr_dir, "conda", "config.yaml")
@@ -175,11 +178,13 @@ def build_image(
175178
command += ["--build-arg", f"http_proxy={os.environ['http_proxy']}"]
176179
if os.environ.get("https_proxy"):
177180
command += ["--build-arg", f"https_proxy={os.environ['https_proxy']}"]
181+
if os.environ.get(CONTAINER_NETWORK):
182+
command += ["--network", os.environ[CONTAINER_NETWORK]]
178183
command += [os.path.abspath(curr_dir)]
179-
logger.info(f"Build image with command {command}")
184+
logger.info("Build image with command %s", command)
180185
proc = run_command(command)
181186
if proc.returncode != 0:
182-
raise RuntimeError(f"Docker build failed.")
187+
raise RuntimeError("Docker build failed.")
183188

184189

185190
def _get_image_name_dockerfile_target(type: str, gpu: bool) -> str:
@@ -343,6 +348,9 @@ def run_container(
343348
logger.warning(f"Image {image} not found. Try pulling it now....")
344349
run_command(["docker", "pull", f"{image}"], None)
345350
try:
351+
kwargs = {}
352+
if CONTAINER_NETWORK in os.environ:
353+
kwargs["network_mode"] = os.environ[CONTAINER_NETWORK]
346354
container = client.containers.run(
347355
image=image,
348356
volumes=bind_volumes,
@@ -351,6 +359,7 @@ def run_container(
351359
detach=True,
352360
entrypoint=entrypoint,
353361
user=0,
362+
**kwargs
354363
# auto_remove=True,
355364
)
356365
logger.info("Container ID: %s", container.id)

0 commit comments

Comments
 (0)