Skip to content

Commit 294ad10

Browse files
felixwang9817felixwang9817
committed
Upload docker image to ECR during feast apply
Signed-off-by: Felix Wang <wangfelix98@gmail.com>
1 parent 2fccf76 commit 294ad10

3 files changed

Lines changed: 48 additions & 0 deletions

File tree

‎sdk/python/feast/feature_store.py‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
update_data_sources_with_inferred_event_timestamp_col,
4646
update_entities_with_inferred_types_from_feature_views,
4747
)
48+
from feast.infra.feature_servers.feature_server import FEATURE_SERVER_IMAGE_FOR_TYPE
4849
from feast.infra.provider import Provider, RetrievalJob, get_provider
4950
from feast.on_demand_feature_view import OnDemandFeatureView
5051
from feast.online_response import OnlineResponse, _infer_online_entity_rows
@@ -1025,6 +1026,48 @@ def serve(self, port: int) -> None:
10251026

10261027
feature_server.start_server(self, port)
10271028

1029+
@log_exceptions_and_usage
1030+
def upload_docker_image(self) -> None:
1031+
"""Upload the docker image for the feature consumption server to the cloud."""
1032+
1033+
# TODO: add error checking and avoid hardcoding the region
1034+
repository_name = "feast-python-server-test"
1035+
feature_server_type = (
1036+
self.config.feature_server.type if self.config.feature_server else None
1037+
)
1038+
if feature_server_type == "aws_lambda":
1039+
import base64
1040+
1041+
import boto3
1042+
import docker
1043+
from botocore.exceptions import ClientError
1044+
1045+
docker_client = docker.from_env()
1046+
image_name = FEATURE_SERVER_IMAGE_FOR_TYPE[feature_server_type]
1047+
docker_client.images.pull(image_name)
1048+
1049+
ecr_client = boto3.client("ecr", region_name="us-west-2")
1050+
try:
1051+
ecr_client.create_repository(repositoryName=repository_name)
1052+
except ClientError:
1053+
pass
1054+
auth_token = ecr_client.get_authorization_token()["authorizationData"][0][
1055+
"authorizationToken"
1056+
]
1057+
username, password = base64.b64decode(auth_token).decode("utf-8").split(":")
1058+
1059+
sts_client = boto3.client("sts")
1060+
aws_account = sts_client.get_caller_identity()["Account"]
1061+
ecr_address = f"{aws_account}.dkr.ecr.us-west-2.amazonaws.com"
1062+
docker_client.login(
1063+
username=username, password=password, registry=ecr_address
1064+
)
1065+
1066+
# Pushing will likely take several minutes.
1067+
image = docker_client.images.get(image_name)
1068+
image.tag(f"{ecr_address}/{repository_name}:latest")
1069+
docker_client.api.push(f"{ecr_address}/{repository_name}:latest")
1070+
10281071

10291072
def _entity_row_to_key(row: GetOnlineFeaturesRequestV2.EntityRow) -> EntityKeyProto:
10301073
names, values = zip(*row.fields.items())

‎sdk/python/feast/infra/feature_servers/feature_server.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
FEATURE_SERVER_IMAGE_FOR_TYPE = {"aws_lambda": "feastdockerbot/feast-python-server"}

‎sdk/python/feast/repo_operations.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,10 @@ def apply_total(repo_config: RepoConfig, repo_path: Path, skip_source_validation
257257
# Commit the update to the registry only after successful infra update
258258
registry.commit()
259259

260+
repo_config = store.config
261+
if repo_config.feature_server and repo_config.feature_server.enabled:
262+
store.upload_docker_image()
263+
260264

261265
def _tag_registry_entities_for_keep_delete(
262266
project: str, registry: Registry, repo: ParsedRepo

0 commit comments

Comments
 (0)