wandb

1
"""Build and optinally push the launch agent image."""
2
import argparse
3
import os
4

5
from wandb.docker import build, push
6

7
DOCKERFILE_TEMPLATE = """
8
FROM python:3.9-bullseye
9
LABEL maintainer='Weights & Biases <support@wandb.com>'
10

11
# install git
12
RUN apt-get update && apt-get upgrade -y \
13
    && apt-get install -y git \
14
    && apt-get -qy autoremove \
15
    && apt-get clean && rm -r /var/lib/apt/lists/*
16

17
# copy reqs and install for docker cache
18
COPY ../requirements.txt /src/requirements.txt
19
RUN pip install --no-cache-dir -r /src/requirements.txt
20
RUN pip install awscli nbconvert nbformat chardet iso8601 typing_extensions boto3 botocore google-auth google-cloud-compute google-cloud-storage google-cloud-artifact-registry kubernetes
21

22

23
# Copy source code and install
24
COPY .. /src
25
RUN pip install --no-cache-dir "/src[launch]"
26

27
# user set up
28
RUN useradd -m -s /bin/bash --create-home --no-log-init -u 1000 -g 0 launch_agent
29
USER launch_agent
30
WORKDIR /home/launch_agent
31
RUN chown -R launch_agent /home/launch_agent
32

33
{version_section}
34

35
ENTRYPOINT ["wandb", "launch-agent"]
36
"""
37

38
VERSION_SECTION = """
39
# set agent version env var
40
ENV WANDB_AGENT_VERSION={agent_version}
41
"""
42

43
DOCKERIGNORE = """
44
.tox/
45
"""
46

47

48
def parse_args():
49
    """Parse command line arguments."""
50
    parser = argparse.ArgumentParser()
51
    parser.add_argument(
52
        "--push",
53
        action="store_true",
54
        help="Push image after creation. This requires that you enter a tag that includes a registry via --tag",
55
    )
56
    parser.add_argument("--tag", default="wandb-launch-agent", help="Tag for the image")
57
    parser.add_argument(
58
        "--platform", default="linux/amd64", help="Platform to use, e.g. 'linux/amd64'"
59
    )
60
    return parser.parse_args()
61

62

63
def main():
64
    """Build the launch agent image."""
65
    args = parse_args()
66
    build_context = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
67
    dockerfile_path = os.path.join(build_context, "Dockerfile")
68
    dockerignore_path = os.path.join(build_context, ".dockerignore")
69
    # Set the version env var if a custom tag is set
70
    version_section = ""
71
    if args.tag != "wandb-launch-agent":
72
        version_section = VERSION_SECTION.format(agent_version=args.tag)
73
    dockerfile_contents = DOCKERFILE_TEMPLATE.format(version_section=version_section)
74
    with open(dockerfile_path, "w") as f:
75
        f.write(dockerfile_contents)
76
    with open(dockerignore_path, "w") as f:
77
        f.write(DOCKERIGNORE)
78
    build(
79
        tags=[args.tag],
80
        file=dockerfile_path,
81
        context_path=build_context,
82
        platform=args.platform,
83
    )
84
    if args.push:
85
        image, tag = args.tag.split(":")
86
        push(image, tag)
87

88
    # Remove the dockerfui
89
    os.remove(dockerfile_path)
90
    os.remove(dockerignore_path)
91

92

93
if __name__ == "__main__":
94
    main()
95