google-cloud-pipeline-components 2.13.1__py3-none-any.whl → 2.14.0__py3-none-any.whl

google_cloud_pipeline_components/__init__.py

@@ -17,14 +17,13 @@ import warnings
 
 from google_cloud_pipeline_components.version import __version__
 
-if sys.version_info < (3, 8):
+if sys.version_info < (3, 9):
   warnings.warn(
       (
-          'Python 3.7 has reached end-of-life. Google Cloud Pipeline Components'
-          ' will drop support for Python 3.7 on April 23, 2024. To use new'
-          ' versions of the KFP SDK after that date, you will need to upgrade'
-          ' to Python >= 3.8. See https://devguide.python.org/versions/ for'
-          ' more details.'
+          ' Google Cloud Pipeline Components will drop support for Python 3.8'
+          ' on Oct 1, 2024. To use new versions of the GCPC SDK after that'
+          ' date, you will need to upgrade to Python >= 3.9. See'
+          ' https://devguide.python.org/versions/ for more details.'
       ),
       FutureWarning,
       stacklevel=2,

google_cloud_pipeline_components/_implementation/llm/deployment_graph.py

@@ -34,6 +34,7 @@ PipelineOutput = NamedTuple(
 def pipeline(
     output_adapter_path: str,
     large_model_reference: str,
+    policy_model_reference: str,
     model_display_name: Optional[str] = None,
     deploy_model: bool = True,
     encryption_spec_key_name: str = '',
@@ -45,6 +46,7 @@ def pipeline(
   Args:
     output_adapter_path: Path to the trained model adapter if LoRA tuning was used.
     large_model_reference: Name of the base model. Supported values are `text-bison@001`, `t5-small`, `t5-large`, `t5-xl` and `t5-xxl`. `text-bison@001` and `t5-small` are supported in `us-central1` and `europe-west4`. `t5-large`, `t5-xl` and `t5-xxl` are only supported in `europe-west4`.
+    policy_model_reference: The name of the model for deployment. The name should be in capitalized snake case format.
     model_display_name: Name of the fine-tuned model shown in the Model Registry. If not provided, a default name will be created.
     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.
@@ -68,14 +70,8 @@ def pipeline(
       .set_display_name('Resolve Model Display Name')
   )
 
-  reference_model_metadata = function_based.resolve_reference_model_metadata(
-      large_model_reference=large_model_reference,
-  ).set_display_name('Resolve Model Metadata')
-
   upload_model = function_based.resolve_upload_model(
-      large_model_reference=reference_model_metadata.outputs[
-          'large_model_reference'
-      ]
+      large_model_reference=policy_model_reference,
   ).set_display_name('Resolve Upload Model')
   upload_task = upload_llm_model.refined_upload_llm_model(
       project=_placeholders.PROJECT_ID_PLACEHOLDER,
@@ -90,9 +86,7 @@ def pipeline(
   ).set_display_name('Upload Model')
   deploy_model = function_based.resolve_deploy_model(
       deploy_model=deploy_model,
-      large_model_reference=reference_model_metadata.outputs[
-          'large_model_reference'
-      ],
+      large_model_reference=policy_model_reference,
   ).set_display_name('Resolve Deploy Model')
   deploy_task = deploy_llm_model.deploy_llm_model(
       project=_placeholders.PROJECT_ID_PLACEHOLDER,

google_cloud_pipeline_components/_implementation/llm/env.py

@@ -19,7 +19,7 @@ from google_cloud_pipeline_components._implementation.llm.generated import refin
 
 
 def get_private_image_tag() -> str:
-  return os.getenv('PRIVATE_IMAGE_TAG') or '20240330_0352_RC00'
+  return os.getenv('PRIVATE_IMAGE_TAG') or refined_image_versions.IMAGE_TAG
 
 
 def get_autosxs_image_tag() -> str:

google_cloud_pipeline_components/_implementation/llm/function_based.py

@@ -231,8 +231,8 @@ def resolve_reference_model_metadata(
               'gs://vertex-llm-restricted/cloud-llm-restricted/checkpoints/'
               'safe_flan_t5/xxl/v1/checkpoint_1190000/'
           ),
-          reward_model_reference='T5_XL',
-          reward_model_path='gs://t5-data/pretrained_models/t5x/t5_1_1_xl',
+          reward_model_reference='T5_XXL',
+          reward_model_path='gs://t5-data/pretrained_models/t5x/t5_1_1_xxl',
           is_supported=True,
       ),
       'palm-tiny': reference_model_metadata(
@@ -265,8 +265,10 @@ def resolve_reference_model_metadata(
           reference_model_path=(
               'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
           ),
-          reward_model_reference='OTTER',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_otter_pretrain/',
+          reward_model_reference='BISON',
+          reward_model_path=(
+              'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
+          ),
           is_supported=False,  # Deprecated: Use text-bision@001 instead.
       ),
       'text-bison@001': reference_model_metadata(
@@ -274,8 +276,10 @@ def resolve_reference_model_metadata(
           reference_model_path=(
               'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
           ),
-          reward_model_reference='OTTER',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_otter_pretrain/',
+          reward_model_reference='BISON',
+          reward_model_path=(
+              'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
+          ),
           is_supported=True,
       ),
       'text-bison@002': reference_model_metadata(
@@ -292,8 +296,10 @@ def resolve_reference_model_metadata(
           reference_model_path=(
               'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
           ),
-          reward_model_reference='OTTER',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_otter_pretrain/',
+          reward_model_reference='BISON',
+          reward_model_path=(
+              'gs://vertex-rlhf-restricted/pretrained_models/palm/t5x_bison/'
+          ),
           is_supported=True,
       ),
       'elephant': reference_model_metadata(
@@ -372,46 +378,6 @@ def convert_to_delimited_string(items: List[str], delimiter: str = ',') -> str:
   return delimiter.join(items)
 
 
-@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
-def generate_default_instruction(
-    task: str,
-    target_sequence_length: int,
-    instruction_override: str = '',
-) -> str:
-  """Generates a default instruction if no override is provided."""
-  if instruction_override:
-    return instruction_override
-  task = task.lower()
-  if task == 'summarization':
-    return f'Summarize in less than {target_sequence_length} words.'
-
-  elif task == 'question_answer':
-    return f'Answer the question in less than {target_sequence_length} words.'
-
-  else:
-    raise ValueError(
-        f'Task not recognized: {task}. Supported tasks are: "summarization",'
-        ' "question_answer".'
-    )
-
-
-@dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
-def resolve_upload_location(upload_location: Optional[str] = None) -> str:
-  """Gets the region to upload the model.
-
-  Args:
-    upload_location: User-specified region to upload the model to.
-
-  Returns:
-    Where to upload the model. If no location is specified, the model will be
-    uploaded to the region where the pipeline is running.
-  """
-  # pylint: disable=g-import-not-at-top
-  import os
-  # pylint: enable=g-import-not-at-top
-  return upload_location or os.environ['CLOUD_ML_REGION']
-
-
 @dsl.component(base_image=_image.GCPC_IMAGE_TAG, install_kfp_package=False)
 def resolve_regional_endpoint(upload_location: str) -> str:
   """Gets the regional endpoint used to upload a model to the registry.

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 = '20240327_1338'
+IMAGE_TAG = '20240425_1027_RC00'

google_cloud_pipeline_components/_implementation/llm/reinforcement_learning_graph.py

@@ -41,6 +41,14 @@ def pipeline(
     input_reward_adapter_path: str,
     input_preference_dataset_path: str,
     large_model_reference: str,
+    reward_model_reference: str,
+    policy_model_reference: str,
+    policy_model_path: str,
+    machine_type: str,
+    tuning_location: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    rl_image_uri: str,
     prompt_sequence_length: int = 512,
     target_sequence_length: int = 64,
     lora_dim: int = 1,
@@ -51,10 +59,10 @@ def pipeline(
     kl_coeff: float = 0.1,
     instruction: Optional[str] = None,
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
-    accelerator_type: str = 'GPU',
     location: str = _placeholders.LOCATION_PLACEHOLDER,
     tensorboard_resource_id: str = '',
     encryption_spec_key_name: str = '',
+    num_microbatches: int = 0,
 ) -> PipelineOutput:
   # fmt: off
   """Trains a reward model.
@@ -64,6 +72,14 @@ def pipeline(
     input_reward_adapter_path: Path to the reward LoRA adapter to use during reinforcement learning.
     input_preference_dataset_path: Path to preference dataset used by the reward model.
     large_model_reference: Name of the base model. Supported values are `text-bison@001`, `t5-small`, `t5-large`, `t5-xl` and `t5-xxl`. `text-bison@001` and `t5-small` are supported in `us-central1` and `europe-west4`. `t5-large`, `t5-xl` and `t5-xxl` are only supported in `europe-west4`.
+    reward_model_reference: Name of the reward model. The name should be in capitalized snake case format.
+    policy_model_reference: Name of the policy model. The name should be in capitalized snake case format.
+    policy_model_path: The model checkpoint path to the reinforcer model.
+    machine_type: The type of the machine to provision for the custom job. Must be a valid GCE instance type and compatible with the accelerator type.
+    tuning_location: The GCP region to run the custom job.
+    accelerator_type: Specific accelerator type for the custom job.
+    accelerator_count: The number of accelerator.
+    rl_image_uri: Docker image URI to use for the reinforcement learning training job.
     prompt_sequence_length: Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most 8192. Default value is 512.
     target_sequence_length: Maximum tokenized sequence length for target text. Higher values increase memory overhead. This value should be at most 1024. Default value is 64.
     lora_dim: The rank of the LoRA adapter. If >0, then use LoRA-tuning. If =0, then use full-tuning. Default is 1.
@@ -74,7 +90,6 @@ def pipeline(
     kl_coeff: Coefficient for KL penalty. This regularizes the policy model and penalizes if it diverges from its initial distribution. If set to 0, the reference language model is not loaded into memory. Default value is 0.1.
     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.
-    accelerator_type: One of 'TPU' or 'GPU'. If 'TPU' is specified, tuning components run in europe-west4. Otherwise tuning components run in us-central1 on GPUs. Default is 'GPU'.
     location: Location used to run non-tuning components, i.e. components that do not require accelerators. If not specified the location used to run the pipeline will be used.
     tensorboard_resource_id: Optional tensorboard resource id in format `projects/{project_number}/locations/{location}/tensorboards/{tensorboard_id}`. If provided, tensorboard metrics will be uploaded to this location.
     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.
@@ -85,14 +100,6 @@ def pipeline(
   """
   # fmt: on
   prompt_column = 'input_text'
-  machine_spec = function_based.resolve_machine_spec(
-      accelerator_type=accelerator_type,
-      use_test_spec=env.get_use_test_machine_spec(),
-  ).set_display_name('Resolve Machine Spec')
-
-  reference_model_metadata = function_based.resolve_reference_model_metadata(
-      large_model_reference=large_model_reference,
-  ).set_display_name('Resolve Model Metadata')
 
   processed_dataset = preprocess_chat_dataset.preprocess_chat_dataset(
       large_model_reference=large_model_reference,
@@ -109,30 +116,18 @@ def pipeline(
           # Target field name does not matter because this field is not used.
           targets_field_name='non_existent_targets_field_name',
           output_split_name=env.TRAIN_SPLIT,
-          large_model_reference=reference_model_metadata.outputs[
-              'large_model_reference'
-          ],
+          large_model_reference=policy_model_reference,
           instruction=instruction,
           encryption_spec_key_name=encryption_spec_key_name,
       )
       .set_display_name('Import Prompt Dataset')
       .set_caching_options(False)
   )
-  rl_image_uri = function_based.resolve_private_refined_image_uri(
-      accelerator_type=machine_spec.outputs['accelerator_type'],
-  ).set_display_name('Resolve Reinforcer Image URI')
-  num_microbatches = function_based.resolve_num_microbatches(
-      large_model_reference=reference_model_metadata.outputs[
-          'large_model_reference'
-      ]
-  ).set_display_name('Resolve Number of Microbatches')
   rl_model = (
       reinforcer.reinforcer(
           project=project,
-          location=machine_spec.outputs['tuning_location'],
-          input_reference_model_path=reference_model_metadata.outputs[
-              'reference_model_path'
-          ],
+          location=tuning_location,
+          input_reference_model_path=policy_model_path,
           input_reward_model_path=input_reward_model_path,
           input_reward_adapter_path=input_reward_adapter_path,
           input_dataset_path=prompt_dataset_importer.outputs[
@@ -140,16 +135,12 @@ def pipeline(
           ],
           input_preference_dataset_path=input_preference_dataset_path,
           train_steps=reinforcement_learning_train_steps,
-          accelerator_type=machine_spec.outputs['accelerator_type'],
-          accelerator_count=machine_spec.outputs['accelerator_count'],
-          large_model_reference=reference_model_metadata.outputs[
-              'large_model_reference'
-          ],
-          reward_model_reference=reference_model_metadata.outputs[
-              'reward_model_reference'
-          ],
-          machine_type=machine_spec.outputs['machine_type'],
-          image_uri=rl_image_uri.output,
+          accelerator_type=accelerator_type,
+          accelerator_count=accelerator_count,
+          large_model_reference=policy_model_reference,
+          reward_model_reference=reward_model_reference,
+          machine_type=machine_type,
+          image_uri=rl_image_uri,
           inputs_sequence_length=prompt_sequence_length,
           targets_sequence_length=target_sequence_length,
           batch_size=batch_size,
@@ -157,7 +148,7 @@ def pipeline(
           kl_coeff=kl_coeff,
           lora_dim=lora_dim,
           reward_lora_dim=reward_lora_dim,
-          num_microbatches=num_microbatches.output,
+          num_microbatches=num_microbatches,
           encryption_spec_key_name=encryption_spec_key_name,
           tensorboard_resource_id=tensorboard_resource_id,
       )

google_cloud_pipeline_components/_implementation/llm/reward_model_graph.py

@@ -26,7 +26,6 @@ import kfp
 
 PipelineOutput = NamedTuple(
     'Outputs',
-    reward_model_base_path=str,
     reward_model_adapter_path=str,
     reward_dataset_path=str,
 )
@@ -39,6 +38,13 @@ PipelineOutput = NamedTuple(
 def pipeline(
     preference_dataset: str,
     large_model_reference: str,
+    reward_model_reference: str,
+    reward_model_path: str,
+    machine_type: str,
+    tuning_location: str,
+    accelerator_type: str,
+    accelerator_count: int,
+    reward_model_image_uri: str,
     prompt_sequence_length: int = 512,
     target_sequence_length: int = 64,
     batch_size: int = 64,
@@ -48,10 +54,10 @@ def pipeline(
     eval_dataset: Optional[str] = None,
     instruction: Optional[str] = None,
     project: str = _placeholders.PROJECT_ID_PLACEHOLDER,
-    accelerator_type: str = 'GPU',
     location: str = _placeholders.LOCATION_PLACEHOLDER,
     tensorboard_resource_id: str = '',
     encryption_spec_key_name: str = '',
+    num_microbatches: int = 0,
 ) -> PipelineOutput:
   # fmt: off
   """Trains a reward model.
@@ -59,6 +65,13 @@ def pipeline(
   Args:
     preference_dataset: Cloud storage path to a human preference JSONL dataset used to train a reward model. Each example in a preference dataset must contain `candidate_0` and `candidate_1` fields that contain candidate responses, `choice` that specifies the preferred candidate and either `input_text` (if tuning a text model) or `messages` (if tuning a chat model). Chat datasets must contain at least 1 message in a `messages` field. Each message must be valid JSON that contains `author` and `content` fields, where valid `author` values are `user` and `assistant` and `content` must be non-empty. Each row may contain multiple messages, but the first and last author must be the `user`. An optional `context` field may be provided for each example in a chat dataset. If provided, the `context` will preprended to the message `content`. The `instruction` serves as the default context. (Useful if most messages use the same system-level context.) Any context provided in the example will override the default value.
     large_model_reference: Name of the base model. Supported values are `text-bison@001`, `t5-small`, `t5-large`, `t5-xl` and `t5-xxl`. `text-bison@001` and `t5-small` are supported in `us-central1` and `europe-west4`. `t5-large`, `t5-xl` and `t5-xxl` are only supported in `europe-west4`.
+    reward_model_reference: Name of the base model. The name should be in capitalized snake case format.
+    reward_model_path: The model checkpoint path for the reward model.
+    machine_type: The type of the machine to provision for the custom job. Must be a valid GCE instance type and compatible with the accelerator type.
+    tuning_location: The GCP region to run the custom job.
+    accelerator_type: Specific accelerator type for the custom job.
+    accelerator_count: The number of accelerator.
+    reward_model_image_uri: Docker image URI to use for the reward model training job.
     prompt_sequence_length: Maximum tokenized sequence length for input text. Higher values increase memory overhead. This value should be at most 8192. Default value is 512.
     target_sequence_length:  Maximum tokenized sequence length for target text. Higher values increase memory overhead. This value should be at most 1024. Default value is 64.
     batch_size: Number of examples in each finetuning step. Default is 64.
@@ -67,13 +80,12 @@ def pipeline(
     reward_model_train_steps: Number of steps to use when training a reward model. Default value is 1000.
     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.
-    accelerator_type: One of 'TPU' or 'GPU'. If 'TPU' is specified, tuning components run in europe-west4. Otherwise tuning components run in us-central1 on GPUs. Default is 'GPU'.
     location: Location used to run non-tuning components, i.e. components that do not require accelerators. If not specified the location used to run the pipeline will be used.
     tensorboard_resource_id: Optional tensorboard resource id in format `projects/{project_number}/locations/{location}/tensorboards/{tensorboard_id}`. If provided, tensorboard metrics will be uploaded to this location.
     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.
+    num_microbatches: The number of microbatches to break the total batch size into during training.
 
   Returns:
-    reward_model_base_path: Path to the base model used by the reward model.
     reward_model_adapter_path: Path to the output LoRA adapter.
     reward_dataset_path: Preference dataset use for tuning the reward model.
   """
@@ -81,14 +93,6 @@ def pipeline(
   prompt_column = 'input_text'
   candidate_columns = ['candidate_0', 'candidate_1']
   choice_column = 'choice'
-  machine_spec = function_based.resolve_machine_spec(
-      accelerator_type=accelerator_type,
-      use_test_spec=env.get_use_test_machine_spec(),
-  ).set_display_name('Resolve Machine Spec')
-
-  reference_model_metadata = function_based.resolve_reference_model_metadata(
-      large_model_reference=large_model_reference,
-  ).set_display_name('Resolve Model Metadata')
 
   processed_preference_dataset = (
       preprocess_chat_dataset.preprocess_chat_dataset(
@@ -113,9 +117,7 @@ def pipeline(
           comma_separated_candidates_field_names=comma_separated_candidates_field_names.output,
           choice_field_name=choice_column,
           split=env.TRAIN_SPLIT,
-          large_model_reference=reference_model_metadata.outputs[
-              'reward_model_reference'
-          ],
+          large_model_reference=reward_model_reference,
           instruction=instruction,
           encryption_spec_key_name=encryption_spec_key_name,
       )
@@ -132,9 +134,7 @@ def pipeline(
           comma_separated_candidates_field_names=comma_separated_candidates_field_names.output,
           choice_field_name=choice_column,
           split=env.TRAIN_SPLIT,
-          large_model_reference=reference_model_metadata.outputs[
-              'reward_model_reference'
-          ],
+          large_model_reference=reward_model_reference,
           instruction=instruction,
           encryption_spec_key_name=encryption_spec_key_name,
       )
@@ -142,21 +142,11 @@ def pipeline(
       .set_caching_options(False)
   )
 
-  reward_model_image_uri = function_based.resolve_private_refined_image_uri(
-      accelerator_type=machine_spec.outputs['accelerator_type'],
-  ).set_display_name('Resolve Reward Model Image URI')
-  num_microbatches = function_based.resolve_num_microbatches(
-      large_model_reference=reference_model_metadata.outputs[
-          'reward_model_reference'
-      ]
-  ).set_display_name('Resolve Number of Microbatches')
   reward_model = (
       reward_model_trainer.reward_model_trainer(
           project=project,
-          location=machine_spec.outputs['tuning_location'],
-          input_model_path=reference_model_metadata.outputs[
-              'reward_model_path'
-          ],
+          location=tuning_location,
+          input_model_path=reward_model_path,
           input_dataset_path=preference_dataset_importer.outputs[
               'output_dataset_path'
           ],
@@ -164,19 +154,17 @@ def pipeline(
               'output_dataset_path'
           ],
           train_steps=reward_model_train_steps,
-          accelerator_type=machine_spec.outputs['accelerator_type'],
-          accelerator_count=machine_spec.outputs['accelerator_count'],
-          large_model_reference=reference_model_metadata.outputs[
-              'reward_model_reference'
-          ],
-          machine_type=machine_spec.outputs['machine_type'],
-          image_uri=reward_model_image_uri.output,
+          accelerator_type=accelerator_type,
+          accelerator_count=accelerator_count,
+          large_model_reference=reward_model_reference,
+          machine_type=machine_type,
+          image_uri=reward_model_image_uri,
           inputs_sequence_length=prompt_sequence_length,
           targets_sequence_length=target_sequence_length,
           batch_size=batch_size,
           learning_rate_multiplier=reward_model_learning_rate_multiplier,
           lora_dim=lora_dim,
-          num_microbatches=num_microbatches.output,
+          num_microbatches=num_microbatches,
           encryption_spec_key_name=encryption_spec_key_name,
           tensorboard_resource_id=tensorboard_resource_id,
       )
@@ -185,9 +173,6 @@ def pipeline(
   )
 
   return PipelineOutput(
-      reward_model_base_path=reference_model_metadata.outputs[
-          'reward_model_path'
-      ],
       reward_model_adapter_path=reward_model.outputs['output_adapter_path'],
       reward_dataset_path=preference_dataset_importer.outputs[
           'output_dataset_path'

google_cloud_pipeline_components/_implementation/llm/rlhf_preprocessor.py

@@ -23,24 +23,65 @@ from kfp import dsl
 
 @dsl.container_component
 def rlhf_preprocessor(
+    large_model_reference: str,
+    accelerator_type: str,
+    use_test_spec: bool,
+    project: str,
+    location: str,
+    artifact_registry: str,
+    tag: str,
     gcp_resources: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
     has_tensorboard_id: dsl.OutputPath(bool),  # pytype: disable=invalid-annotation
     has_inference_dataset: dsl.OutputPath(bool),  # pytype: disable=invalid-annotation
+    metadata_large_model_reference: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_reference_model_path: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_reward_model_reference: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_reward_model_path: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_machine_type: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_tuning_location: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    metadata_accelerator_type: dsl.OutputPath(str),  # pytype: disable=invalid-annotation
+    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
+    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', ''),
 ) -> dsl.ContainerSpec:  # pylint: disable=g-doc-args
+  # fmt: off
   """Preprocess RLHF pipeline inputs.
 
   Args:
+    large_model_reference: The model for fine tuning.
+    accelerator_type: Specific accelerator type for the job.
+    use_test_spec: Whether to use a lower resource machine for testing.
+    project: Project that contains the artifact registry.
+    location: Region that contains the artifact registry.
+    artifact_registry: Registry that contains Docker images.
+    tag: Image tag.
+    use_experimental_image:  Whether to use refined experimental image.
     evaluation_dataset: Path to evaluation data.
     tensorboard_resource_id: TensorBoard resource id.
+    metadata_large_model_reference: The base model for fine tuning. The name should be in capitalized snake case format.
+    metadata_reference_model_path: The model checkpoint path for the reinforcer model
+    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.
 
   Returns:
     gcp_resources: GCP resources that can be used to track the custom job.
     has_tensorboard_id: Whether a tensorboard id is provided.
     has_inference_dataset: Whether inference data are provided.
+    metadata_machine_type: The type of the machine to provision for the custom job.
+    metadata_tuning_location: The GCP region to run the custom job.
+    metadata_accelerator_type: Specific accelerator type for the custom job.
+    metadata_accelerator_count: The number of accelerator.
+    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.
   """
+  # fmt: on
   return gcpc_utils.build_serverless_customjob_container_spec(
       project=_placeholders.PROJECT_ID_PLACEHOLDER,
       location=_placeholders.LOCATION_PLACEHOLDER,
@@ -52,8 +93,27 @@ def rlhf_preprocessor(
               '--app_name=rlhf_preprocessor',
               f'--evaluation_dataset={evaluation_dataset}',
               f'--tensorboard_resource_id={tensorboard_resource_id}',
+              f'--large_model_reference={large_model_reference}',
+              f'--input_reference_model_path={input_reference_model_path}',
+              f'--accelerator_type={accelerator_type}',
+              f'--use_test_spec={use_test_spec}',
+              f'--project={project}',
+              f'--location={location}',
+              f'--artifact_registry={artifact_registry}',
+              f'--tag={tag}',
+              f'--use_experimental_image={use_experimental_image}',
               f'--has_tensorboard_id_path={has_tensorboard_id}',
               f'--has_inference_dataset_path={has_inference_dataset}',
+              f'--metadata_large_model_reference_path={metadata_large_model_reference}',
+              f'--metadata_reference_model_path_path={metadata_reference_model_path}',
+              f'--metadata_reward_model_reference_path={metadata_reward_model_reference}',
+              f'--metadata_reward_model_path_path={metadata_reward_model_path}',
+              f'--metadata_machine_type_path={metadata_machine_type}',
+              f'--metadata_tuning_location_path={metadata_tuning_location}',
+              f'--metadata_accelerator_type_path={metadata_accelerator_type}',
+              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}',
           ],
       ),
       gcp_resources=gcp_resources,

google_cloud_pipeline_components/_implementation/llm/validate_pipeline.py

@@ -79,8 +79,19 @@ def validate_pipeline(
     # ]
     # [ Check CMEK
     supported_pipeline_regions = {
+        'asia-northeast1',
+        'asia-northeast3',
+        'asia-southeast1',
+        'europe-west1',
+        'europe-west2',
+        'europe-west3',
         'europe-west4',
+        'europe-west9',
+        'northamerica-northeast1',
         'us-central1',
+        'us-east4',
+        'us-west1',
+        'us-west4',
     }
     if location not in supported_pipeline_regions:
       raise ValueError(

google_cloud_pipeline_components/_placeholders.py

@@ -13,9 +13,38 @@
 # limitations under the License.
 """Placeholders for use in component authoring."""
 
-# prefer not using PIPELINE_TASK_ prefix like KFP does for reduced verbosity
+# prefer not using PIPELINE_TASK_ or PIPELINE_ prefix like KFP does for reduced
+# verbosity
 PROJECT_ID_PLACEHOLDER = "{{$.pipeline_google_cloud_project_id}}"
+"""A placeholder used to obtain Google Cloud project id where the pipeline
+executes. The placeholder value is set at pipeline runtime.
+"""
 LOCATION_PLACEHOLDER = "{{$.pipeline_google_cloud_location}}"
+"""A placeholder used to obtain Google Cloud location where the pipeline
+executes. The placeholder value is set at pipeline runtime.
+"""
+SERVICE_ACCOUNT_PLACEHOLDER = "{{$.pipeline_service_account}}"
+"""A placeholder used to obtain service account that is defined in [PipelineJob](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.pipelineJobs).
+If PipelineJob doesn't have a service account set, this placeholder will be resolved to default service account.
+The placeholder value is set at pipeline runtime.
+"""
+NETWORK_PLACEHOLDER = "{{$.pipeline_network}}"
+"""A placeholder used to obtain network that is defined in [PipelineJob](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.pipelineJobs).
+If PipelineJob doesn't have a network set, this placeholder will be empty. The
+placeholder value is set at pipeline runtime.
+"""
+PERSISTENT_RESOURCE_ID_PLACEHOLDER = "{{$.pipeline_persistent_resource_id}}"
+"""A placeholder used to obtain persistent resource id that is defined in
+PipelineJob [RuntimeConfig](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.pipelineJobs#PipelineJob.RuntimeConfig).
+If PipelineJob doesn't have a persistent resource id, this placeholder will be
+empty. The placeholder value is set at pipeline runtime.
+"""
+ENCRYPTION_SPEC_KMS_KEY_NAME_PLACEHOLDER = "{{$.pipeline_encryption_key_name}}"
+"""A placeholder used to obtain kmsKeyName that is defined in
+PipelineJob's [EncryptionSpec](https://cloud.google.com/vertex-ai/docs/reference/rest/v1/EncryptionSpec).
+If PipelineJob doesn't have a encryption key name, this placeholder will be
+empty. The placeholder value is set at pipeline runtime.
+"""
 
 
 # omit placeholder type annotation to avoid dependency on KFP SDK internals

google_cloud_pipeline_components/preview/automl/forecasting/forecasting_ensemble.py

@@ -72,7 +72,7 @@ def automl_forecasting_ensemble(
   # fmt: on
   job_id = dsl.PIPELINE_JOB_ID_PLACEHOLDER
   task_id = dsl.PIPELINE_TASK_ID_PLACEHOLDER
-  image_uri = 'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240214_1325'
+  image_uri = 'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240419_0625'
   display_name = f'automl-forecasting-ensemble-{job_id}-{task_id}'
 
   error_file_path = f'{root_dir}/{job_id}/{task_id}/error.pb'

google_cloud_pipeline_components/preview/automl/forecasting/forecasting_stage_1_tuner.py

@@ -99,14 +99,14 @@ def automl_forecasting_stage_1_tuner(
                       ' 1, "machine_spec": {"machine_type": "n1-standard-8"},'
                       ' "container_spec": {"image_uri":"'
                   ),
-                  'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240214_1325',
+                  'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240419_0625',
                   '", "args": ["forecasting_mp_l2l_stage_1_tuner',
                   '", "--region=',
                   location,
                   '", "--transform_output_path=',
                   transform_output.uri,
                   '", "--training_docker_uri=',
-                  'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240214_1325',
+                  'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240419_0625',
                   '", "--reduce_search_space_mode=',
                   reduce_search_space_mode,
                   f'", "--component_id={dsl.PIPELINE_TASK_ID_PLACEHOLDER}',