feat: add support for Predict Request Response Logging in Endpoint SDK

vertex-sdk-bot · copybara-github · commit 372ab8dd59b3 · 2022-12-21T18:27:30.000-08:00
PiperOrigin-RevId: 497049904
diff --git a/google/cloud/aiplatform/models.py b/google/cloud/aiplatform/models.py
@@ -279,6 +279,9 @@ def create(
         sync=True,
         create_request_timeout: Optional[float] = None,
         endpoint_id: Optional[str] = None,
+        enable_request_response_logging=False,
+        request_response_logging_sampling_rate: Optional[float] = None,
+        request_response_logging_bq_destination_table: Optional[str] = None,
     ) -> "Endpoint":
         """Creates a new endpoint.
 
@@ -339,12 +342,18 @@ def create(
                 is populated based on a query string argument, such as
                 ``?endpoint_id=12345``. This is the fallback for fields
                 that are not included in either the URI or the body.
+            enable_request_response_logging (bool):
+                Optional. Whether to enable request & response logging for this endpoint.
+            request_response_logging_sampling_rate (float):
+                Optional. The request response logging sampling rate. If not set, default is 0.0.
+            request_response_logging_bq_destination_table (str):
+                Optional. The request response logging bigquery destination. If not set, will create a table with name:
+                ``bq://{project_id}.logging_{endpoint_display_name}_{endpoint_id}.request_response_logging``.
 
         Returns:
             endpoint (aiplatform.Endpoint):
                 Created endpoint.
         """
-
         api_client = cls._instantiate_client(location=location, credentials=credentials)
 
         if not display_name:
@@ -357,6 +366,17 @@ def create(
         project = project or initializer.global_config.project
         location = location or initializer.global_config.location
 
+        predict_request_response_logging_config = None
+        if enable_request_response_logging:
+            predict_request_response_logging_config = (
+                gca_endpoint_compat.PredictRequestResponseLoggingConfig(
+                    enabled=True,
+                    sampling_rate=request_response_logging_sampling_rate,
+                    bigquery_destination=gca_io_compat.BigQueryDestination(
+                        output_uri=request_response_logging_bq_destination_table
+                    ),
+                )
+            )
         return cls._create(
             api_client=api_client,
             display_name=display_name,
@@ -372,6 +392,7 @@ def create(
             sync=sync,
             create_request_timeout=create_request_timeout,
             endpoint_id=endpoint_id,
+            predict_request_response_logging_config=predict_request_response_logging_config,
         )
 
     @classmethod
@@ -391,6 +412,9 @@ def _create(
         sync=True,
         create_request_timeout: Optional[float] = None,
         endpoint_id: Optional[str] = None,
+        predict_request_response_logging_config: Optional[
+            gca_endpoint_compat.PredictRequestResponseLoggingConfig
+        ] = None,
     ) -> "Endpoint":
         """Creates a new endpoint by calling the API client.
 
@@ -453,6 +477,8 @@ def _create(
                 is populated based on a query string argument, such as
                 ``?endpoint_id=12345``. This is the fallback for fields
                 that are not included in either the URI or the body.
+            predict_request_response_logging_config (aiplatform.endpoint.PredictRequestResponseLoggingConfig):
+                Optional. The request response logging configuration for online prediction.
 
         Returns:
             endpoint (aiplatform.Endpoint):
@@ -469,6 +495,7 @@ def _create(
             labels=labels,
             encryption_spec=encryption_spec,
             network=network,
+            predict_request_response_logging_config=predict_request_response_logging_config,
         )
 
         operation_future = api_client.create_endpoint(
diff --git a/tests/unit/aiplatform/test_endpoints.py b/tests/unit/aiplatform/test_endpoints.py
@@ -49,6 +49,7 @@
     prediction_service as gca_prediction_service,
     endpoint_service as gca_endpoint_service,
     encryption_spec as gca_encryption_spec,
+    io as gca_io,
 )
 
 
@@ -200,6 +201,19 @@
 
 _TEST_LABELS = {"my_key": "my_value"}
 
+_TEST_REQUEST_RESPONSE_LOGGING_SAMPLING_RATE = 0.1
+_TEST_REQUEST_RESPONSE_LOGGING_BQ_DEST = (
+    output_uri
+) = f"bq://{_TEST_PROJECT}/test_dataset/test_table"
+_TEST_REQUEST_RESPONSE_LOGGING_CONFIG = (
+    gca_endpoint.PredictRequestResponseLoggingConfig(
+        enabled=True,
+        sampling_rate=_TEST_REQUEST_RESPONSE_LOGGING_SAMPLING_RATE,
+        bigquery_destination=gca_io.BigQueryDestination(
+            output_uri=_TEST_REQUEST_RESPONSE_LOGGING_BQ_DEST
+        ),
+    )
+)
 
 """
 ----------------------------------------------------------------------------
@@ -853,6 +867,32 @@ def test_create_with_labels(self, create_endpoint_mock, sync):
             timeout=None,
         )
 
+    @pytest.mark.usefixtures("get_endpoint_mock")
+    @pytest.mark.parametrize("sync", [True, False])
+    def test_create_with_request_response_logging(self, create_endpoint_mock, sync):
+        my_endpoint = models.Endpoint.create(
+            display_name=_TEST_DISPLAY_NAME,
+            enable_request_response_logging=True,
+            request_response_logging_sampling_rate=_TEST_REQUEST_RESPONSE_LOGGING_SAMPLING_RATE,
+            request_response_logging_bq_destination_table=_TEST_REQUEST_RESPONSE_LOGGING_BQ_DEST,
+            sync=sync,
+            create_request_timeout=None,
+        )
+        if not sync:
+            my_endpoint.wait()
+
+        expected_endpoint = gca_endpoint.Endpoint(
+            display_name=_TEST_DISPLAY_NAME,
+            predict_request_response_logging_config=_TEST_REQUEST_RESPONSE_LOGGING_CONFIG,
+        )
+        create_endpoint_mock.assert_called_once_with(
+            parent=_TEST_PARENT,
+            endpoint=expected_endpoint,
+            endpoint_id=None,
+            metadata=(),
+            timeout=None,
+        )
+
     @pytest.mark.usefixtures("get_endpoint_mock")
     def test_update_endpoint(self, update_endpoint_mock):
         endpoint = models.Endpoint(_TEST_ENDPOINT_NAME)
