daggerml.contrib.funks

 1from __future__ import annotations
 2
 3from daggerml.contrib import api
 4
 5
 6def _run(*cmd: str) -> None:
 7    import subprocess
 8
 9    from daggerml.api import DmlRepoError
10
11    proc = subprocess.run(cmd, check=False)
12    if proc.returncode == 0:
13        return
14    raise DmlRepoError(
15        f"Command failed (exit code {proc.returncode}): {' '.join(cmd)}. See the execution logs for command output."
16    )
17
18
19def _gzip_file(source: str, destination: str) -> None:
20    import gzip
21    import shutil
22
23    with open(source, "rb") as src, gzip.open(destination, "wb") as dst:
24        shutil.copyfileobj(src, dst)
25
26
27def _remove_docker_image(image: str) -> None:
28    import subprocess
29
30    subprocess.run(("docker", "image", "rm", "-f", image), check=False, capture_output=True, text=True)
31
32
33@api.funkify(uri="script", adapter="local", extra_objs=(_run, _gzip_file, _remove_docker_image))
34def docker_build(dag, context_tarball, build_flags=(), repo=None):
35    from contextlib import chdir
36    from tempfile import TemporaryDirectory
37    from uuid import uuid4
38
39    from daggerml import Uri
40    from daggerml.contrib.s3 import S3Store
41
42    build_flags = tuple(build_flags.value())
43
44    store = S3Store()
45    tag = uuid4().hex
46    local_image = f"dml:{tag}"
47    with TemporaryDirectory(prefix="dml-docker-build-") as build_dir:
48        store.untar(context_tarball.value(), build_dir)
49        try:
50            with chdir(build_dir):
51                _run("docker", "build", *build_flags, "-t", local_image, ".")
52                repo = repo.value() if repo is not None else None
53                if repo is not None:
54                    remote_image = f"{repo.uri}:{tag}"
55                    try:
56                        _run("docker", "tag", local_image, remote_image)
57                        _run("docker", "push", remote_image)
58                    finally:
59                        _remove_docker_image(remote_image)
60                    return dag.put(Uri(remote_image), name="remote-image")
61                image_tar = "./image.tar"
62                _run("docker", "save", "-o", str(image_tar), local_image)
63                compressed_image_tar = "./image.tar.gz"
64                _gzip_file(image_tar, compressed_image_tar)
65                return store.put(filepath=compressed_image_tar, suffix=".tar.gz")
66        finally:
67            _remove_docker_image(local_image)
68
69
70__all__ = ["docker_build"]