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

google_cloud_pipeline_components/_implementation/starry_net/train/component.py

@@ -0,0 +1,220 @@
+# 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.
+"""Container Component for training STARRY-Net."""
+
+from google_cloud_pipeline_components import _placeholders
+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 train(
+    gcp_resources: dsl.OutputPath(str),
+    trainer_dir: dsl.Output[dsl.Artifact],  # pytype: disable=unsupported-operands
+    num_epochs: int,
+    backcast_length: int,
+    forecast_length: int,
+    train_end_date: str,
+    csv_data_path: str,
+    csv_static_covariates_path: str,
+    static_covariates_vocab_path: str,
+    train_tf_record_patterns: str,
+    val_tf_record_patterns: str,
+    test_tf_record_patterns: str,
+    n_decomposition_plots: int,
+    n_val_windows: int,
+    n_test_windows: int,
+    test_set_stride: int,
+    nan_threshold: float,
+    zero_threshold: float,
+    cleaning_activation_regularizer_coeff: float,
+    change_point_activation_regularizer_coeff: float,
+    change_point_output_regularizer_coeff: float,
+    alpha_upper_bound: float,
+    beta_upper_bound: float,
+    phi_lower_bound: float,
+    b_fixed_val: int,
+    b0_fixed_val: int,
+    phi_fixed_val: int,
+    quantiles: str,
+    use_static_covariates: bool,
+    static_covariate_names: str,
+    model_blocks: str,
+    freeze_point_forecasts: bool,
+    machine_type: str,
+    accelerator_type: str,
+    docker_region: str,
+    location: str,
+    job_id: str,
+    project: str,
+    encryption_spec_key_name: str,
+):
+  # fmt: off
+  """Trains a STARRY-Net model.
+
+  Args:
+    gcp_resources: Serialized JSON of ``gcp_resources`` which tracks the
+      CustomJob.
+    trainer_dir: The gcp bucket path where training artifacts are saved.
+    num_epochs: The number of epochs to train for.
+    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.
+    csv_data_path: The path to the training data csv.
+    csv_static_covariates_path: The path to the static covariates csv.
+    static_covariates_vocab_path: The path to the master static covariates vocab
+      json.
+    train_tf_record_patterns: The glob patterns to the tf records to use for
+      training.
+    val_tf_record_patterns: The glob patterns to the tf records to use for
+      validation.
+    test_tf_record_patterns: The glob patterns to the tf records to use for
+      testing.
+    n_decomposition_plots: How many decomposition plots to save to tensorboard.
+    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.
+    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.
+    cleaning_activation_regularizer_coeff: The regularization coefficient for
+      the cleaning param estimator's final layer's activation in the cleaning
+      block.
+    change_point_activation_regularizer_coeff: The regularization coefficient
+      for the change point param estimator's final layer's activation in the
+      change_point block.
+    change_point_output_regularizer_coeff: The regularization coefficient
+      for the change point param estimator's output in the change_point block.
+    alpha_upper_bound: The upper bound for data smooth parameter alpha in the
+      trend block.
+    beta_upper_bound: The upper bound for data smooth parameter beta in the
+      trend block.
+    phi_lower_bound: The lower bound for damping param phi in the trend block.
+    b_fixed_val: The fixed value for b in the trend block. If set to anything
+      other than -1, the trend block will not learn to provide estimates
+      but use the fixed value directly.
+    b0_fixed_val: The fixed value for b0 in the trend block. If set to
+      anything other than -1, the trend block will not learn to provide
+      estimates but use the fixed value directly.
+    phi_fixed_val: The fixed value for phi in the trend block. If set to
+      anything other than -1, the trend block will not learn to provide
+      estimates but use the fixed value directly.
+    quantiles: The stringified tuple of quantiles to learn in the quantile
+      block, e.g., 0.5,0.9,0.95. This should always start with 0.5,
+      representing the point forecasts.
+    use_static_covariates: Whether to use static covariates.
+    static_covariate_names: The stringified tuple of names of the static
+      covariates.
+    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`.
+    freeze_point_forecasts: Whether or not to do two pass training, where
+      first the point forecast model is trained, then the quantile block is,
+      added, all preceding blocks are frozen, and the quantile block is trained.
+      This should always be True if quantiles != [0.5].
+    machine_type: The machine type.
+    accelerator_type: The accelerator type.
+    docker_region: The docker region, used to determine which image to use.
+    location: Location for creating the custom training job. If not set,
+      defaults to us-central1.
+    job_id: The pipeline job id.
+    project: Project to create the custom training job in. Defaults to
+      the project in which the PipelineJob is run.
+    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.
+    trainer_dir: The gcp bucket path where training artifacts are saved.
+  """
+  job_name = f'trainer-{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),
+                  'accelerator_type': str(accelerator_type),
+                  'accelerator_count': 1,
+              },
+              '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/trainer:{version.TRAINER_VERSION}',
+                  'args': [
+                      f'--vertex_experiment_dir={trainer_dir.path}',
+                      f'--vertex_job_id={job_id}',
+                      '--config=analysis/trafficforecast/starryn/experiments/configs/vertex.py',
+                      f'--config.num_epochs={num_epochs}',
+                      f'--config.freeze_point_forecasts={freeze_point_forecasts}',
+                      f'--config.callbacks.tensorboard.n_decomposition_plots={n_decomposition_plots}',
+                      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.train_path={csv_data_path}',
+                      f'--config.datasets.static_covariates_path={csv_static_covariates_path}',
+                      f'--config.datasets.static_covariates_vocab_path={static_covariates_vocab_path}',
+                      f'--config.datasets.train_tf_record_patterns={train_tf_record_patterns}',
+                      f'--config.datasets.val_tf_record_patterns={val_tf_record_patterns}',
+                      f'--config.datasets.test_tf_record_patterns={test_tf_record_patterns}',
+                      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.regularizer_coeff={cleaning_activation_regularizer_coeff}',
+                      f'--config.model.activation_regularizer_coeff={change_point_activation_regularizer_coeff}',
+                      f'--config.model.output_regularizer_coeff={change_point_output_regularizer_coeff}',
+                      f'--config.model.alpha_upper_bound={alpha_upper_bound}',
+                      f'--config.model.beta_upper_bound={beta_upper_bound}',
+                      f'--config.model.phi_lower_bound={phi_lower_bound}',
+                      f'--config.model.b_fixed_val={b_fixed_val}',
+                      f'--config.model.b0_fixed_val={b0_fixed_val}',
+                      f'--config.model.phi_fixed_val={phi_fixed_val}',
+                      f'--config.model.quantiles={quantiles}',
+                      f'--config.model.use_static_covariates_trend={use_static_covariates}',
+                      f'--config.model.use_static_covariates_calendar={use_static_covariates}',
+                      f'--config.model.static_cov_names={static_covariate_names}',
+                      f'--config.model.blocks_list={model_blocks}',
+                  ],
+              },
+          }]
+      }
+    }
+  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/upload_decomposition_plots/__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/upload_decomposition_plots/component.py

@@ -0,0 +1,64 @@
+# 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 upload decomposition plots component."""
+
+from kfp import dsl
+
+
+@dsl.component(
+    packages_to_install=[
+        'google-cloud-aiplatform[tensorboard]',
+        'protobuf==3.20.*',
+    ]
+)
+def upload_decomposition_plots(
+    project: str,
+    location: str,
+    tensorboard_id: str,
+    display_name: str,
+    trainer_dir: dsl.InputPath(),
+) -> dsl.Artifact:
+  # fmt: off
+  """Uploads decomposition plots to Tensorboard.
+
+  Args:
+    project: The project where the pipeline is run. Defaults to current project.
+    location: The location where the pipeline components are run.
+    tensorboard_id: The tensorboard instance ID.
+    display_name: The diplay name of the job.
+    trainer_dir: The directory where training artifacts where stored.
+
+  Returns:
+    A dsl.Artifact where the URI is the URI where decomposition plots can be
+    viewed.
+  """
+  import os  # pylint: disable=g-import-not-at-top
+  from google.cloud import aiplatform  # pylint: disable=g-import-not-at-top
+
+  log_dir = os.path.join(trainer_dir, 'tensorboard', 'r=1:gc=0')
+  project_number = os.environ['CLOUD_ML_PROJECT_ID']
+  aiplatform.init(project=project, location=location)
+  aiplatform.upload_tb_log(
+      tensorboard_id=tensorboard_id,
+      tensorboard_experiment_name=display_name,
+      logdir=log_dir,
+      experiment_display_name=display_name,
+      description=f'Tensorboard for {display_name}',
+  )
+  uri = (
+      f'https://{location}.tensorboard.googleusercontent.com/experiment/'
+      f'projects+{project_number}+locations+{location}+tensorboards+'
+      f'{tensorboard_id}+experiments+{display_name}/#images'
+  )
+  return dsl.Artifact(uri=uri)

google_cloud_pipeline_components/_implementation/starry_net/upload_model/__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/upload_model/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.
+"""Starry Net Upload Model Component."""
+
+import os
+
+from kfp import components
+
+# TODO(b/346580764)
+upload_model = components.load_component_from_file(
+    os.path.join(os.path.dirname(__file__), 'upload_model.yaml')
+)

google_cloud_pipeline_components/_implementation/starry_net/upload_model/upload_model.yaml

@@ -0,0 +1,37 @@
+name: model_upload
+inputs:
+- {name: project, type: String}
+- {name: location, type: String, default: "us-central1"}
+- {name: display_name, type: String}
+- {name: description, type: String, optional: true, default: ''}
+- {name: unmanaged_container_model, type: google.UnmanagedContainerModel, optional: true}
+- {name: encryption_spec_key_name, type: String, optional: true, default: ''}
+- {name: labels, type: JsonObject, optional: true, default: '{}'}
+- {name: parent_model, type: google.VertexModel, optional: true}
+outputs:
+- {name: model, type: google.VertexModel}
+- {name: gcp_resources, type: String}
+implementation:
+  container:
+    image: gcr.io/ml-pipeline/automl-tables-private:1.0.17
+    command: [python3, -u, -m, launcher]
+    args: [
+      --type, UploadModel,
+      --payload,
+      concat: [
+          '{',
+          '"display_name": "', {inputValue: display_name}, '"',
+          ', "description": "', {inputValue: description}, '"',
+          ', "encryption_spec": {"kms_key_name":"', {inputValue: encryption_spec_key_name}, '"}',
+          ', "labels": ', {inputValue: labels},
+          '}'
+      ],
+      --project, {inputValue: project},
+      --location, {inputValue: location},
+      --gcp_resources, {outputPath: gcp_resources},
+      --executor_input, "{{$}}",
+      {if: {
+        cond: {isPresent: parent_model},
+        then: ["--parent_model_name", "{{$.inputs.artifacts['parent_model'].metadata['resourceName']}}",]
+      }},
+    ]

google_cloud_pipeline_components/_implementation/starry_net/version.py

@@ -0,0 +1,18 @@
+# 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.
+"""Version constants for starry net components."""
+
+DATAPREP_VERSION = '20240722_2225_RC00'
+PREDICTOR_VERSION = '20240723_0542_RC00'
+TRAINER_VERSION = '20240723_0542_RC00'

google_cloud_pipeline_components/container/preview/custom_job/remote_runner.py

@@ -32,6 +32,27 @@ def insert_system_labels_into_payload(payload):
   return json.dumps(job_spec)
 
 
+def cast_accelerator_count_to_int(payload):
+  """Casts accelerator_count from string to an int."""
+
+  job_spec = json.loads(payload)
+  # TODO(b/353577594): accelerator_count placeholder is not resolved to int.
+  # Need to typecast to int to avoid type mismatch error. Can remove when fix
+  # placeholder resolution.
+  if (
+      'accelerator_count'
+      in job_spec['job_spec']['worker_pool_specs'][0]['machine_spec']
+  ):
+    job_spec['job_spec']['worker_pool_specs'][0]['machine_spec'][
+        'accelerator_count'
+    ] = int(
+        job_spec['job_spec']['worker_pool_specs'][0]['machine_spec'][
+            'accelerator_count'
+        ]
+    )
+  return json.dumps(job_spec)
+
+
 def create_custom_job_with_client(job_client, parent, job_spec):
   create_custom_job_fn = None
   try:
@@ -86,6 +107,7 @@ def create_custom_job(
     # Create custom job if it does not exist
     job_name = remote_runner.check_if_job_exists()
     if job_name is None:
+      payload = cast_accelerator_count_to_int(payload)
       job_name = remote_runner.create_job(
           create_custom_job_with_client,
           insert_system_labels_into_payload(payload),

google_cloud_pipeline_components/container/utils/error_surfacing.py

@@ -0,0 +1,45 @@
+# 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.
+
+"""Utilities for surface user defined error messages."""
+
+import json
+import os
+from google.protobuf import json_format
+from google_cloud_pipeline_components.proto import task_error_pb2
+
+
+def write_customized_error(
+    executor_input: str, error: task_error_pb2.TaskError
+):
+  """Writes a TaskError customized by the author of the pipelines to a JSON file ('executor_error.json') in the output directory specified in the executor input.
+
+  Args:
+    executor_input: JSON string containing executor input data.
+    error: TaskError protocol buffer message.
+  """
+  executor_input_json = json.loads(executor_input)
+  os.makedirs(
+      os.path.dirname(executor_input_json['outputs']['outputFile']),
+      exist_ok=True,
+  )
+  executor_out_path = executor_input_json['outputs']['outputFile']
+  directory_path = os.path.dirname(executor_out_path)
+  executor_error_path = os.path.join(directory_path, 'executor_error.json')
+  error_dict = json_format.MessageToDict(error)
+  with open(
+      executor_error_path,
+      'w',
+  ) as f:
+    json.dump(error_dict, f)

google_cloud_pipeline_components/container/v1/model/get_model/remote_runner.py

@@ -12,13 +12,33 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Remote runner for Get Model based on the Vertex AI SDK."""
-
+import contextlib
+from typing import Tuple, Type, Union
 from google.api_core.client_options import ClientOptions
 from google.cloud import aiplatform_v1 as aip_v1
 from google_cloud_pipeline_components.container.utils import artifact_utils
+from google_cloud_pipeline_components.container.utils import error_surfacing
+from google_cloud_pipeline_components.proto import task_error_pb2
 from google_cloud_pipeline_components.types import artifact_types
 
 
+@contextlib.contextmanager
+def catch_write_and_raise(
+    executor_input: str,
+    exception_types: Union[
+        Type[Exception], Tuple[Type[Exception], ...]
+    ] = Exception,
+):
+  """Context manager to catch specified exceptions, log them using error_surfacing, and then re-raise."""
+  try:
+    yield
+  except exception_types as e:
+    task_error = task_error_pb2.TaskError()
+    task_error.error_message = str(e)
+    error_surfacing.write_customized_error(executor_input, task_error)
+    raise
+
+
 def get_model(
     executor_input,
     model_name: str,
@@ -26,11 +46,16 @@ def get_model(
     location: str,
 ) -> None:
   """Get model."""
-  if not location or not project:
-    raise ValueError(
-        'Model resource name must be in the format'
-        ' projects/{project}/locations/{location}/models/{model_name}'
-    )
+  with catch_write_and_raise(
+      executor_input=executor_input,
+      exception_types=ValueError,
+  ):
+    if not location or not project:
+      model_name_error_message = (
+          'Model resource name must be in the format'
+          ' projects/{project}/locations/{location}/models/{model_name}'
+      )
+      raise ValueError(model_name_error_message)
   api_endpoint = location + '-aiplatform.googleapis.com'
   vertex_uri_prefix = f'https://{api_endpoint}/v1/'
   model_resource_name = (
@@ -40,7 +65,11 @@ def get_model(
   client_options = ClientOptions(api_endpoint=api_endpoint)
   client = aip_v1.ModelServiceClient(client_options=client_options)
   request = aip_v1.GetModelRequest(name=model_resource_name)
-  get_model_response = client.get_model(request)
+  with catch_write_and_raise(
+      executor_input=executor_input,
+      exception_types=Exception,
+  ):
+    get_model_response = client.get_model(request)
   resp_model_name_without_version = get_model_response.name.split('@', 1)[0]
   model_resource_name = (
       f'{resp_model_name_without_version}@{get_model_response.version_id}'

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:20240419_0625'
+  image_uri = 'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240710_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:20240419_0625',
+                  'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240710_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:20240419_0625',
+                  'us-docker.pkg.dev/vertex-ai-restricted/automl-tabular/forecasting-training:20240710_0625',
                   '", "--reduce_search_space_mode=',
                   reduce_search_space_mode,
                   f'", "--component_id={dsl.PIPELINE_TASK_ID_PLACEHOLDER}',

google_cloud_pipeline_components/preview/automl/forecasting/forecasting_stage_2_tuner.py

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