Bring your own image

# Team A's image: data preparation
#
# This team owns this entire Dockerfile. They control Python version, WORKDIR,
# and PYTHONPATH. Flyte has no say here.
#
# Pure BYOI constraint: workflow code must be baked in because there is no
# code bundle. Every code change requires rebuilding and pushing this image.

FROM python:3.11-slim

# System deps this team needs
RUN apt-get update && apt-get install -y --no-install-recommends \
    libpq-dev \
    && rm -rf /var/lib/apt/lists/*

# Team A's WORKDIR. Python will find modules here because PYTHONPATH includes it.
WORKDIR /app

# Team A's Python packages. These are their own dependencies — Flyte doesn't
# install anything on top in pure BYOI mode.
RUN pip install --no-cache-dir \
    flyte \
    pandas==2.1.4 \
    pyarrow==14.0.1

# Bake the workflow code into the image.
# In pure BYOI, this is the ONLY way Flyte can find your task functions.
# Downside: every edit to tasks.py requires a new image tag + CI build.
COPY workflow_code/ /app/workflow_code/

# /app is on PYTHONPATH so `import workflow_code.tasks` resolves at runtime.
ENV PYTHONPATH=/app

# Team B's image: model training (GPU workload)
#
# Intentionally different from Team A:
#   - Different base: CUDA runtime instead of debian slim
#   - Different Python version: 3.10 (team B's standard)
#   - Different WORKDIR: /workspace (not /app)
#   - Different PYTHONPATH: /workspace
#   - Different system packages: CUDA tools
#
# Pure BYOI: Flyte injects nothing at runtime. Everything must be baked in.
# Each team owns their filesystem layout entirely.
#
# In practice this would be:
#   FROM nvidia/cuda:12.1.0-cudnn8-runtime-ubuntu22.04
# Using python:3.10-slim here so you can build/test without a GPU machine.

FROM python:3.10-slim

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# Team B uses /workspace, not /app. Each team controls their own layout.
WORKDIR /workspace

RUN pip install --no-cache-dir \
    flyte \
    torch==2.1.2 \
    numpy==1.26.4

# Bake in workflow code at /workspace/workflow_code/.
# Python finds it because PYTHONPATH includes /workspace.
COPY workflow_code/ /workspace/workflow_code/

# /workspace is on PYTHONPATH — same import path as Team A's image despite
# the different WORKDIR. Both images expose `import workflow_code.tasks`.
ENV PYTHONPATH=/workspace

docker build -f data_prep/Dockerfile -t <your-registry>/data-prep:latest .
docker build -f training/Dockerfile  -t <your-registry>/training:latest  .
docker push <your-registry>/data-prep:latest
docker push <your-registry>/training:latest

import flyte

# Image refs are resolved at runtime via init_from_config(images=...) in main.py.
# from_ref_name() is used instead of hardcoding the URI because this file is
# COPYed into both images — hardcoding would create a circular reference (the
# image baked into itself would reference its own tag).
env_train = flyte.TaskEnvironment(
    name="training",
    image=flyte.Image.from_ref_name("training"),
)

env_data = flyte.TaskEnvironment(
    name="data-prep",
    image=flyte.Image.from_ref_name("data-prep"),
    depends_on=[env_train],
)

from workflow_code.envs import env_data, env_train


@env_train.task
async def train(processed: str) -> float:
    # Runs in Team B's container: python 3.10, WORKDIR /workspace
    # torch is available (installed in training/Dockerfile)
    return float(len(processed))


@env_data.task
async def prepare(raw: str = "Hello World") -> float:
    # Runs in Team A's container: python 3.11, WORKDIR /app
    # pandas is available (installed in data_prep/Dockerfile)
    processed = raw.strip().lower()
    return await train(processed)

# /// script
# requires-python = ">=3.10"
# dependencies = ["flyte"]
# ///
"""Pure BYOI entry point.

workflow_code/ is baked into both images via their Dockerfiles.
Flyte runs each container as a black box — no code bundle is sent.

Build and push both images first (from v2_guide/pure_byoi/):
  docker build -f data_prep/Dockerfile -t <your-registry>/data-prep:latest .
  docker build -f training/Dockerfile  -t <your-registry>/training:latest  .
  docker push <your-registry>/data-prep:latest
  docker push <your-registry>/training:latest

Run (from v2_guide/pure_byoi/):
  uv run main.py
"""

from workflow_code.tasks import prepare

import flyte

DATA_PREP_IMAGE = "<your-registry>/data-prep:latest"
TRAINING_IMAGE = "<your-registry>/training:latest"

if __name__ == "__main__":
    # No root_dir — Flyte does not inject code. The images contain everything.
    flyte.init_from_config(
        project="flytesnacks",
        domain="development",
        images=(
            f"data-prep={DATA_PREP_IMAGE}",
            f"training={TRAINING_IMAGE}",
        ),
    )

    # Development: run the pipeline.
    # copy_style="none": no code bundle, the image IS the deployment.
    run = flyte.with_runcontext(copy_style="none", version="dev").run(prepare, raw="Hello World")
    print(run.url)
    run.wait()

    # Production: register task environments against the cluster.
    # from workflow_code.envs import env_data, env_train
    # flyte.deploy(env_data, copy_style="none", version="1.0.0")
    # flyte.deploy(env_train, copy_style="none", version="1.0.0")

uv run main.py

# Team A's image: data preparation
# Uses conda. Team A has no knowledge of Flyte.
# Python is at /opt/conda/bin/python.
# No PYTHONPATH set — teams don't know Flyte needs it.

FROM continuumio/miniconda3:latest

RUN conda install -y -c conda-forge pandas==2.2.3 pyarrow==19.0.1 && \
    conda clean -afy

WORKDIR /app

# Team B's image: model training
# Uses a pip venv at /opt/venv. Team B has no knowledge of Flyte.
# Python is at /opt/venv/bin/python.
# No PYTHONPATH set, PATH doesn't include /opt/venv/bin.

FROM python:3.10-slim

RUN apt-get update && apt-get install -y --no-install-recommends \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

RUN python -m venv /opt/venv && \
    /opt/venv/bin/pip install --no-cache-dir \
        torch==2.1.2+cpu \
        numpy==1.26.4 \
    --extra-index-url https://download.pytorch.org/whl/cpu

WORKDIR /workspace

import pathlib

import flyte

# Root of the remote_builder example (one level up from this file)
HERE = pathlib.Path(__file__).parent.parent

REGISTRY = "<your-registry>"  # e.g. "ghcr.io/your-org" or "123456789.dkr.ecr.us-east-1.amazonaws.com"

# ── Team B: pip venv ────────────────────────────────────────────────────────────
# Base image uses a venv at /opt/venv — python at /opt/venv/bin/python.
# The venv is not activated: PATH doesn't include /opt/venv/bin.
# Flyte adapts: installs flyte via venv's pip, adds venv to PATH, sets PYTHONPATH.
# $PATH expands at Docker build time to the base image's PATH value.
env_train_image = (
    flyte.Image.from_base("<your-registry>/training-base:latest")
    .clone(name="<your-org>/<your-image>", registry=REGISTRY, extendable=True)
    .with_commands(["/opt/venv/bin/pip install flyte"])
    .with_env_vars(
        {
            "PATH": "/opt/venv/bin:$PATH",
            "PYTHONPATH": "/workspace",  # /workspace is WORKDIR
        }
    )
    .with_code_bundle()
)

# ── Team A: conda ───────────────────────────────────────────────────────────────
# Base image uses conda — python at /opt/conda/bin/python.
# conda's own Dockerfile already adds /opt/conda/bin to PATH, so python is findable.
# Flyte adapts: installs flyte via conda's pip, sets PYTHONPATH.
env_data_image = (
    flyte.Image.from_base("<your-registry>/data-prep-base:latest")
    .clone(name="<your-org>/<your-image>", registry=REGISTRY, extendable=True)
    .with_commands(["/opt/conda/bin/pip install flyte"])
    .with_env_vars({"PYTHONPATH": "/app"})  # /app is WORKDIR; code bundle extracts here
    .with_code_bundle()
)

env_train = flyte.TaskEnvironment(name="training", image=env_train_image)
env_data = flyte.TaskEnvironment(name="data-prep", image=env_data_image, depends_on=[env_train])

from tasks.envs import env_data, env_train


@env_train.task
async def train(processed: str) -> float:
    # Runs in /opt/venv — torch and numpy available from Team B's base image
    return float(len(processed))


@env_data.task
async def prepare(raw: str = "Hello World") -> float:
    # Runs in conda env — pandas and pyarrow available from Team A's base image
    processed = raw.strip().lower()
    return await train(processed)

# /// script
# requires-python = ">=3.10"
# dependencies = ["flyte"]
# ///
"""Remote builder BYOI: two teams, two Flyte-unaware images, Flyte fills the gaps.

Each team owns their base image — built with their preferred package manager,
no knowledge of Flyte. The Flyte engineer adapts by installing flyte via the
right pip, fixing PATH, and setting PYTHONPATH.

  data_prep:  continuumio/miniconda3  — conda env, python at /opt/conda/bin/python
  training:   python:3.10-slim        — venv at /opt/venv, python at /opt/venv/bin/python

Build and push base images first (from v2_guide/remote_builder/):
  docker build -f data_prep/Dockerfile -t <your-registry>/data-prep-base:latest data_prep/
  docker build -f training/Dockerfile  -t <your-registry>/training-base:latest  training/
  docker push <your-registry>/data-prep-base:latest
  docker push <your-registry>/training-base:latest

Run (from v2_guide/remote_builder/):
  uv run main.py
"""

import pathlib

import flyte

HERE = pathlib.Path(__file__).parent

if __name__ == "__main__":
    flyte.init_from_config(root_dir=HERE)

    # Development: fast code iteration — image only rebuilds when base image changes.
    # run = flyte.with_runcontext(version="dev").run(prepare, raw="Hello World")
    # print(run.url)
    # run.wait()

    # Production: bake code into both images, pin to a version.
    from tasks.envs import env_data, env_train

    flyte.deploy(env_data, copy_style="none", version="3.0.0")
    flyte.deploy(env_train, copy_style="none", version="3.0.0")

uv run main.py