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

google_cloud_pipeline_components/_implementation/starry_net/set_eval_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.

google_cloud_pipeline_components/_implementation/starry_net/set_eval_args/component.py

@@ -0,0 +1,76 @@
+# 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 Set Eval Args Component."""
+
+from typing import List, NamedTuple
+
+from kfp import dsl
+
+
+@dsl.component
+def set_eval_args(
+    big_query_source: dsl.Input[dsl.Artifact], quantiles: List[float]
+) -> NamedTuple(
+    'EvalArgs',
+    big_query_source=str,
+    forecasting_type=str,
+    quantiles=List[float],
+    prediction_score_column=str,
+):
+  # fmt: off
+  """Creates Evaluation args.
+
+  Args:
+    big_query_source: The BQ Table containing the test set.
+    quantiles: The quantiles the model was trained to output.
+
+  Returns:
+    A NamedTuple containing big_query_source as a string, forecasting_type
+    used for evaluation step, quantiles in the format expected by the evaluation
+    job, and the prediction_score_column used to evaluate.
+  """
+  outputs = NamedTuple(
+      'EvalArgs',
+      big_query_source=str,
+      forecasting_type=str,
+      quantiles=List[float],
+      prediction_score_column=str)
+
+  def set_forecasting_type_for_eval(quantiles: List[float]) -> str:
+    if quantiles and quantiles[-1] != 0.5:
+      return 'quantile'
+    return 'point'
+
+  def set_quantiles_for_eval(quantiles: List[float]) -> List[float]:
+    updated_q = [q for q in quantiles if q != 0.5]
+    if updated_q:
+      updated_q = [0.5] + updated_q
+    return updated_q
+
+  def set_prediction_score_column(
+      quantiles: List[float]) -> str:
+    updated_q = [q for q in quantiles if q != 0.5]
+    if updated_q:
+      return 'predicted_x.quantile_predictions'
+    return 'predicted_x.value'
+
+  project_id = big_query_source.metadata['projectId']
+  dataset_id = big_query_source.metadata['datasetId']
+  table_id = big_query_source.metadata['tableId']
+  return outputs(
+      f'bq://{project_id}.{dataset_id}.{table_id}',  # pylint: disable=too-many-function-args
+      set_forecasting_type_for_eval(quantiles),  # pylint: disable=too-many-function-args
+      set_quantiles_for_eval(quantiles),  # pylint: disable=too-many-function-args
+      set_prediction_score_column(quantiles),  # pylint: disable=too-many-function-args
+  )

google_cloud_pipeline_components/_implementation/starry_net/set_test_set/__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/set_test_set/component.py

@@ -0,0 +1,48 @@
+# 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 Set Test Set Component."""
+
+from typing import NamedTuple
+
+from kfp import dsl
+
+
+@dsl.component(packages_to_install=['tensorflow==2.11.0'])
+def set_test_set(
+    dataprep_dir: dsl.InputPath(),
+) -> NamedTuple('TestSetArtifact', uri=str, artifact=dsl.Artifact):
+  # fmt: off
+  """Creates test set artifact.
+
+  Args:
+    dataprep_dir: The bucket where dataprep artifacts are stored.
+
+  Returns:
+    The test set dsl.Artifact.
+  """
+  import os  # pylint: disable=g-import-not-at-top
+  import json  # 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(dataprep_dir, 'big_query_test_set.json')
+  ) as f:
+    metadata = json.load(f)
+  project = metadata['projectId']
+  dataset = metadata['datasetId']
+  table = metadata['tableId']
+  output = NamedTuple('TestSetArtifact', uri=str, artifact=dsl.Artifact)
+  uri = f'bq://{project}.{dataset}.{table}'
+  artifact = dsl.Artifact(uri=uri, metadata=metadata)
+  return output(uri, artifact)  # pylint: disable=too-many-function-args

google_cloud_pipeline_components/_implementation/starry_net/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.

google_cloud_pipeline_components/_implementation/starry_net/set_tfrecord_args/component.py

@@ -0,0 +1,70 @@
+# 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 to set TFRecord args."""
+
+from typing import List, NamedTuple
+
+from kfp import dsl
+
+
+@dsl.component
+def set_tfrecord_args(
+    dataprep_dir: dsl.InputPath(),
+    static_covariates: List[str],
+) -> NamedTuple(
+    'TfrecordArgs',
+    static_covariates_vocab_path=str,
+    train_tf_record_patterns=str,
+    val_tf_record_patterns=str,
+    test_tf_record_patterns=str,
+):
+  # fmt: off
+  """Creates Trainer TFRecord args.
+
+  Args:
+    dataprep_dir: The dataprep directory where dataprep artifacts are stored.
+    static_covariates: The static covariates to train the model with.
+
+  Returns:
+    A NamedTuple containing the path to the static covariates covabulary, and
+    the tf record patterns for the train, validation, and test sets.
+  """
+
+  outputs = NamedTuple(
+      'TfrecordArgs',
+      static_covariates_vocab_path=str,
+      train_tf_record_patterns=str,
+      val_tf_record_patterns=str,
+      test_tf_record_patterns=str,
+  )
+
+  if static_covariates and dataprep_dir:
+    static_covariates_vocab_path = f'{dataprep_dir}/static_covariate_vocab.json'
+  else:
+    static_covariates_vocab_path = ''
+  if dataprep_dir:
+    train_tf_record_patterns = f"('{dataprep_dir}/tf_records/train*',)"
+    val_tf_record_patterns = f"('{dataprep_dir}/tf_records/val*',)"
+    test_tf_record_patterns = (
+        f"('{dataprep_dir}/tf_records/test_path_for_plot*',)")
+  else:
+    train_tf_record_patterns = '()'
+    val_tf_record_patterns = '()'
+    test_tf_record_patterns = '()'
+  return outputs(
+      static_covariates_vocab_path,  # pylint: disable=too-many-function-args
+      train_tf_record_patterns,  # pylint: disable=too-many-function-args
+      val_tf_record_patterns,  # pylint: disable=too-many-function-args
+      test_tf_record_patterns,  # pylint: disable=too-many-function-args
+  )

google_cloud_pipeline_components/_implementation/starry_net/set_train_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.

google_cloud_pipeline_components/_implementation/starry_net/set_train_args/component.py

@@ -0,0 +1,90 @@
+# 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 to set training args."""
+
+from typing import List, NamedTuple
+
+from kfp import dsl
+
+
+@dsl.component
+def set_train_args(
+    quantiles: List[float],
+    model_blocks: List[str],
+    static_covariates: List[str],
+) -> NamedTuple(
+    'TrainArgs',
+    quantiles=str,
+    use_static_covariates=bool,
+    static_covariate_names=str,
+    model_blocks=str,
+    freeze_point_forecasts=bool,
+):
+  # fmt: off
+  """Creates Trainer model args.
+
+  Args:
+    quantiles: The list of floats representing quantiles. Leave blank if
+      only training to produce point forecasts.
+    model_blocks: The list of model blocks to use in the order they will appear
+      in the model. Possible values are `cleaning`, `change_point`, `trend`,
+      `hour_of_week`, `day_of_week`, `day_of_year`, `week_of_year`,
+      `month_of_year`, `residual`.
+    static_covariates: The list of strings of static covariate names.
+
+  Returns:
+    A NamedTuple containing the quantiles formatted as expected by the train
+    job, a bool indicating whether the job should train with static covariates,
+    the model blocks formatted as expected by the train job, and a bool
+    indicating whether or not to do two-pass training, fist training for point
+    forecsats and then quantiles.
+  """
+
+  outputs = NamedTuple(
+      'TrainArgs',
+      quantiles=str,
+      use_static_covariates=bool,
+      static_covariate_names=str,
+      model_blocks=str,
+      freeze_point_forecasts=bool,
+  )
+
+  def set_quantiles(input_list: List[float]) -> str:
+    if not input_list or input_list[0] != 0.5:
+      input_list = [0.5] + input_list
+    if len(input_list) == 1:
+      return str(input_list).replace('[', '(').replace(']', ',)')
+    return str(input_list).replace('[', '(').replace(']', ')')
+
+  def maybe_update_model_blocks(
+      quantiles: List[float], model_blocks: List[str]) -> List[str]:
+    updated_q = [q for q in quantiles if q != 0.5]
+    model_blocks = [b for b in model_blocks if b != 'quantile']
+    if updated_q:
+      model_blocks.append('quantile')
+    return [f'{b}-hybrid' if '_of_' in b else b for b in model_blocks]
+
+  def create_name_tuple_from_list(input_list: List[str]) -> str:
+    if len(input_list) == 1:
+      return str(input_list).replace('[', '(').replace(']', ',)')
+    return str(input_list).replace('[', '(').replace(']', ')')
+
+  return outputs(
+      set_quantiles(quantiles),  # pylint: disable=too-many-function-args
+      True if static_covariates else False,  # pylint: disable=too-many-function-args
+      create_name_tuple_from_list(static_covariates),  # pylint: disable=too-many-function-args
+      create_name_tuple_from_list(  # pylint: disable=too-many-function-args
+          maybe_update_model_blocks(quantiles, model_blocks)),
+      True if quantiles and quantiles[-1] != 0.5 else False,  # pylint: disable=too-many-function-args
+  )

google_cloud_pipeline_components/_implementation/starry_net/train/__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/train/component.py

@@ -0,0 +1,209 @@
+# 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,
+    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.
+    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.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,59 @@
+# 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]'])
+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.