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", 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.executors.docker import _image_lock
41    from daggerml.contrib.s3 import S3Store
42
43    build_flags = tuple(build_flags.value())
44
45    store = S3Store()
46    tag = uuid4().hex
47    local_image = f"daggerml-archive:{tag}"
48    with TemporaryDirectory(prefix="dml-docker-build-") as build_dir:
49        store.untar(context_tarball.value(), build_dir)
50        with chdir(build_dir):
51            repo = repo.value() if repo is not None else None
52            image_tar = "./image.tar"
53            with _image_lock():
54                try:
55                    _run("docker", "build", *build_flags, "-t", local_image, ".")
56                    if repo is not None:
57                        remote_image = f"{repo.uri}:{tag}"
58                        try:
59                            _run("docker", "tag", local_image, remote_image)
60                            _run("docker", "push", remote_image)
61                        finally:
62                            _remove_docker_image(remote_image)
63                    else:
64                        _run("docker", "save", "-o", str(image_tar), local_image)
65                finally:
66                    _remove_docker_image(local_image)
67            if repo is not None:
68                return dag.put(Uri(remote_image), name="remote-image")
69            compressed_image_tar = "./image.tar.gz"
70            _gzip_file(image_tar, compressed_image_tar)
71            return store.put(filepath=compressed_image_tar, suffix=".tar.gz")
72
73
74__all__ = ["docker_build"]