feat(components): Add CMEK validation to preview.llm.infer_pipeline

PiperOrigin-RevId: 612956960

components/google-cloud/RELEASE.md

@@ -2,6 +2,7 @@
 * Add `v1.automl.forecasting.learn_to_learn_forecasting_pipeline`, `v1.automl.forecasting.sequence_to_sequence_forecasting_pipeline`, `v1.automl.forecasting.temporal_fusion_transformer_forecasting_pipeline`, `v1.automl.forecasting.time_series_dense_encoder_forecasting_pipeline` as Forecasting on Pipelines moves to GA.
 * Fix bug in `preview.llm.rlhf_pipeline` that caused wrong output artifact to be used for inference after training.
 * Fix issue where AutoSxS was not propagating location to all sub-components.
+* Add CMEK support to `preview.llm.infer_pipeline`.
 
 ## Release 2.10.0
 * Fix the missing output of pipeline remote runner. `AutoMLImageTrainingJobRunOp` now passes the model artifacts correctly to downstream components.
@@ -11,7 +12,6 @@
 * Bump supported KFP versions to `kfp>=2.6.0,<=2.7.0`.
 * Apply latest GCPC image vulnerability resolutions (base OS and software updates).
 * Add CMEK support to `preview.llm.rlhf_pipeline` when tuning in `us-central1` with GPUs.
-
 ## Release 2.9.0
 * Use `large_model_reference` for `model_reference_name` when uploading models from `preview.llm.rlhf_pipeline` instead of hardcoding value as `text-bison@001`.
 * Disable caching when resolving model display names for RLHF-tuned models so a unique name is generated on each `preview.llm.rlhf_pipeline` run.

components/google-cloud/google_cloud_pipeline_components/_implementation/llm/bulk_inferrer.py

@@ -37,6 +37,7 @@ def bulk_inferrer(
     output_prediction_gcs_path: kfp.dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     gcp_resources: kfp.dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     sampling_strategy: str = 'greedy',
+    encryption_spec_key_name: str = '',
 ) -> kfp.dsl.ContainerSpec:  # pylint: disable=g-doc-args
   """Performs bulk inference.
 
@@ -56,6 +57,10 @@ def bulk_inferrer(
     input_dataset_path: Path to dataset to use for inference.
     sampling_strategy: The sampling strategy for inference.
     dataset_split: Perform inference on this split of the input dataset.
+    encryption_spec_key_name: Customer-managed encryption key. If this is set,
+      then all resources created by the CustomJob will be encrypted with the
+      provided encryption key. Note that this is not supported for TPU at the
+      moment.
 
   Returns:
     output_prediction: Where to save the output prediction.
@@ -83,6 +88,7 @@ def bulk_inferrer(
               f'--output_prediction={output_prediction}',
               f'--output_prediction_gcs_path={output_prediction_gcs_path}',
           ],
+          encryption_spec_key_name=encryption_spec_key_name,
       ),
       gcp_resources=gcp_resources,
   )

components/google-cloud/google_cloud_pipeline_components/preview/llm/infer/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2023 The Kubeflow Authors. All Rights Reserved.
+#
+# Licensed 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.

components/google-cloud/google_cloud_pipeline_components/preview/llm/infer/component.py

@@ -42,6 +42,7 @@ def infer_pipeline(
     instruction: Optional[str] = None,
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
     location: str = _placeholders.LOCATION_PLACEHOLDER,
+    encryption_spec_key_name: str = '',
 ) -> PipelineOutput:
   # fmt: off
   """Uses a large-language model to perform bulk inference on a prompt dataset.
@@ -56,6 +57,7 @@ def infer_pipeline(
     instruction: This field lets the model know what task it needs to perform. Base models have been trained over a large set of varied instructions. You can give a simple and intuitive description of the task and the model will follow it, e.g. "Classify this movie review as positive or negative" or "Translate this sentence to Danish". Do not specify this if your dataset already prepends the instruction to the inputs field.
     project: Project used to run custom jobs. If not specified the project used to run the pipeline will be used.
     location: Location used to run custom jobs. If not specified the location used to run the pipeline will be used.
+    encryption_spec_key_name: Customer-managed encryption key. If this is set, then all resources created by the CustomJob will be encrypted with the provided encryption key. Note that this is not supported for TPU at the moment.
 
   Returns:
     Cloud storage path to output predictions.
@@ -94,6 +96,7 @@ def infer_pipeline(
               'large_model_reference'
           ],
           instruction=resolved_text_instruction.output,
+          encryption_spec_key_name=encryption_spec_key_name,
       )
       .set_display_name('Import Prompt Dataset')
       .set_caching_options(False)
@@ -118,6 +121,7 @@ def infer_pipeline(
       accelerator_count=machine_spec.outputs['accelerator_count'],
       machine_type=machine_spec.outputs['machine_type'],
       image_uri=bulk_inferrer_image_uri.output,
+      encryption_spec_key_name=encryption_spec_key_name,
   ).set_display_name('Bulk Inferrer')
 
   return PipelineOutput(

components/google-cloud/google_cloud_pipeline_components/preview/llm/rlhf/component.py

@@ -167,6 +167,7 @@ def rlhf_pipeline(
           prompt_sequence_length=prompt_sequence_length,
           target_sequence_length=target_sequence_length,
           instruction=instruction,
+          encryption_spec_key_name=encryption_spec_key_name,
       )
 
   llm_model_handler = deployment_graph.pipeline(