consolidate training and serving image builds

Makefile

@@ -91,22 +91,11 @@ destroy-infra: ## DESTROY the Terraform infrastructure in your project. Requires
 	terraform init -backend-config='bucket=${VERTEX_PROJECT_ID}-tfstate' && \
 	terraform destroy -var 'project_id=${VERTEX_PROJECT_ID}' -var 'region=${VERTEX_LOCATION}'
 
-build-training-container: ## Build and push training container image using Docker
+target ?= training
+build-container: ## Build and push training/serving container image using Docker. Specify target=<training|serving>
 	@ cd model && \
-	poetry export -f requirements.txt -o training/requirements.txt && \
-	cd training && \
 	gcloud builds submit . \
-	--tag=${TRAINING_CONTAINER_IMAGE} \
 	--region=${VERTEX_LOCATION} \
 	--project=${VERTEX_PROJECT_ID} \
-	--gcs-source-staging-dir=gs://${VERTEX_PROJECT_ID}-staging/source
-
-build-serving-container: ## Build and push serving container image using Docker
-	@ cd model && \
-	poetry export --with serving -f requirements.txt -o serving/requirements.txt && \
-	cd serving && \
-	gcloud builds submit . \
-	--tag=${SERVING_CONTAINER_IMAGE} \
-	--region=${VERTEX_LOCATION} \
-	--project=${VERTEX_PROJECT_ID} \
-	--gcs-source-staging-dir=gs://${VERTEX_PROJECT_ID}-staging/source
+	--gcs-source-staging-dir=gs://${VERTEX_PROJECT_ID}-staging/source \
+	--substitutions=_DOCKER_TARGET=${target},_DESTINATION_IMAGE_URI=${CONTAINER_IMAGE_REGISTRY}/${target}:${RESOURCE_SUFFIX}

env.sh.example

@@ -24,5 +24,4 @@ export RESOURCE_SUFFIX=default
 # Leave as-is
 export VERTEX_SA_EMAIL=vertex-pipelines@${VERTEX_PROJECT_ID}.iam.gserviceaccount.com
 export VERTEX_PIPELINE_ROOT=gs://${VERTEX_PROJECT_ID}-pl-root
-export TRAINING_CONTAINER_IMAGE=${VERTEX_LOCATION}-docker.pkg.dev/${VERTEX_PROJECT_ID}/vertex-images/training:${RESOURCE_SUFFIX}
-export SERVING_CONTAINER_IMAGE=${VERTEX_LOCATION}-docker.pkg.dev/${VERTEX_PROJECT_ID}/vertex-images/serving:${RESOURCE_SUFFIX}
+export CONTAINER_IMAGE_REGISTRY=${VERTEX_LOCATION}-docker.pkg.dev/${VERTEX_PROJECT_ID}/vertex-images

model/Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.9.16-slim AS builder
+
+ENV PIP_NO_CACHE_DIR=off \
+    PIP_DISABLE_PIP_VERSION_CHECK=on \
+    PIP_DEFAULT_TIMEOUT=100
+
+ARG POETRY_VERSION=1.5.1
+
+COPY pyproject.toml pyproject.toml
+COPY poetry.lock poetry.lock
+
+RUN pip install poetry==${POETRY_VERSION}
+RUN poetry install
+
+FROM builder AS training
+
+COPY training/train.py training/train.py
+
+FROM builder AS serving
+
+RUN poetry install --with serving
+COPY serving/main.py serving/main.py
+
+CMD exec uvicorn serving.main:app --host "0.0.0.0" --port "$AIP_HTTP_PORT"

model/cloudbuild.yaml

@@ -0,0 +1,7 @@
+---
+steps:
+  - name: 'gcr.io/kaniko-project/executor:latest'
+    args:
+      - --destination=${_DESTINATION_IMAGE_URI}
+      - --target=${_DOCKER_TARGET}
+      - --cache=true

pipelines/src/pipelines/training/pipeline.py

@@ -22,8 +22,10 @@
 from bigquery_components import extract_bq_to_dataset
 from vertex_components import upload_model
 
-TRAINING_IMAGE = os.environ["TRAINING_CONTAINER_IMAGE"]
-SERVING_IMAGE = os.environ["SERVING_CONTAINER_IMAGE"]
+CONTAINER_IMAGE_REGISTRY = os.environ["CONTAINER_IMAGE_REGISTRY"]
+RESOURCE_SUFFIX = os.environ.get("RESOURCE_SUFFIX", "default")
+TRAINING_IMAGE = f"{CONTAINER_IMAGE_REGISTRY}/training:{RESOURCE_SUFFIX}"
+SERVING_IMAGE = f"{CONTAINER_IMAGE_REGISTRY}/serving:{RESOURCE_SUFFIX}"
 
 
 @dsl.container_component
@@ -40,7 +42,7 @@ def train(
         image=TRAINING_IMAGE,
         command=["python"],
         args=[
-            "train.py",
+            "training/train.py",
             "--train-data",
             train_data.path,
             "--valid-data",