google-cloud-pipeline-components 2.14.1__py3-none-any.whl → 2.16.0__py3-none-any.whl

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 = '20240506_1707'
+IMAGE_TAG = '20240623_1707'

google_cloud_pipeline_components/_implementation/model_evaluation/llm_evaluation_preprocessor/component.py

@@ -38,6 +38,9 @@ def evaluation_dataset_preprocessor_internal(
     output_dirs: dsl.OutputPath(list),
     gcp_resources: dsl.OutputPath(str),
     input_field_name: str = 'input_text',
+    role_field_name: str = 'role',
+    target_field_name: str = 'ground_truth',
+    model_name: str = 'publishers/google/model/text-bison@002',
     display_name: str = 'llm_evaluation_dataset_preprocessor_component',
     machine_type: str = 'e2-highmem-16',
     service_account: str = '',
@@ -56,6 +59,11 @@ def evaluation_dataset_preprocessor_internal(
       gcs_source_uris: A json escaped list of GCS URIs of the input eval dataset.
       input_field_name: The field name of the input eval dataset instances that
         contains the input prompts to the LLM.
+      role_field_name: The field name of the role for input eval dataset instances
+        that contains the input prompts to the LLM.
+      target_field_name: The field name of the target for input eval dataset
+        instances.
+      model_name: Name of the model being used to create model-specific schemas.
       machine_type: The machine type of this custom job. If not set, defaulted
         to `e2-highmem-16`. More details:
         https://cloud.google.com/compute/docs/machine-resource
@@ -92,6 +100,11 @@ def evaluation_dataset_preprocessor_internal(
               f'--eval_dataset_preprocessor={True}',
               f'--gcs_source_uris={gcs_source_uris}',
               f'--input_field_name={input_field_name}',
+              f'--role_field_name={role_field_name}',
+              (
+                  f'--target_field_name={target_field_name}'
+                  f'--model_name={model_name}'
+              ),
               f'--output_dirs={output_dirs}',
               '--executor_input={{$.json_escape[1]}}',
           ],
@@ -109,6 +122,9 @@ def llm_evaluation_dataset_preprocessor_graph_component(
     location: str,
     gcs_source_uris: List[str],
     input_field_name: str = 'input_text',
+    role_field_name: str = 'role',
+    target_field_name: str = 'ground_truth',
+    model_name: str = 'publishers/google/model/text-bison@002',
     display_name: str = 'llm_evaluation_dataset_preprocessor_component',
     machine_type: str = 'e2-standard-4',
     service_account: str = '',
@@ -126,6 +142,11 @@ def llm_evaluation_dataset_preprocessor_graph_component(
       gcs_source_uris: A list of GCS URIs of the input eval dataset.
       input_field_name: The field name of the input eval dataset instances that
         contains the input prompts to the LLM.
+      role_field_name: The field name of the role for input eval dataset
+        instances that contains the input prompts to the LLM.
+      target_field_name: The field name of the target for input eval dataset
+        instances.
+      model_name: Name of the model being used to create model-specific schemas.
       display_name: The name of the Evaluation job.
       machine_type: The machine type of this custom job. If not set, defaulted
         to `e2-standard-4`. More details:
@@ -163,6 +184,9 @@ def llm_evaluation_dataset_preprocessor_graph_component(
           input_list=gcs_source_uris
       ).output,
       input_field_name=input_field_name,
+      role_field_name=role_field_name,
+      target_field_name=target_field_name,
+      model_name=model_name,
       display_name=display_name,
       machine_type=machine_type,
       service_account=service_account,

google_cloud_pipeline_components/_implementation/starry_net/__init__.py

@@ -0,0 +1,41 @@
+# Copyright 2024 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.
+from google_cloud_pipeline_components._implementation.starry_net.dataprep.component import dataprep as DataprepOp
+from google_cloud_pipeline_components._implementation.starry_net.evaluation.component import evaluation as EvaluationOp
+from google_cloud_pipeline_components._implementation.starry_net.get_training_artifacts.component import get_training_artifacts as GetTrainingArtifactsOp
+from google_cloud_pipeline_components._implementation.starry_net.maybe_set_tfrecord_args.component import maybe_set_tfrecord_args as MaybeSetTfrecordArgsOp
+from google_cloud_pipeline_components._implementation.starry_net.set_dataprep_args.component import set_dataprep_args as SetDataprepArgsOp
+from google_cloud_pipeline_components._implementation.starry_net.set_eval_args.component import set_eval_args as SetEvalArgsOp
+from google_cloud_pipeline_components._implementation.starry_net.set_test_set.component import set_test_set as SetTestSetOp
+from google_cloud_pipeline_components._implementation.starry_net.set_tfrecord_args.component import set_tfrecord_args as SetTfrecordArgsOp
+from google_cloud_pipeline_components._implementation.starry_net.set_train_args.component import set_train_args as SetTrainArgsOp
+from google_cloud_pipeline_components._implementation.starry_net.train.component import train as TrainOp
+from google_cloud_pipeline_components._implementation.starry_net.upload_decomposition_plots.component import upload_decomposition_plots as UploadDecompositionPlotsOp
+from google_cloud_pipeline_components._implementation.starry_net.upload_model.component import upload_model as UploadModelOp
+
+
+__all__ = [
+    'DataprepOp',
+    'EvaluationOp',
+    'GetTrainingArtifactsOp',
+    'MaybeSetTfrecordArgsOp',
+    'SetDataprepArgsOp',
+    'SetEvalArgsOp',
+    'SetTestSetOp',
+    'SetTfrecordArgsOp',
+    'SetTrainArgsOp',
+    'TrainOp',
+    'UploadDecompositionPlotsOp',
+    'UploadModelOp',
+]

google_cloud_pipeline_components/_implementation/{model_evaluation/import_evaluation → starry_net/dataprep}/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2023 The Kubeflow Authors. All Rights Reserved.
+# Copyright 2024 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.
@@ -11,4 +11,3 @@
 # 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.
-"""Google Cloud Pipeline Evaluation Import Evaluation Component."""

google_cloud_pipeline_components/_implementation/starry_net/dataprep/component.py

@@ -0,0 +1,173 @@
+# Copyright 2024 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.
+"""Starry Net component for data preparation."""
+
+from google_cloud_pipeline_components import utils
+from google_cloud_pipeline_components._implementation.starry_net import version
+from kfp import dsl
+
+
+@dsl.container_component
+def dataprep(
+    gcp_resources: dsl.OutputPath(str),
+    dataprep_dir: dsl.Output[dsl.Artifact],  # pytype: disable=unsupported-operands
+    backcast_length: int,
+    forecast_length: int,
+    train_end_date: str,
+    n_val_windows: int,
+    n_test_windows: int,
+    test_set_stride: int,
+    model_blocks: str,
+    bigquery_source: str,
+    ts_identifier_columns: str,
+    time_column: str,
+    static_covariate_columns: str,
+    static_covariates_vocab_path: str,  # pytype: disable=unused-argument
+    target_column: str,
+    machine_type: str,
+    docker_region: str,
+    location: str,
+    project: str,
+    job_id: str,
+    job_name_prefix: str,
+    num_workers: int,
+    max_num_workers: int,
+    disk_size_gb: int,
+    test_set_only: bool,
+    bigquery_output: str,
+    nan_threshold: float,
+    zero_threshold: float,
+    gcs_source: str,
+    gcs_static_covariate_source: str,
+    encryption_spec_key_name: str,
+):
+  # fmt: off
+  """Runs Dataprep for training and evaluating a STARRY-Net model.
+
+  Args:
+    gcp_resources: Serialized JSON of ``gcp_resources`` which tracks the
+      CustomJob.
+    dataprep_dir: The gcp bucket path where all dataprep artifacts
+      are saved.
+    backcast_length: The length of the input window to feed into the model.
+    forecast_length: The length of the forecast horizon.
+    train_end_date: The last date of data to use in the training set. All
+      subsequent dates are part of the test set.
+    n_val_windows: The number of windows to use for the val set. If 0, no
+      validation set is used.
+    n_test_windows: The number of windows to use for the test set. Must be >= 1.
+    test_set_stride: The number of timestamps to roll forward when
+      constructing the val and test sets.
+    model_blocks: The stringified tuple of blocks to use in the order
+      that they appear in the model. Possible values are `cleaning`,
+      `change_point`, `trend`, `hour_of_week-hybrid`, `day_of_week-hybrid`,
+      `day_of_year-hybrid`, `week_of_year-hybrid`, `month_of_year-hybrid`,
+      `residual`, `quantile`.
+    bigquery_source: The BigQuery source of the data.
+    ts_identifier_columns: The columns that identify unique time series in the BigQuery
+      data source.
+    time_column: The column with timestamps in the BigQuery source.
+    static_covariate_columns: The names of the staic covariates.
+    static_covariates_vocab_path: The path to the master static covariates vocab
+      json.
+    target_column: The target column in the Big Query data source.
+    machine_type: The machine type of the dataflow workers.
+    docker_region: The docker region, used to determine which image to use.
+    location: The location where the job is run.
+    project: The name of the project.
+    job_id: The pipeline job id.
+    job_name_prefix: The name of the dataflow job name prefix.
+    num_workers: The initial number of workers in the dataflow job.
+    max_num_workers: The maximum number of workers in the dataflow job.
+    disk_size_gb: The disk size of each dataflow worker.
+    test_set_only: Whether to only create the test set BigQuery table or also
+      to create TFRecords for traiing and validation.
+    bigquery_output: The BigQuery dataset where the test set is written in the
+      form bq://project.dataset.
+    nan_threshold: Series having more nan / missing values than
+      nan_threshold (inclusive) in percentage for either backtest or forecast
+      will not be sampled in the training set (including missing due to
+      train_start and train_end). All existing nans are replaced by zeros.
+    zero_threshold: Series having more 0.0 values than zero_threshold
+      (inclusive) in percentage for either backtest or forecast will not be
+      sampled in the training set.
+    gcs_source: The path the csv file of the data source.
+    gcs_static_covariate_source: The path to the csv file of static covariates.
+    encryption_spec_key_name: Customer-managed encryption key options for the
+      CustomJob. If this is set, then all resources created by the CustomJob
+      will be encrypted with the provided encryption key.
+
+  Returns:
+    gcp_resources: Serialized JSON of ``gcp_resources`` which tracks the
+      CustomJob.
+    dataprep_dir: The gcp bucket path where all dataprep artifacts
+      are saved.
+  """
+  job_name = f'{job_name_prefix}-{job_id}'
+  payload = {
+      'display_name': job_name,
+      'encryption_spec': {
+          'kms_key_name': str(encryption_spec_key_name),
+      },
+      'job_spec': {
+          'worker_pool_specs': [{
+              'replica_count': '1',
+              'machine_spec': {
+                  'machine_type': str(machine_type),
+              },
+              'disk_spec': {
+                  'boot_disk_type': 'pd-ssd',
+                  'boot_disk_size_gb': 100,
+              },
+              'container_spec': {
+                  'image_uri': f'{docker_region}-docker.pkg.dev/vertex-ai-restricted/starryn/dataprep:captain_{version.DATAPREP_VERSION}',
+                  'args': [
+                      '--config=starryn/experiments/configs/vertex.py',
+                      f'--config.datasets.backcast_length={backcast_length}',
+                      f'--config.datasets.forecast_length={forecast_length}',
+                      f'--config.datasets.train_end_date={train_end_date}',
+                      f'--config.datasets.n_val_windows={n_val_windows}',
+                      f'--config.datasets.val_rolling_window_size={test_set_stride}',
+                      f'--config.datasets.n_test_windows={n_test_windows}',
+                      f'--config.datasets.test_rolling_window_size={test_set_stride}',
+                      f'--config.datasets.nan_threshold={nan_threshold}',
+                      f'--config.datasets.zero_threshold={zero_threshold}',
+                      f'--config.model.static_cov_names={static_covariate_columns}',
+                      f'--config.model.blocks_list={model_blocks}',
+                      f'--bigquery_source={bigquery_source}',
+                      f'--bigquery_output={bigquery_output}',
+                      f'--gcs_source={gcs_source}',
+                      f'--gcs_static_covariate_source={gcs_static_covariate_source}',
+                      f'--ts_identifier_columns={ts_identifier_columns}',
+                      f'--time_column={time_column}',
+                      f'--target_column={target_column}',
+                      f'--job_id={job_name}',
+                      f'--num_workers={num_workers}',
+                      f'--max_num_workers={max_num_workers}',
+                      f'--root_bucket={dataprep_dir.uri}',
+                      f'--disk_size={disk_size_gb}',
+                      f'--machine_type={machine_type}',
+                      f'--test_set_only={test_set_only}',
+                      f'--image_uri={docker_region}-docker.pkg.dev/vertex-ai-restricted/starryn/dataprep:replica_{version.DATAPREP_VERSION}',
+                  ],
+              },
+          }]
+      }
+    }
+  return utils.build_serverless_customjob_container_spec(
+      project=project,
+      location=location,
+      custom_job_payload=payload,
+      gcp_resources=gcp_resources,
+    )

google_cloud_pipeline_components/_implementation/starry_net/evaluation/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 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.

google_cloud_pipeline_components/_implementation/starry_net/evaluation/component.py

@@ -0,0 +1,23 @@
+# Copyright 2024 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.
+"""StarryNet Evaluation Component."""
+
+import os
+
+from kfp import components
+
+# TODO(b/346580764)
+evaluation = components.load_component_from_file(
+    os.path.join(os.path.dirname(__file__), 'evaluation.yaml')
+)

google_cloud_pipeline_components/_implementation/starry_net/evaluation/evaluation.yaml

@@ -0,0 +1,197 @@
+name: model_evaluation_forecasting
+description: |
+  Computes a google.ForecastingMetrics Artifact, containing evaluation metrics given a model's prediction results.
+  Creates a dataflow job with Apache Beam and TFMA to compute evaluation metrics.
+  Supports point forecasting and quantile forecasting for tabular data.
+  Args:
+      project (str):
+          Project to run evaluation container.
+      location (Optional[str]):
+          Location for running the evaluation.
+          If not set, defaulted to `us-central1`.
+      root_dir (str):
+          The GCS directory for keeping staging files.
+          A random subdirectory will be created under the directory to keep job info for resuming
+          the job in case of failure.
+      predictions_format (Optional[str]):
+          The file format for the batch prediction results. `jsonl` is currently the only allowed
+          format.
+          If not set, defaulted to `jsonl`.
+      predictions_gcs_source (Optional[system.Artifact]):
+          An artifact with its URI pointing toward a GCS directory with prediction or explanation
+          files to be used for this evaluation.
+          For prediction results, the files should be named "prediction.results-*".
+          For explanation results, the files should be named "explanation.results-*".
+      predictions_bigquery_source (Optional[google.BQTable]):
+          BigQuery table with prediction or explanation data to be used for this evaluation.
+          For prediction results, the table column should be named "predicted_*".
+      ground_truth_format(Optional[str]):
+          Required for custom tabular and non tabular data.
+          The file format for the ground truth files. `jsonl` is currently the only allowed format.
+          If not set, defaulted to `jsonl`.
+      ground_truth_gcs_source(Optional[Sequence[str]]):
+          Required for custom tabular and non tabular data.
+          The GCS uris representing where the ground truth is located.
+          Used to provide ground truth for each prediction instance when they are not part of the batch prediction jobs prediction instance.
+      ground_truth_bigquery_source(Optional[str]):
+          Required for custom tabular.
+          The BigQuery table uri representing where the ground truth is located.
+          Used to provide ground truth for each prediction instance when they are not part of the batch prediction jobs prediction instance.
+      target_field_name (str):
+          The full name path of the features target field in the predictions file.
+          Formatted to be able to find nested columns, delimited by `.`.
+          Alternatively referred to as the ground truth (or ground_truth_column) field.
+      model (Optional[google.VertexModel]):
+          The Model used for predictions job.
+          Must share the same ancestor Location.
+      prediction_score_column (Optional[str]):
+          Optional. The column name of the field containing batch prediction scores.
+          Formatted to be able to find nested columns, delimited by `.`.
+          If not set, defaulted to `prediction.value` for a `point` forecasting_type and
+          `prediction.quantile_predictions` for a `quantile` forecasting_type.
+      forecasting_type (Optional[str]):
+          Optional. If the problem_type is `forecasting`, then the forecasting type being addressed
+          by this regression evaluation run. `point` and `quantile` are the supported types.
+          If not set, defaulted to `point`.
+      forecasting_quantiles (Optional[Sequence[Float]]):
+          Required for a `quantile` forecasting_type.
+          The list of quantiles in the same order appeared in the quantile prediction score column.
+          If one of the quantiles is set to `0.5f`, point evaluation will be set on that index.
+      example_weight_column (Optional[str]):
+          Optional. The column name of the field containing example weights.
+            Each value of positive_classes provided.
+      point_evaluation_quantile (Optional[Float]):
+          Required for a `quantile` forecasting_type.
+          A quantile in the list of forecasting_quantiles that will be used for point evaluation
+          metrics.
+      dataflow_service_account (Optional[str]):
+          Optional. Service account to run the dataflow job.
+          If not set, dataflow will use the default woker service account.
+          For more details, see https://cloud.google.com/dataflow/docs/concepts/security-and-permissions#default_worker_service_account
+      dataflow_disk_size (Optional[int]):
+          Optional. The disk size (in GB) of the machine executing the evaluation run.
+          If not set, defaulted to `50`.
+      dataflow_machine_type (Optional[str]):
+          Optional. The machine type executing the evaluation run.
+          If not set, defaulted to `n1-standard-4`.
+      dataflow_workers_num (Optional[int]):
+          Optional. The number of workers executing the evaluation run.
+          If not set, defaulted to `10`.
+      dataflow_max_workers_num (Optional[int]):
+          Optional. The max number of workers executing the evaluation run.
+          If not set, defaulted to `25`.
+      dataflow_subnetwork (Optional[str]):
+          Dataflow's fully qualified subnetwork name, when empty the default subnetwork will be
+          used. More details:
+          https://cloud.google.com/dataflow/docs/guides/specifying-networks#example_network_and_subnetwork_specifications
+      dataflow_use_public_ips (Optional[bool]):
+          Specifies whether Dataflow workers use public IP addresses.
+      encryption_spec_key_name (Optional[str]):
+          Customer-managed encryption key.
+  Returns:
+      evaluation_metrics (google.ForecastingMetrics):
+          google.ForecastingMetrics artifact representing the forecasting evaluation metrics in GCS.
+inputs:
+  - { name: project, type: String }
+  - { name: location, type: String, default: "us-central1" }
+  - { name: root_dir, type: system.Artifact }
+  - { name: predictions_format, type: String, default: "jsonl" }
+  - { name: predictions_gcs_source, type: Artifact, optional: True }
+  - { name: predictions_bigquery_source, type: google.BQTable, optional: True }
+  - { name: ground_truth_format, type: String, default: "jsonl" }
+  - { name: ground_truth_gcs_source, type: JsonArray, default: "[]" }
+  - { name: ground_truth_bigquery_source, type: String, default: "" }
+  - { name: target_field_name, type: String }
+  - { name: model, type: google.VertexModel, optional: True }
+  - { name: prediction_score_column, type: String, default: "" }
+  - { name: forecasting_type, type: String, default: "point" }
+  - { name: forecasting_quantiles, type: JsonArray, default: "[0.5]" }
+  - { name: example_weight_column, type: String, default: "" }
+  - { name: point_evaluation_quantile, type: Float, default: 0.5 }
+  - { name: dataflow_service_account, type: String, default: "" }
+  - { name: dataflow_disk_size, type: Integer, default: 50 }
+  - { name: dataflow_machine_type, type: String, default: "n1-standard-4" }
+  - { name: dataflow_workers_num, type: Integer, default: 1 }
+  - { name: dataflow_max_workers_num, type: Integer, default: 5 }
+  - { name: dataflow_subnetwork, type: String, default: "" }
+  - { name: dataflow_use_public_ips, type: Boolean, default: "true" }
+  - { name: encryption_spec_key_name, type: String, default: "" }
+outputs:
+  - { name: evaluation_metrics, type: google.ForecastingMetrics }
+  - { name: gcp_resources, type: String }
+implementation:
+  container:
+    image: gcr.io/ml-pipeline/model-evaluation:v0.9
+    command:
+      - python
+      - /main.py
+    args:
+      - --setup_file
+      - /setup.py
+      - --json_mode
+      - "true"
+      - --project_id
+      - { inputValue: project }
+      - --location
+      - { inputValue: location }
+      - --problem_type
+      - "forecasting"
+      - --forecasting_type
+      - { inputValue: forecasting_type }
+      - --forecasting_quantiles
+      - { inputValue: forecasting_quantiles }
+      - --point_evaluation_quantile
+      - { inputValue: point_evaluation_quantile }
+      - --batch_prediction_format
+      - { inputValue: predictions_format }
+      - if:
+          cond: {isPresent: predictions_gcs_source}
+          then:
+          - --batch_prediction_gcs_source
+          - "{{$.inputs.artifacts['predictions_gcs_source'].uri}}"
+      - if:
+          cond: {isPresent: predictions_bigquery_source}
+          then:
+          - --batch_prediction_bigquery_source
+          - "bq://{{$.inputs.artifacts['predictions_bigquery_source'].metadata['projectId']}}.{{$.inputs.artifacts['predictions_bigquery_source'].metadata['datasetId']}}.{{$.inputs.artifacts['predictions_bigquery_source'].metadata['tableId']}}"
+      - if:
+          cond: {isPresent: model}
+          then:
+          - --model_name
+          - "{{$.inputs.artifacts['model'].metadata['resourceName']}}"
+      - --ground_truth_format
+      - { inputValue: ground_truth_format }
+      - --ground_truth_gcs_source
+      - { inputValue: ground_truth_gcs_source }
+      - --ground_truth_bigquery_source
+      - { inputValue: ground_truth_bigquery_source }
+      - --root_dir
+      - "{{$.inputs.artifacts['root_dir'].uri}}"
+      - --target_field_name
+      - "instance.{{$.inputs.parameters['target_field_name']}}"
+      - --prediction_score_column
+      - { inputValue: prediction_score_column }
+      - --dataflow_job_prefix
+      - "evaluation-{{$.pipeline_job_uuid}}-{{$.pipeline_task_uuid}}"
+      - --dataflow_service_account
+      - { inputValue: dataflow_service_account }
+      - --dataflow_disk_size
+      - { inputValue: dataflow_disk_size }
+      - --dataflow_machine_type
+      - { inputValue: dataflow_machine_type }
+      - --dataflow_workers_num
+      - { inputValue: dataflow_workers_num }
+      - --dataflow_max_workers_num
+      - { inputValue: dataflow_max_workers_num }
+      - --dataflow_subnetwork
+      - { inputValue: dataflow_subnetwork }
+      - --dataflow_use_public_ips
+      - { inputValue: dataflow_use_public_ips }
+      - --kms_key_name
+      - { inputValue: encryption_spec_key_name }
+      - --output_metrics_gcs_path
+      - { outputUri: evaluation_metrics }
+      - --gcp_resources
+      - { outputPath: gcp_resources }
+      - --executor_input
+      - "{{$}}"

google_cloud_pipeline_components/_implementation/starry_net/get_training_artifacts/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 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.

google_cloud_pipeline_components/_implementation/starry_net/get_training_artifacts/component.py

@@ -0,0 +1,62 @@
+# Copyright 2024 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.
+"""StarryNet get training artifacts component."""
+
+from typing import NamedTuple
+
+from kfp import dsl
+
+
+@dsl.component(packages_to_install=['tensorflow==2.11.0'])
+def get_training_artifacts(
+    docker_region: str,
+    trainer_dir: dsl.InputPath(),
+) -> NamedTuple(
+    'TrainingArtifacts',
+    image_uri=str,
+    artifact_uri=str,
+    prediction_schema_uri=str,
+    instance_schema_uri=str,
+):
+  # fmt: off
+  """Gets the artifact URIs from the training job.
+
+  Args:
+    docker_region: The region from which the training docker image is pulled.
+    trainer_dir: The directory where training artifacts where stored.
+
+  Returns:
+    A NamedTuple containing the image_uri for the prediction server,
+    the artifact_uri with model artifacts, the prediction_schema_uri,
+    and the instance_schema_uri.
+  """
+  import os  # pylint: disable=g-import-not-at-top
+  import tensorflow as tf  # pylint: disable=g-import-not-at-top
+
+  with tf.io.gfile.GFile(os.path.join(trainer_dir, 'trainer.txt')) as f:
+    private_dir = f.read().strip()
+
+  outputs = NamedTuple(
+      'TrainingArtifacts',
+      image_uri=str,
+      artifact_uri=str,
+      prediction_schema_uri=bool,
+      instance_schema_uri=str,
+  )
+  return outputs(
+      f'{docker_region}-docker.pkg.dev/vertex-ai/starryn/predictor:20240723_0542_RC00',  # pylint: disable=too-many-function-args
+      private_dir,  # pylint: disable=too-many-function-args
+      os.path.join(private_dir, 'predict_schema.yaml'),  # pylint: disable=too-many-function-args
+      os.path.join(private_dir, 'instance_schema.yaml'),  # pylint: disable=too-many-function-args
+  )

google_cloud_pipeline_components/_implementation/starry_net/maybe_set_tfrecord_args/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2024 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.