[DRAFT] Let the local builder accept user-provided build context managers

[redartera]

Description

This change introduces a new optional build_ctx_manager class variable that allows users to provide a context manager to handle runtime setup and cleanup during Docker image builds.

This enables:

• Authentication with container registries
• Secret injection and environment variable setup
• Probably a few more use-cases (e.g. some cleanup if needed, etc...)

The context manager is called at the start of build_image() and and has a no-op default when not provided.

Example usage

Creating the context manager

import contextlib
import os
import subprocess
from textwrap import dedent

from flyte._internal.imagebuild.docker_builder import DockerImageBuilder


# This will hydrate the "GITHUB_CONFIG" secret during the image build
@contextlib.contextmanager
def init_gh_config():
    token = subprocess.check_output(["gh", "auth", "token"]).decode("utf-8").strip()
    # Create a git config that replaces https://github.com with https://${GITHUB_TOKEN}@github.com
    CFG = dedent(f"""\
    [url "https://{token}@github.com"]
        insteadOf = https://github.com
    """)
    os.environ["GITHUB_CONFIG"] = CFG
    try:
        yield
    finally:
        del os.environ["GITHUB_CONFIG"]


# This will ensure we are logged in ECR
@contextlib.contextmanager
def init_ecr_login():
    # Login to ECR
    subprocess.check_output(
        ["docker", "login", "--username", "AWS", "--password-stdin", "1234567890.dkr.ecr.us-west-2.amazonaws.com"],
        input=subprocess.check_output(["aws", "ecr", "get-login-password", "--region", "us-west-2"]).decode("utf-8").strip()
    )
    yield



@contextlib.contextmanager
def init_all():
    with init_gh_config(), init_ecr_login():
            yield


DockerImageBuilder.build_ctx_manager = init_all()

Building an image

import pathlib

import flyte as flv2


_PROJECT_ROOT = pathlib.Path(__file__).parent

image = (
    flv2.Image.from_debian_base()
    .with_uv_project(
        pyproject_file=_PROJECT_ROOT / "pyproject.toml",
        uvlock=_PROJECT_ROOT / "uv.lock",
        # Notice here how the secret 'GITHUB_CONFIG' is bound to this layer
        secret_mounts=flv2.Secret(key="GITHUB_CONFIG", mount=pathlib.Path("/root/.gitconfig")),
    )
)

env = flv2.TaskEnvironment(name="example", image=image)


@env.task
def example_task():
    print("Hello, world!")

[redartera]

P.S: Given this would be a user-facing API - users needing to do from flyte._internal... import ... may not be ideal - but we could create a simple public API function that sets DockerImageBuilder.build_ctx_manager under the hood - currently still a draft

[kumare3]

@redartera let us think about this and we will get back to you soon