diff --git a/sdk/python/feast/infra/aws.py b/sdk/python/feast/infra/aws.py index 04b4abb48b3..1d454c6a2b8 100644 --- a/sdk/python/feast/infra/aws.py +++ b/sdk/python/feast/infra/aws.py @@ -63,47 +63,64 @@ def update_infra( entities_to_keep: Sequence[Entity], partial: bool, ): - self.online_store.update( - config=self.repo_config, - tables_to_delete=tables_to_delete, - tables_to_keep=tables_to_keep, - entities_to_keep=entities_to_keep, - entities_to_delete=entities_to_delete, - partial=partial, - ) - - if self.repo_config.feature_server and self.repo_config.feature_server.enabled: - if not enable_aws_lambda_feature_server(self.repo_config): - raise ExperimentalFeatureNotEnabled(FLAG_AWS_LAMBDA_FEATURE_SERVER_NAME) - - # Since the AWS Lambda feature server will attempt to load the registry, we - # only allow the registry to be in S3. - registry_path = ( - self.repo_config.registry - if isinstance(self.repo_config.registry, str) - else self.repo_config.registry.path + try: + self.online_store.update( + config=self.repo_config, + tables_to_delete=tables_to_delete, + tables_to_keep=tables_to_keep, + entities_to_keep=entities_to_keep, + entities_to_delete=entities_to_delete, + partial=partial, ) - registry_store_class = get_registry_store_class_from_scheme(registry_path) - if registry_store_class != S3RegistryStore: - raise IncompatibleRegistryStoreClass( - registry_store_class.__name__, S3RegistryStore.__name__ - ) - ecr_client = boto3.client("ecr") - docker_image_version = _get_docker_image_version() - repository_uri = self._create_or_get_repository_uri(ecr_client) - # Only download & upload the docker image if it doesn't already exist in ECR - if not ecr_client.batch_get_image( - repositoryName=AWS_LAMBDA_FEATURE_SERVER_REPOSITORY, - imageIds=[{"imageTag": docker_image_version}], - ).get("images"): - image_uri = self._upload_docker_image( - ecr_client, repository_uri, docker_image_version + if ( + self.repo_config.feature_server + and self.repo_config.feature_server.enabled + ): + if not enable_aws_lambda_feature_server(self.repo_config): + raise ExperimentalFeatureNotEnabled( + FLAG_AWS_LAMBDA_FEATURE_SERVER_NAME + ) + + # Since the AWS Lambda feature server will attempt to load the registry, we + # only allow the registry to be in S3. + registry_path = ( + self.repo_config.registry + if isinstance(self.repo_config.registry, str) + else self.repo_config.registry.path ) - else: - image_uri = f"{repository_uri}:{docker_image_version}" - - self._deploy_feature_server(project, image_uri) + registry_store_class = get_registry_store_class_from_scheme( + registry_path + ) + if registry_store_class != S3RegistryStore: + raise IncompatibleRegistryStoreClass( + registry_store_class.__name__, S3RegistryStore.__name__ + ) + + ecr_client = boto3.client("ecr") + docker_image_version = _get_docker_image_version() + repository_uri = self._create_or_get_repository_uri(ecr_client) + # Only download & upload the docker image if it doesn't already exist in ECR + if not ecr_client.batch_get_image( + repositoryName=AWS_LAMBDA_FEATURE_SERVER_REPOSITORY, + imageIds=[{"imageTag": docker_image_version}], + ).get("images"): + image_uri = self._upload_docker_image( + ecr_client, repository_uri, docker_image_version + ) + else: + image_uri = f"{repository_uri}:{docker_image_version}" + + self._deploy_feature_server(project, image_uri) + except Exception: + self.rollback_infra( + project=project, + tables_to_delete=tables_to_delete, + tables_to_keep=tables_to_keep, + entities_to_delete=entities_to_delete, + entities_to_keep=entities_to_keep, + ) + raise def _deploy_feature_server(self, project: str, image_uri: str): _logger.info("Deploying feature server...") diff --git a/sdk/python/feast/infra/passthrough_provider.py b/sdk/python/feast/infra/passthrough_provider.py index c6dda628172..986cc490e26 100644 --- a/sdk/python/feast/infra/passthrough_provider.py +++ b/sdk/python/feast/infra/passthrough_provider.py @@ -6,6 +6,11 @@ from tqdm import tqdm from feast.entity import Entity +from feast.errors import ( + EntityNotFoundException, + FeatureTableNotFoundException, + FeatureViewNotFoundException, +) from feast.feature_table import FeatureTable from feast.feature_view import FeatureView from feast.infra.offline_stores.offline_store import RetrievalJob @@ -48,13 +53,70 @@ def update_infra( partial: bool, ): set_usage_attribute("provider", self.__class__.__name__) + try: + self.online_store.update( + config=self.repo_config, + tables_to_delete=tables_to_delete, + tables_to_keep=tables_to_keep, + entities_to_keep=entities_to_keep, + entities_to_delete=entities_to_delete, + partial=partial, + ) + except Exception: + self.rollback_infra( + project=project, + tables_to_delete=tables_to_delete, + tables_to_keep=tables_to_keep, + entities_to_delete=entities_to_delete, + entities_to_keep=entities_to_keep, + ) + raise + + def rollback_infra( + self, + project: str, + tables_to_delete: Sequence[Union[FeatureTable, FeatureView]], + tables_to_keep: Sequence[Union[FeatureTable, FeatureView]], + entities_to_delete: Sequence[Entity], + entities_to_keep: Sequence[Entity], + ): + set_usage_attribute("provider", self.__class__.__name__) + registry = self.repo_config.get_registry_config() + + tables_to_recreate = list(tables_to_delete) + tables_to_revert: List[Union[FeatureTable, FeatureView]] = [] + tables_to_remove: List[Union[FeatureTable, FeatureView]] = [] + for table in tables_to_keep: + # Need to remove tables which didn't already exist and revert ones which + # might have been changed. + try: + tables_to_revert.append(registry.get_feature_view(table.name, project)) + except FeatureViewNotFoundException: + try: + tables_to_revert.append( + registry.get_feature_table(table.name, project) + ) + except FeatureTableNotFoundException: + tables_to_remove.append(table) + + entities_to_recreate = list(entities_to_delete) + entities_to_revert: List[Entity] = [] + entities_to_remove: List[Entity] = [] + for entity in entities_to_keep: + # Need to remove entities which didn't already exist and revert ones which + # might have been changed. + try: + entities_to_revert.append(registry.get_entity(entity.name, project)) + except EntityNotFoundException: + entities_to_remove.append(entity) + self.online_store.update( config=self.repo_config, - tables_to_delete=tables_to_delete, - tables_to_keep=tables_to_keep, - entities_to_keep=entities_to_keep, - entities_to_delete=entities_to_delete, - partial=partial, + tables_to_delete=tables_to_remove, + tables_to_keep=tables_to_revert + tables_to_recreate, + entities_to_keep=entities_to_revert + entities_to_recreate, + entities_to_delete=entities_to_remove, + partial=False, ) def teardown_infra(