feat(components): migrate function_based resolve_regional_endpoint to rlhf_preprocessor component

PiperOrigin-RevId: 629315370

Author: Googler

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

@@ -39,6 +39,7 @@ def pipeline(
     deploy_model: bool = True,
     encryption_spec_key_name: str = '',
     upload_location: str = _placeholders.LOCATION_PLACEHOLDER,
+    regional_endpoint: str = '',
 ) -> PipelineOutput:
   # fmt: off
   """Uploads a tuned language model and (optionally) deploys it to an endpoint.
@@ -51,16 +52,13 @@ def pipeline(
     deploy_model: Whether to deploy the model to an endpoint in `us-central1`. Default is True.
     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.
     upload_location: Region to upload and deploy the model to. Default is the location used to run the pipeline components.
+    regional_endpoint: Regional endpoint to upload the model.
 
   Returns:
     model_resource_name: Path to the model uploaded to the Model Registry. This will be an empty string if the model was not deployed.
     endpoint_resource_name: Path the Online Prediction Endpoint. This will be an empty string if the model was not deployed.
   """
   # fmt: on
-  regional_endpoint = function_based.resolve_regional_endpoint(
-      upload_location=upload_location
-  ).set_display_name('Resolve Regional Endpoint')
-
   display_name = (
       function_based.resolve_model_display_name(
           large_model_reference=large_model_reference,
@@ -76,7 +74,7 @@ def pipeline(
   upload_task = upload_llm_model.refined_upload_llm_model(
       project=_placeholders.PROJECT_ID_PLACEHOLDER,
       location=upload_location,
-      regional_endpoint=regional_endpoint.output,
+      regional_endpoint=regional_endpoint,
       artifact_uri=output_adapter_path,
       model_display_name=display_name.output,
       model_reference_name=large_model_reference,
@@ -93,7 +91,7 @@ def pipeline(
       location=upload_location,
       model_resource_name=upload_task.outputs['model_resource_name'],
       display_name=display_name.output,
-      regional_endpoint=regional_endpoint.output,
+      regional_endpoint=regional_endpoint,
       deploy_model=deploy_model.output,
       encryption_spec_key_name=encryption_spec_key_name,
   ).set_display_name('Deploy Model')

components/google-cloud/google_cloud_pipeline_components/_implementation/llm/generated/refined_image_versions.py

@@ -17,4 +17,4 @@
 DO NOT EDIT - This file is generated, manual changes will be overridden.
 """
 
-IMAGE_TAG = '20240428_1707'
+IMAGE_TAG = '20240429_1553_RC00'

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

@@ -45,11 +45,13 @@ def rlhf_preprocessor(
     metadata_accelerator_count: dsl.OutputPath(int),  # pytype: disable=invalid-annotation
     metadata_refined_image_uri: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     metadata_num_microbatches: dsl.OutputPath(int),  # pytype: disable=invalid-annotation
+    metadata_upload_location: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     use_experimental_image: bool = False,
     evaluation_dataset: str = '',
     tensorboard_resource_id: str = '',
     input_reference_model_path: str = '',
     image_uri: str = utils.get_default_image_uri('refined_cpu', ''),
+    upload_location: str = '',
 ) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
   # fmt: off
   """Preprocess RLHF pipeline inputs.
@@ -70,6 +72,7 @@ def rlhf_preprocessor(
     metadata_reward_model_reference:  The base model for training reward model. The name should be in capitalized snake case format.
     metadata_reward_model_path: The model checkpoint path for the reward model.
     image_uri: Docker image URI to use for the custom job.
+    upload_location: Region where the model will be uploaded.
 
   Returns:
     gcp_resources: GCP resources that can be used to track the custom job.
@@ -82,6 +85,7 @@ def rlhf_preprocessor(
     metadata_refined_image_uri: Docker image URI to use for the custom job.
     metadata_num_microbatches: Number of microbatches to break the total batch
       size into during training.
+    metadata_upload_location: Regional endpoint.
   """
   # fmt: on
   return gcpc_utils.build_serverless_customjob_container_spec(
@@ -104,6 +108,7 @@ def rlhf_preprocessor(
               f'--artifact_registry={artifact_registry}',
               f'--tag={tag}',
               f'--use_experimental_image={use_experimental_image}',
+              f'--upload_location={upload_location}',
               f'--has_tensorboard_id_path={has_tensorboard_id}',
               f'--has_inference_dataset_path={has_inference_dataset}',
               f'--metadata_candidate_columns_string_path={metadata_candidate_columns_string}',
@@ -117,6 +122,7 @@ def rlhf_preprocessor(
               f'--metadata_accelerator_count_path={metadata_accelerator_count}',
               f'--metadata_refined_image_uri_path={metadata_refined_image_uri}',
               f'--metadata_num_microbatches_path={metadata_num_microbatches}',
+              f'--metadata_upload_location_path={metadata_upload_location}',
           ],
       ),
       gcp_resources=gcp_resources,

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

@@ -106,6 +106,7 @@ def rlhf_pipeline(
       tag=env.get_private_image_tag(),
       evaluation_dataset=eval_dataset,
       tensorboard_resource_id=tensorboard_resource_id,
+      upload_location=location,
   ).set_display_name('Preprocess Inputs')
   num_microbatches = preprocess_metadata.outputs['metadata_num_microbatches']
 
@@ -233,6 +234,7 @@ def rlhf_pipeline(
       deploy_model=deploy_model,
       encryption_spec_key_name=encryption_spec_key_name,
       upload_location=location,
+      regional_endpoint=preprocess_metadata.outputs['metadata_upload_location'],
   ).set_display_name('Upload and Deploy Tuned Model')
 
   return PipelineOutput(