| # |
| # Licensed to the Apache Software Foundation (ASF) under one or more |
| # contributor license agreements. See the NOTICE file distributed with |
| # this work for additional information regarding copyright ownership. |
| # The ASF licenses this file to You under the Apache License, Version 2.0 |
| # (the "License"); you may not use this file except in compliance with |
| # the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| |
| """A connector for sending API requests to the GCP Recommendations AI |
| API (https://cloud.google.com/recommendations). |
| """ |
| |
| from __future__ import absolute_import |
| |
| from typing import Sequence |
| from typing import Tuple |
| |
| from google.api_core.retry import Retry |
| |
| from apache_beam import pvalue |
| from apache_beam.metrics import Metrics |
| from apache_beam.options.pipeline_options import GoogleCloudOptions |
| from apache_beam.transforms import DoFn |
| from apache_beam.transforms import ParDo |
| from apache_beam.transforms import PTransform |
| from apache_beam.transforms.util import GroupIntoBatches |
| from cachetools.func import ttl_cache |
| |
| # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports |
| try: |
| from google.cloud import recommendationengine |
| except ImportError: |
| raise ImportError( |
| 'Google Cloud Recommendation AI not supported for this execution ' |
| 'environment (could not import google.cloud.recommendationengine).') |
| # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports |
| |
| __all__ = [ |
| 'CreateCatalogItem', |
| 'WriteUserEvent', |
| 'ImportCatalogItems', |
| 'ImportUserEvents', |
| 'PredictUserEvent' |
| ] |
| |
| FAILED_CATALOG_ITEMS = "failed_catalog_items" |
| |
| |
| @ttl_cache(maxsize=128, ttl=3600) |
| def get_recommendation_prediction_client(): |
| """Returns a Recommendation AI - Prediction Service client.""" |
| _client = recommendationengine.PredictionServiceClient() |
| return _client |
| |
| |
| @ttl_cache(maxsize=128, ttl=3600) |
| def get_recommendation_catalog_client(): |
| """Returns a Recommendation AI - Catalog Service client.""" |
| _client = recommendationengine.CatalogServiceClient() |
| return _client |
| |
| |
| @ttl_cache(maxsize=128, ttl=3600) |
| def get_recommendation_user_event_client(): |
| """Returns a Recommendation AI - UserEvent Service client.""" |
| _client = recommendationengine.UserEventServiceClient() |
| return _client |
| |
| |
| class CreateCatalogItem(PTransform): |
| """Creates catalogitem information. |
| The ``PTranform`` returns a PCollectionTuple with a PCollections of |
| successfully and failed created CatalogItems. |
| |
| Example usage:: |
| |
| pipeline | CreateCatalogItem( |
| project='example-gcp-project', |
| catalog_name='my-catalog') |
| """ |
| def __init__( |
| self, |
| project: str = None, |
| retry: Retry = None, |
| timeout: float = 120, |
| metadata: Sequence[Tuple[str, str]] = None, |
| catalog_name: str = "default_catalog"): |
| """Initializes a :class:`CreateCatalogItem` transform. |
| |
| Args: |
| project (str): Optional. GCP project name in which the catalog |
| data will be imported. |
| retry: Optional. Designation of what |
| errors, if any, should be retried. |
| timeout (float): Optional. The amount of time, in seconds, to wait |
| for the request to complete. |
| metadata: Optional. Strings which |
| should be sent along with the request as metadata. |
| catalog_name (str): Optional. Name of the catalog. |
| Default: 'default_catalog' |
| """ |
| self.project = project |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.catalog_name = catalog_name |
| |
| def expand(self, pcoll): |
| if self.project is None: |
| self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project |
| if self.project is None: |
| raise ValueError( |
| """GCP project name needs to be specified in "project" pipeline |
| option""") |
| return pcoll | ParDo( |
| _CreateCatalogItemFn( |
| self.project, |
| self.retry, |
| self.timeout, |
| self.metadata, |
| self.catalog_name)) |
| |
| |
| class _CreateCatalogItemFn(DoFn): |
| def __init__( |
| self, |
| project: str = None, |
| retry: Retry = None, |
| timeout: float = 120, |
| metadata: Sequence[Tuple[str, str]] = None, |
| catalog_name: str = None): |
| self._client = None |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.parent = f"projects/{project}/locations/global/catalogs/{catalog_name}" |
| self.counter = Metrics.counter(self.__class__, "api_calls") |
| |
| def setup(self): |
| if self._client is None: |
| self._client = get_recommendation_catalog_client() |
| |
| def process(self, element): |
| catalog_item = recommendationengine.CatalogItem(element) |
| request = recommendationengine.CreateCatalogItemRequest( |
| parent=self.parent, catalog_item=catalog_item) |
| |
| try: |
| created_catalog_item = self._client.create_catalog_item( |
| request=request, |
| retry=self.retry, |
| timeout=self.timeout, |
| metadata=self.metadata) |
| |
| self.counter.inc() |
| yield recommendationengine.CatalogItem.to_dict(created_catalog_item) |
| except Exception: |
| yield pvalue.TaggedOutput( |
| FAILED_CATALOG_ITEMS, |
| recommendationengine.CatalogItem.to_dict(catalog_item)) |
| |
| |
| class ImportCatalogItems(PTransform): |
| """Imports catalogitems in bulk. |
| The `PTransform` returns a PCollectionTuple with PCollections of |
| successfully and failed imported CatalogItems. |
| |
| Example usage:: |
| |
| pipeline |
| | ImportCatalogItems( |
| project='example-gcp-project', |
| catalog_name='my-catalog') |
| """ |
| def __init__( |
| self, |
| max_batch_size: int = 5000, |
| project: str = None, |
| retry: Retry = None, |
| timeout: float = 120, |
| metadata: Sequence[Tuple[str, str]] = None, |
| catalog_name: str = "default_catalog"): |
| """Initializes a :class:`ImportCatalogItems` transform |
| |
| Args: |
| batch_size (int): Required. Maximum number of catalogitems per |
| request. |
| project (str): Optional. GCP project name in which the catalog |
| data will be imported. |
| retry: Optional. Designation of what |
| errors, if any, should be retried. |
| timeout (float): Optional. The amount of time, in seconds, to wait |
| for the request to complete. |
| metadata: Optional. Strings which |
| should be sent along with the request as metadata. |
| catalog_name (str): Optional. Name of the catalog. |
| Default: 'default_catalog' |
| """ |
| self.max_batch_size = max_batch_size |
| self.project = project |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.catalog_name = catalog_name |
| |
| def expand(self, pcoll): |
| if self.project is None: |
| self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project |
| if self.project is None: |
| raise ValueError( |
| 'GCP project name needs to be specified in "project" pipeline option') |
| return ( |
| pcoll | GroupIntoBatches.WithShardedKey(self.max_batch_size) | ParDo( |
| _ImportCatalogItemsFn( |
| self.project, |
| self.retry, |
| self.timeout, |
| self.metadata, |
| self.catalog_name))) |
| |
| |
| class _ImportCatalogItemsFn(DoFn): |
| def __init__( |
| self, |
| project=None, |
| retry=None, |
| timeout=120, |
| metadata=None, |
| catalog_name=None): |
| self._client = None |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.parent = f"projects/{project}/locations/global/catalogs/{catalog_name}" |
| self.counter = Metrics.counter(self.__class__, "api_calls") |
| |
| def setup(self): |
| if self._client is None: |
| self.client = get_recommendation_catalog_client() |
| |
| def process(self, element): |
| catalog_items = [recommendationengine.CatalogItem(e) for e in element[1]] |
| catalog_inline_source = recommendationengine.CatalogInlineSource( |
| {"catalog_items": catalog_items}) |
| input_config = recommendationengine.InputConfig( |
| catalog_inline_source=catalog_inline_source) |
| |
| request = recommendationengine.ImportCatalogItemsRequest( |
| parent=self.parent, input_config=input_config) |
| |
| try: |
| operation = self._client.import_catalog_items( |
| request=request, |
| retry=self.retry, |
| timeout=self.timeout, |
| metadata=self.metadata) |
| self.counter.inc(len(catalog_items)) |
| yield operation.result() |
| except Exception: |
| yield pvalue.TaggedOutput(FAILED_CATALOG_ITEMS, catalog_items) |
| |
| |
| class WriteUserEvent(PTransform): |
| """Write user event information. |
| The `PTransform` returns a PCollectionTuple with PCollections of |
| successfully and failed written UserEvents. |
| |
| Example usage:: |
| |
| pipeline |
| | WriteUserEvent( |
| project='example-gcp-project', |
| catalog_name='my-catalog', |
| event_store='my_event_store') |
| """ |
| def __init__( |
| self, |
| project: str = None, |
| retry: Retry = None, |
| timeout: float = 120, |
| metadata: Sequence[Tuple[str, str]] = None, |
| catalog_name: str = "default_catalog", |
| event_store: str = "default_event_store"): |
| """Initializes a :class:`WriteUserEvent` transform. |
| |
| Args: |
| project (str): Optional. GCP project name in which the catalog |
| data will be imported. |
| retry: Optional. Designation of what |
| errors, if any, should be retried. |
| timeout (float): Optional. The amount of time, in seconds, to wait |
| for the request to complete. |
| metadata: Optional. Strings which |
| should be sent along with the request as metadata. |
| catalog_name (str): Optional. Name of the catalog. |
| Default: 'default_catalog' |
| event_store (str): Optional. Name of the event store. |
| Default: 'default_event_store' |
| """ |
| self.project = project |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.catalog_name = catalog_name |
| self.event_store = event_store |
| |
| def expand(self, pcoll): |
| if self.project is None: |
| self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project |
| if self.project is None: |
| raise ValueError( |
| 'GCP project name needs to be specified in "project" pipeline option') |
| return pcoll | ParDo( |
| _WriteUserEventFn( |
| self.project, |
| self.retry, |
| self.timeout, |
| self.metadata, |
| self.catalog_name, |
| self.event_store)) |
| |
| |
| class _WriteUserEventFn(DoFn): |
| FAILED_USER_EVENTS = "failed_user_events" |
| |
| def __init__( |
| self, |
| project=None, |
| retry=None, |
| timeout=120, |
| metadata=None, |
| catalog_name=None, |
| event_store=None): |
| self._client = None |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.parent = f"projects/{project}/locations/global/catalogs/"\ |
| f"{catalog_name}/eventStores/{event_store}" |
| self.counter = Metrics.counter(self.__class__, "api_calls") |
| |
| def setup(self): |
| if self._client is None: |
| self._client = get_recommendation_user_event_client() |
| |
| def process(self, element): |
| user_event = recommendationengine.UserEvent(element) |
| request = recommendationengine.WriteUserEventRequest( |
| parent=self.parent, user_event=user_event) |
| |
| try: |
| created_user_event = self._client.write_user_event(request) |
| self.counter.inc() |
| yield recommendationengine.UserEvent.to_dict(created_user_event) |
| except Exception: |
| yield pvalue.TaggedOutput( |
| self.FAILED_USER_EVENTS, |
| recommendationengine.UserEvent.to_dict(user_event)) |
| |
| |
| class ImportUserEvents(PTransform): |
| """Imports userevents in bulk. |
| The `PTransform` returns a PCollectionTuple with PCollections of |
| successfully and failed imported UserEvents. |
| |
| Example usage:: |
| |
| pipeline |
| | ImportUserEvents( |
| project='example-gcp-project', |
| catalog_name='my-catalog', |
| event_store='my_event_store') |
| """ |
| def __init__( |
| self, |
| max_batch_size: int = 5000, |
| project: str = None, |
| retry: Retry = None, |
| timeout: float = 120, |
| metadata: Sequence[Tuple[str, str]] = None, |
| catalog_name: str = "default_catalog", |
| event_store: str = "default_event_store"): |
| """Initializes a :class:`WriteUserEvent` transform. |
| |
| Args: |
| batch_size (int): Required. Maximum number of catalogitems |
| per request. |
| project (str): Optional. GCP project name in which the catalog |
| data will be imported. |
| retry: Optional. Designation of what |
| errors, if any, should be retried. |
| timeout (float): Optional. The amount of time, in seconds, to wait |
| for the request to complete. |
| metadata: Optional. Strings which |
| should be sent along with the request as metadata. |
| catalog_name (str): Optional. Name of the catalog. |
| Default: 'default_catalog' |
| event_store (str): Optional. Name of the event store. |
| Default: 'default_event_store' |
| """ |
| self.max_batch_size = max_batch_size |
| self.project = project |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.catalog_name = catalog_name |
| self.event_store = event_store |
| |
| def expand(self, pcoll): |
| if self.project is None: |
| self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project |
| if self.project is None: |
| raise ValueError( |
| 'GCP project name needs to be specified in "project" pipeline option') |
| return ( |
| pcoll | GroupIntoBatches.WithShardedKey(self.max_batch_size) | ParDo( |
| _ImportUserEventsFn( |
| self.project, |
| self.retry, |
| self.timeout, |
| self.metadata, |
| self.catalog_name, |
| self.event_store))) |
| |
| |
| class _ImportUserEventsFn(DoFn): |
| FAILED_USER_EVENTS = "failed_user_events" |
| |
| def __init__( |
| self, |
| project=None, |
| retry=None, |
| timeout=120, |
| metadata=None, |
| catalog_name=None, |
| event_store=None): |
| self._client = None |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.parent = f"projects/{project}/locations/global/catalogs/"\ |
| f"{catalog_name}/eventStores/{event_store}" |
| self.counter = Metrics.counter(self.__class__, "api_calls") |
| |
| def setup(self): |
| if self._client is None: |
| self.client = get_recommendation_user_event_client() |
| |
| def process(self, element): |
| |
| user_events = [recommendationengine.UserEvent(e) for e in element[1]] |
| user_event_inline_source = recommendationengine.UserEventInlineSource( |
| {"user_events": user_events}) |
| input_config = recommendationengine.InputConfig( |
| user_event_inline_source=user_event_inline_source) |
| |
| request = recommendationengine.ImportUserEventsRequest( |
| parent=self.parent, input_config=input_config) |
| |
| try: |
| operation = self._client.write_user_event(request) |
| self.counter.inc(len(user_events)) |
| yield recommendationengine.PredictResponse.to_dict(operation.result()) |
| except Exception: |
| yield pvalue.TaggedOutput(self.FAILED_USER_EVENTS, user_events) |
| |
| |
| class PredictUserEvent(PTransform): |
| """Make a recommendation prediction. |
| The `PTransform` returns a PCollection |
| |
| Example usage:: |
| |
| pipeline |
| | PredictUserEvent( |
| project='example-gcp-project', |
| catalog_name='my-catalog', |
| event_store='my_event_store', |
| placement_id='recently_viewed_default') |
| """ |
| def __init__( |
| self, |
| project: str = None, |
| retry: Retry = None, |
| timeout: float = 120, |
| metadata: Sequence[Tuple[str, str]] = None, |
| catalog_name: str = "default_catalog", |
| event_store: str = "default_event_store", |
| placement_id: str = None): |
| """Initializes a :class:`PredictUserEvent` transform. |
| |
| Args: |
| project (str): Optional. GCP project name in which the catalog |
| data will be imported. |
| retry: Optional. Designation of what |
| errors, if any, should be retried. |
| timeout (float): Optional. The amount of time, in seconds, to wait |
| for the request to complete. |
| metadata: Optional. Strings which |
| should be sent along with the request as metadata. |
| catalog_name (str): Optional. Name of the catalog. |
| Default: 'default_catalog' |
| event_store (str): Optional. Name of the event store. |
| Default: 'default_event_store' |
| placement_id (str): Required. ID of the recommendation engine |
| placement. This id is used to identify the set of models that |
| will be used to make the prediction. |
| """ |
| self.project = project |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.placement_id = placement_id |
| self.catalog_name = catalog_name |
| self.event_store = event_store |
| if placement_id is None: |
| raise ValueError('placement_id must be specified') |
| else: |
| self.placement_id = placement_id |
| |
| def expand(self, pcoll): |
| if self.project is None: |
| self.project = pcoll.pipeline.options.view_as(GoogleCloudOptions).project |
| if self.project is None: |
| raise ValueError( |
| 'GCP project name needs to be specified in "project" pipeline option') |
| return pcoll | ParDo( |
| _PredictUserEventFn( |
| self.project, |
| self.retry, |
| self.timeout, |
| self.metadata, |
| self.catalog_name, |
| self.event_store, |
| self.placement_id)) |
| |
| |
| class _PredictUserEventFn(DoFn): |
| FAILED_PREDICTIONS = "failed_predictions" |
| |
| def __init__( |
| self, |
| project=None, |
| retry=None, |
| timeout=120, |
| metadata=None, |
| catalog_name=None, |
| event_store=None, |
| placement_id=None): |
| self._client = None |
| self.retry = retry |
| self.timeout = timeout |
| self.metadata = metadata |
| self.name = f"projects/{project}/locations/global/catalogs/"\ |
| f"{catalog_name}/eventStores/{event_store}/placements/"\ |
| f"{placement_id}" |
| self.counter = Metrics.counter(self.__class__, "api_calls") |
| |
| def setup(self): |
| if self._client is None: |
| self._client = get_recommendation_prediction_client() |
| |
| def process(self, element): |
| user_event = recommendationengine.UserEvent(element) |
| request = recommendationengine.PredictRequest( |
| name=self.name, user_event=user_event) |
| |
| try: |
| prediction = self._client.predict(request) |
| self.counter.inc() |
| yield [ |
| recommendationengine.PredictResponse.to_dict(p) |
| for p in prediction.pages |
| ] |
| except Exception: |
| yield pvalue.TaggedOutput(self.FAILED_PREDICTIONS, user_event) |
