diff --git a/components/google-cloud/google_cloud_pipeline_components/aiplatform/__init__.py b/components/google-cloud/google_cloud_pipeline_components/aiplatform/__init__.py index 2dc8ac4b742..663ec9b8c38 100644 --- a/components/google-cloud/google_cloud_pipeline_components/aiplatform/__init__.py +++ b/components/google-cloud/google_cloud_pipeline_components/aiplatform/__init__.py @@ -17,17 +17,36 @@ from google_cloud_pipeline_components.aiplatform import utils __all__ = [ - 'ImageDatasetCreateOp', 'TabularDatasetCreateOp', 'TextDatasetCreateOp', - 'VideoDatasetCreateOp', 'ImageDatasetExportDataOp', - 'TabularDatasetExportDataOp', 'TextDatasetExportDataOp', - 'VideoDatasetExportDataOp', 'ImageDatasetImportDataOp', - 'TextDatasetImportDataOp', 'VideoDatasetImportDataOp', - 'CustomContainerTrainingJobRunOp', 'CustomPythonPackageTrainingJobRunOp', - 'AutoMLImageTrainingJobRunOp', 'AutoMLTextTrainingJobRunOp', - 'AutoMLTabularTrainingJobRunOp', 'AutoMLVideoTrainingJobRunOp', - 'ModelDeployOp', 'ModelBatchPredictOp', 'ModelUploadOp', 'EndpointCreateOp' + 'ImageDatasetCreateOp', + 'TabularDatasetCreateOp', + 'TextDatasetCreateOp', + 'VideoDatasetCreateOp', + 'ImageDatasetExportDataOp', + 'TabularDatasetExportDataOp', + 'TextDatasetExportDataOp', + 'VideoDatasetExportDataOp', + 'ImageDatasetImportDataOp', + 'TextDatasetImportDataOp', + 'VideoDatasetImportDataOp', + 'CustomContainerTrainingJobRunOp', + 'CustomPythonPackageTrainingJobRunOp', + 'AutoMLImageTrainingJobRunOp', + 'AutoMLTextTrainingJobRunOp', + 'AutoMLTabularTrainingJobRunOp', + 'AutoMLVideoTrainingJobRunOp', + 'ModelDeployOp', + 'ModelBatchPredictOp', + 'ModelUploadOp', + 'EndpointCreateOp', + 'TimeSeriesDatasetCreateOp', + 'TimeSeriesDatasetExportDataOp', + 'AutoMLForecastingTrainingJobRunOp', ] +TimeSeriesDatasetCreateOp = utils.convert_method_to_component( + aiplatform_sdk.TimeSeriesDataset, aiplatform_sdk.TimeSeriesDataset.create +) + ImageDatasetCreateOp = utils.convert_method_to_component( aiplatform_sdk.ImageDataset, aiplatform_sdk.ImageDataset.create ) @@ -54,6 +73,11 @@ aiplatform_sdk.TabularDataset.export_data, ) +TimeSeriesDatasetExportDataOp = utils.convert_method_to_component( + aiplatform_sdk.TimeSeriesDataset, + aiplatform_sdk.TimeSeriesDataset.export_data, +) + TextDatasetExportDataOp = utils.convert_method_to_component( aiplatform_sdk.TextDataset, aiplatform_sdk.TextDataset.export_data, @@ -104,6 +128,11 @@ aiplatform_sdk.AutoMLTabularTrainingJob.run, ) +AutoMLForecastingTrainingJobRunOp = utils.convert_method_to_component( + aiplatform_sdk.AutoMLForecastingTrainingJob, + aiplatform_sdk.AutoMLForecastingTrainingJob.run, +) + AutoMLVideoTrainingJobRunOp = utils.convert_method_to_component( aiplatform_sdk.AutoMLVideoTrainingJob, aiplatform_sdk.AutoMLVideoTrainingJob.run,