diff --git a/.pyspelling-ignore-words b/.pyspelling-ignore-words index 3d35dc9b00..9cf0c91d45 100644 --- a/.pyspelling-ignore-words +++ b/.pyspelling-ignore-words @@ -67,6 +67,7 @@ CamelCase ChatOps ClassVar ClickException +CloudSQL Colab ComponentWrapper ComputeTarget @@ -75,6 +76,7 @@ ConnectionError ContainerOp ContainerOps Contextmanager +cryptoKeys Cron CustomFormatter DAGs @@ -108,6 +110,7 @@ Evidently's EvidentlyProfileConfig EvidentlyProfileStep EvidentlyVisualizer +eu ExampleSecretSchema ExamplesRepo ExecutionStatus @@ -177,6 +180,8 @@ KMS KServe Keras KeyError +keyRings +kr Kubeflow KubeflowMetadataStore KubeflowOrchestrator @@ -734,6 +739,7 @@ numericaltargetdrift numpy odbc onerror +orch orchestrator orchestrators os @@ -764,6 +770,7 @@ probabilisticmodelperformance profiler proto protobuf +PWD py pydantic pyenv @@ -805,6 +812,7 @@ seldon sepc serializable serverless +serviceacc serviceAccountName setFormatter setattr diff --git a/docs/book/advanced-guide/execute-pipelines-in-cloud.md b/docs/book/advanced-guide/execute-pipelines-in-cloud.md index 8094456cb0..ebd3d587c4 100644 --- a/docs/book/advanced-guide/execute-pipelines-in-cloud.md +++ b/docs/book/advanced-guide/execute-pipelines-in-cloud.md @@ -316,7 +316,7 @@ new stack with these components that you have just created. zenml container-registry register cloud_registry --flavor= --uri=$PATH_TO_YOUR_CONTAINER_REGISTRY zenml orchestrator register cloud_orchestrator --flavor=kubeflow --custom_docker_base_image_name=YOUR_IMAGE zenml artifact-store register cloud_artifact_store --flavor= --path=$PATH_TO_YOUR_BUCKET - zenml secrets-manager register cloud_secrets_manager --flavor= + zenml secrets-manager register cloud_secrets_manager --flavor= # Register the cloud stack zenml stack register cloud_kubeflow_stack -m cloud_metadata_store -a cloud_artifact_store -o cloud_orchestrator -c cloud_registry -x cloud_secrets_manager diff --git a/docs/book/advanced-guide/integrations.md b/docs/book/advanced-guide/integrations.md index a9cde7cd05..3270d56534 100644 --- a/docs/book/advanced-guide/integrations.md +++ b/docs/book/advanced-guide/integrations.md @@ -35,15 +35,21 @@ These are the third-party integrations that ZenML currently supports: |---------------------------|--------|------------------------|-----------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------| | Apache Airflow | ✅ | Orchestrator | Works for local environment. | [airflow_orchestration](https://github.com/zenml-io/zenml/tree/main/examples/airflow_orchestration) | | Apache Beam | ✅ | Distributed Processing | | | -| AWS | ✅ | Secrets Manager | Use AWS as a secrets manager. | | -| Azure | ✅ | Cloud | Use Azure Blob Storage buckets as ZenML artifact stores and AzureML as a ZenML step operator. | [azureml_step_operator](https://github.com/zenml-io/zenml/tree/main/examples/step_operator_remote_training) | -| BentoML | ⛏ | Cloud | Looking for community implementors. | | +| AWS | ✅ | Container Registry | Use the AWS container registry to store your containers. | | +| AWS | ✅ | Secrets Manager | Use AWS as a secrets manager. | [cloud_secrets_manager](https://github.com/zenml-io/zenml/tree/develop/examples/cloud_secrets_manager) | +| AWS | ✅ | Step Operator | Sagemaker as a ZenML step operator. | [sagemaker_step_operator](https://github.com/zenml-io/zenml/tree/main/examples/step_operator_remote_training) | +| Azure | ✅ | Artifact Store | Use Azure Blob Storage buckets as ZenML artifact stores. | | +| Azure | ✅ | Step Operator | Use AzureML as a step operator to supercharge specific steps. | [azureml_step_operator](https://github.com/zenml-io/zenml/tree/main/examples/step_operator_remote_training) | +| BentoML | ⛏ | Deployment | Looking for community implementors. | | | Dash | ✅ | Visualizer | For Pipeline and PipelineRun visualization objects. | [lineage](https://github.com/zenml-io/zenml/tree/main/examples/lineage) | | Evidently | ✅ | Monitoring | Allows for visualization of drift as well as export of a `Profile` object. | [drift_detection](https://github.com/zenml-io/zenml/tree/release/0.5.7/examples/drift_detection) | | Facets | ✅ | Visualizer | Quickly visualize your datasets using `facets`. | [statistics](https://github.com/zenml-io/zenml/tree/main/examples/statistics) | | Feast | ✅ | Feature Store | Use Feast with Redis for your online features. | [feature_store](https://github.com/zenml-io/zenml/tree/main/examples/feature_store) | -| GCP | ✅ | Cloud | Use GCS buckets as a ZenML artifact store. | | | GitHub | ✅ | Orchestrator | Use GitHub Actions to orchestrate your ZenML pipelines. | [github_actions_orchestration](https://github.com/zenml-io/zenml/tree/main/examples/github_actions_orchestration) | +| GCP | ✅ | Artifact Store | Use GCS buckets as a ZenML artifact store. | | +| GCP | ✅ | Step Secrets Manager | Use the GCP Secret Manager. | [cloud_secrets_manager](https://github.com/zenml-io/zenml/tree/develop/examples/cloud_secrets_manager) | +| GCP | ✅ | Step Operator | Vertex AI as a ZenML step operator. | [vertex_step_operator](https://github.com/zenml-io/zenml/tree/main/examples/step_operator_remote_training) | +| GCP | ✅ | Orchestrator | Execute your ZenML pipelines using Vertex AI Pipelines. | | | Graphviz | ✅ | Visualizer | For Pipeline and PipelineRun visualization objects. | [dag_visualizer](https://github.com/zenml-io/zenml/tree/main/examples/dag_visualizer) | | Great Expectations | ⛏ | Data Validation | Looking for community implementors. | | | Hugging Face | ✅ | Materializer | Use Hugging Face tokenizers, datasets and models. | [huggingface](https://github.com/zenml-io/zenml/tree/main/examples/huggingface) | @@ -58,16 +64,13 @@ These are the third-party integrations that ZenML currently supports: | Plotly | ✅ | Visualizer | For Pipeline and PipelineRun visualization objects. | [lineage](https://github.com/zenml-io/zenml/tree/main/examples/lineage) | | PyTorch | ✅ | Training | | [pytorch](https://github.com/zenml-io/zenml/tree/main/examples/pytorch) | | PyTorch Lightning | ✅ | Training | | | -| Vertex | ✅ | Step Operator | Vertex AI as a ZenML step operator. | [vertex_step_operator](https://github.com/zenml-io/zenml/tree/main/examples/step_operator_remote_training) | -| Vertex | ✅ | Orchestrator | Execute your ZenML pipelines using Vertex AI Pipelines | | S3 | ✅ | Artifact Store | Use S3 buckets as ZenML artifact stores. | [caching chapter](https://docs.zenml.io/v/docs/guides/functional-api/caching) | -| Sagemaker | ✅ | Cloud | Sagemaker as a ZenML step operator. | [sagemaker_step_operator](https://github.com/zenml-io/zenml/tree/main/examples/step_operator_remote_training) | | scikit-learn | ✅ | Training | | [caching chapter](https://docs.zenml.io/v/docs/guides/functional-api/caching) | | scipy | ✅ | Materializer | Use sparse matrices. | | | Seldon Core | ✅ | Deployment | Seldon Core as a model deployer. | [seldon](https://github.com/zenml-io/zenml/tree/main/examples/seldon_deployment) | | Tensorflow | ✅ | Training, Visualizer | TensorBoard support. | [quickstart](https://github.com/zenml-io/zenml/tree/main/examples/quickstart). [kubeflow](https://github.com/zenml-io/zenml/tree/main/examples/kubeflow) | | Weights & Biases Tracking | ✅ | Experiment Tracking | Tracks your pipelines and your training runs. | [wandb_tracking](https://github.com/zenml-io/zenml/tree/main/examples/wandb_tracking) | -| whylogs | ✅ | logging | Integration fully implemented for data logging. | [whylogs](https://github.com/zenml-io/zenml/tree/main/examples/whylogs) | +| whylogs | ✅ | Logging | Integration fully implemented for data logging. | [whylogs](https://github.com/zenml-io/zenml/tree/main/examples/whylogs) | | xgboost | ✅ | Training | Support for `Booster` and `DMatrix` materialization. | [xgboost](https://github.com/zenml-io/zenml/tree/main/examples/xgboost) | ✅ means the integration is already implemented. diff --git a/docs/book/advanced-guide/manage-your-secrets.md b/docs/book/advanced-guide/manage-your-secrets.md index 4c4e2687c0..e0e42a3ebd 100644 --- a/docs/book/advanced-guide/manage-your-secrets.md +++ b/docs/book/advanced-guide/manage-your-secrets.md @@ -305,7 +305,7 @@ then register your secrets manager in the following way: ```shell zenml integration install azure -zenml secrets-manager register AZURE_SECRETS_MANAGER_NAME -f azure_secrets_manager \ +zenml secrets-manager register AZURE_SECRETS_MANAGER_NAME -f azure \ --key_vault_name= # update your default stack, for example diff --git a/docs/book/extending-zenml/orchestrators.md b/docs/book/extending-zenml/orchestrators.md index 2eb417f92c..822daed2ff 100644 --- a/docs/book/extending-zenml/orchestrators.md +++ b/docs/book/extending-zenml/orchestrators.md @@ -79,11 +79,13 @@ Moreover, additional orchestrators can be found in specific `integrations` modules, such as the `AirflowOrchestrator` in the `airflow` integration and the `KubeflowOrchestrator` in the `kubeflow` integration. -| | Flavor | Integration | -|---------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------|-------------| -| [LocalOrchestrator](https://apidocs.zenml.io/latest/api_docs/orchestrators/#zenml.orchestrators.local.local_orchestrator.LocalOrchestrator) | local | `built-in` | -| [AirflowOrchestrator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.airflow.orchestrators.airflow_orchestrator.AirflowOrchestrator) | airflow | airflow | -| [KubeflowOrchestrator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.kubeflow.orchestrators.kubeflow_orchestrator.KubeflowOrchestrator) | kubeflow | kubeflow | +| | Flavor | Integration | +|---------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------|-------------| +| [LocalOrchestrator](https://apidocs.zenml.io/latest/api_docs/orchestrators/#zenml.orchestrators.local.local_orchestrator.LocalOrchestrator) | local | `built-in` | +| [AirflowOrchestrator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.airflow.orchestrators.airflow_orchestrator.AirflowOrchestrator) | airflow | airflow | +| [KubeflowOrchestrator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.kubeflow.orchestrators.kubeflow_orchestrator.KubeflowOrchestrator) | kubeflow | kubeflow | +| [VertexAIOrchestrator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.gcp.orchestrators.vertex_orchestrator.VertexOrchestrator) | vertex | gcp | + If you would like to see the available flavors for artifact stores, you can use the command: diff --git a/docs/book/extending-zenml/secrets-managers.md b/docs/book/extending-zenml/secrets-managers.md index 36b605cf10..398c353b68 100644 --- a/docs/book/extending-zenml/secrets-managers.md +++ b/docs/book/extending-zenml/secrets-managers.md @@ -82,14 +82,15 @@ not intended for production use. For production use cases some more flavors can be found in specific `integrations` modules, such as the `GCPSecretsManager` in the -`gcp_secrets_manager` integration and the `AWSSecretsManager` in the -`aws` integration. - -| | Flavor | Integration | -|-----------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------|---------------------| -| [LocalSecretsManager](https://apidocs.zenml.io/latest/api_docs/secrets_managers/#zenml.secrets_managers.local.local_secrets_manager.LocalSecretsManager) | local | `built-in` | -| [AWSSecretsManager](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.aws.secrets_managers.aws_secrets_manager.AWSSecretsManager) | aws | aws | -| GCPSecretsManager | gcp_secrets_manager | gcp_secrets_manager | +`gcp` integration, the `AWSSecretsManager` in the +`aws` integration and the `AzureSecretsManager` in the `azure` integration. + +| | Flavor | Integration | +|-------------------------------------------------------------------------------------------------------------------------------------------------------------------|--------|-------------| +| [LocalSecretsManager](https://apidocs.zenml.io/latest/api_docs/secrets_managers/#zenml.secrets_managers.local.local_secrets_manager.LocalSecretsManager) | local | `built-in` | +| [AWSSecretsManager](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.aws.secrets_managers.aws_secrets_manager.AWSSecretsManager) | aws | aws | +| [GCPSecretsManager](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.gcp.secrets_managers.gcp_secrets_manager.GCPSecretsManager) | gcp | gcp | +| [AzureSecretsManager](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.azure.secrets_managers.azure_secrets_manager.AzureSecretsManager) | azure | azure | If you would like to see the available flavors for secret managers, you can use the command: diff --git a/docs/book/extending-zenml/step-operators.md b/docs/book/extending-zenml/step-operators.md index 1169b1596b..165a45c612 100644 --- a/docs/book/extending-zenml/step-operators.md +++ b/docs/book/extending-zenml/step-operators.md @@ -72,11 +72,11 @@ and get the complete docstrings, please check the [API docs](https://apidocs.zen You can find step operator implementations for the three big cloud providers in the `azureml`, `sagemaker` and `vertex` integrations. -| | Flavor | Integration | -|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------|-------------| -| [AzureMLStepOperator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.azureml.step_operators.azureml_step_operator.AzureMLStepOperator) | azureml | azureml | -| [SagemakerStepOperator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.sagemaker.step_operators.sagemaker_step_operator.SagemakerStepOperator) | sagemaker | sagemaker | -| [VertexStepOperator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.vertex.step_operators.vertex_step_operator.VertexStepOperator) | vertex | vertex | +| | Flavor | Integration | +|---------------------------------------------------------------------------------------------------------------------------------------------------------------------|-----------|-------------| +| [AzureMLStepOperator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.azure.step_operators.azureml_step_operator.AzureMLStepOperator) | azureml | azure | +| [SagemakerStepOperator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.aws.step_operators.sagemaker_step_operator.SagemakerStepOperator) | sagemaker | aws | +| [VertexStepOperator](https://apidocs.zenml.io/latest/api_docs/integrations/#zenml.integrations.gcp.step_operators.vertex_step_operator.VertexStepOperator) | vertex | gcp | If you would like to see the available flavors for step operators, you can use the command: diff --git a/examples/cloud_secrets_manager/README.md b/examples/cloud_secrets_manager/README.md index 09dc60c38e..60aebf549d 100644 --- a/examples/cloud_secrets_manager/README.md +++ b/examples/cloud_secrets_manager/README.md @@ -26,19 +26,14 @@ you please. In order to run this example, you need to install and initialize ZenML. Within this example You'll be able to choose between using the local yaml based secrets manager, -the [AWS Secrets Manager](https://aws.amazon.com/secrets-manager/) -and the [GCP Secret Manager](https://cloud.google.com/secret-manager).: +the [AWS Secrets Manager](https://aws.amazon.com/secrets-manager/), +the [GCP Secret Manager](https://cloud.google.com/secret-manager) or +the [Azure Key Vault](https://azure.microsoft.com/en-us/services/key-vault/#product-overview) ```shell # install CLI pip install zenml -# either install the aws integrations -zenml integration install aws - -# or alternatively the gcp integration -zenml integration install gcp_secrets_manager - # pull example zenml example pull cloud_secrets_manager cd zenml_examples/cloud_secrets_manager @@ -56,6 +51,8 @@ to make sure everything is set up properly. ```shell +zenml integration install aws + zenml secrets-manager register aws_secrets_manager --flavor=aws zenml stack register secrets_stack -m default -o default -a default -x aws_secrets_manager --set ``` @@ -71,15 +68,30 @@ and get the `project_id` which will need to be specified when you register the secrets manager. ```shell +zenml integration install gcp + zenml secrets-manager register gcp_secrets_manager --flavor=gcp_secrets_manager --project_id=PROJECT_ID zenml stack register secrets_stack -m default -o default -a default -x gcp_secrets_manager --set ``` +### 🥞 Set up your stack for Azure + +To get going with Azure you will need to install and configure the +[Azure CLI](https://docs.microsoft.com/en-us/cli/azure/install-azure-cli) +with the correct credentials to access the Azure secrets manager. + +```shell +zenml integration install azure + +zenml secrets-manager register azure_key_vault --flavor=azure_key_vault --key_vault_name= +zenml stack register secrets_stack -m default -o default -a default -x azure_key_vault --set +``` + ### Or stay on a local stack In case you run into issues with either of the clouds, feel free to use a local -secret manager. Just replace `--flavor=aws`/`--flavor=gcp_secret_manager` with `--flavor=local` to -use a file based version of a secret manager. Be aware that this is not +secret manager. Just replace `--flavor=aws`/`--flavor=gcp_secrets_manager`/`--flavor=azure_key_vault` +with `--flavor=local` to use a file based version of a secret manager. Be aware that this is not a recommended location to store sensitive information. diff --git a/examples/step_operator_remote_training/README.md b/examples/step_operator_remote_training/README.md index 9ccf36e9ea..0b23541e0d 100644 --- a/examples/step_operator_remote_training/README.md +++ b/examples/step_operator_remote_training/README.md @@ -1,22 +1,32 @@ # 🧮 Train models on remote environments -This example shows how you can use the `StepOperator` class to run your training jobs on remote backends. +This example shows how you can use the `StepOperator` class to run your training +jobs on remote backends. -The step operator defers the execution of individual steps in a pipeline to specialized runtime environments that are optimized for Machine Learning workloads. +The step operator defers the execution of individual steps in a pipeline to +specialized runtime environments that are optimized for Machine Learning +workloads. ## 🗺 Overview -Here we train a simple sklearn classifier on the MNIST dataset using one of three step operators: + +Here we train a simple sklearn classifier on the MNIST dataset using one of +three step operators: - AWS Sagemaker - GCP Vertex AI - Microsoft AzureML -Currently, step operators only work with a local orchestrator but support for cloud orchestrators is on the way soon! +Currently, step operators only work with a local orchestrator but support for +cloud orchestrators is on the way soon! # 🖥 Run it locally + ## 👣 Step-by-Step -### 📄 Prerequisites -In order to run this example, you need to install and initialize ZenML and the necessary integrations: + +### 📄 Prerequisites + +In order to run this example, you need to install and initialize ZenML and the +necessary integrations: ```shell # install CLI @@ -33,29 +43,39 @@ cd zenml_examples/step_operator_remote_training zenml init ``` -Each type of step operator has their own pre-requisites. +Each type of step operator has their own prerequisites. -Before running this example, you must set up the individual cloud providers in a certain way. The complete guide can be found in the [docs](https://docs.zenml.io/advanced-guide/cloud/step-operators). +Before running this example, you must set up the individual cloud providers in a +certain way. The complete guide can be found in +the [docs](https://docs.zenml.io/advanced-guide/cloud/step-operators). -Please jump to the section of -the step operator you would like to run on: +Please jump to the section applicable to +the step operator you would like to use: ### 🌿 Sagemaker -Sagemaker offers specialized compute instances to run your training jobs and has a beautiful UI to track and manage your models and logs. You can now use ZenML to submit individual steps to be run on compute instances managed by Amazon Sagemaker. + +Sagemaker offers specialized compute instances to run your training jobs and has +a beautiful UI to track and manage your models and logs. You can now use ZenML +to submit individual steps to be run on compute instances managed by Amazon +Sagemaker. The stack will consist of: -* The **local metadata store** which will track the configuration of your -executions. + +* The **local metadata store** which will track the configuration of your + executions. * The **local orchestrator** which will be executing your pipelines steps. * An **S3 artifact store** which will be responsible for storing the -artifacts of your pipeline. -* The **Sagemaker step operator** which will be utilized to run the training step on Sagemaker. + artifacts of your pipeline. +* The **Sagemaker step operator** which will be utilized to run the training + step on Sagemaker. -To configure resources for the step operators, please follow [this guide](https://docs.zenml.io/advanced-guide/cloud/step-operators) and then proceed with the following steps: +To configure resources for the step operators, please +follow [this guide](https://docs.zenml.io/advanced-guide/cloud/step-operators) +and then proceed with the following steps: ```bash # install ZenML integrations -zenml integration install aws s3 sagemaker +zenml integration install aws s3 zenml artifact-store register s3_store \ --flavor=s3 \ @@ -85,26 +105,29 @@ zenml stack register sagemaker_stack \ ### 🪟 Microsoft AzureML -[AzureML](https://azure.microsoft.com/en-us/services/machine-learning/) -offers specialized compute instances to run your training jobs and -has a beautiful UI to track and manage your models and logs. You can now use -ZenML to submit individual steps to be run on compute instances managed by +[AzureML](https://azure.microsoft.com/en-us/services/machine-learning/) +offers specialized compute instances to run your training jobs and +has a beautiful UI to track and manage your models and logs. You can now use +ZenML to submit individual steps to be run on compute instances managed by AzureML. -The stack will consist of: +The stack will consist of: -* The **local metadata store** which will track the configuration of your -executions. +* The **local metadata store** which will track the configuration of your + executions. * The **local orchestrator** which will be executing your pipelines steps. * An **azure artifact store** which will be responsible for storing the -artifacts of your pipeline. -* The **azureml step operator** which will be utilized to run the training step on Azure. + artifacts of your pipeline. +* The **azureml step operator** which will be utilized to run the training step + on Azure. -To configure resources for the step operators, please follow [this guide](https://docs.zenml.io/advanced-guide/cloud/step-operators) and then proceed with the following steps: +To configure resources for the step operators, please +follow [this guide](https://docs.zenml.io/advanced-guide/cloud/step-operators) +and then proceed with the following steps: ```bash # install ZenML integrations -zenml integration install azure azureml +zenml integration install azure zenml artifact-store register azure_store \ --flavor=azure \ @@ -128,25 +151,31 @@ zenml stack register azureml_stack \ ### 📐 GCP Vertex AI -[Vertex AI](https://cloud.google.com/vertex-ai) offers specialized compute to run -[custom training jobs](https://cloud.google.com/vertex-ai/docs/training/custom-training) -and has a beautiful UI to track and manage your models and logs. You can now use ZenML to submit an individual step to -run on a managed training job managed on Vertex AI. +[Vertex AI](https://cloud.google.com/vertex-ai) offers specialized compute to +run +[custom training jobs](https://cloud.google.com/vertex-ai/docs/training/custom-training) +and has a beautiful UI to track and manage your models and logs. You can now use +ZenML to submit an individual step to +run on a managed training job managed on Vertex AI. The stack will consist of: -* The **local metadata store** which will track the configuration of your -executions. + +* The **local metadata store** which will track the configuration of your + executions. * The **local orchestrator** which will be executing your pipelines steps. * A **GCP Bucket artifact store** which will be responsible for storing the -artifacts of your pipeline. -* The **Vertex AI step operator** which will be utilized to run the training step -on GCP. + artifacts of your pipeline. +* The **Vertex AI step operator** which will be used to run the training + step + on GCP. -To configure resources for the step operators, please follow [this guide](https://docs.zenml.io/advanced-guide/cloud/step-operators) and then proceed with the following steps: +To configure resources for the step operators, please +follow [this guide](https://docs.zenml.io/advanced-guide/cloud/step-operators) +and then proceed with the following steps: ```bash # install ZenML integrations -zenml integration install gcp vertex +zenml integration install gcp zenml artifact-store register gcp_store \ --flavor=gcp \ @@ -174,6 +203,7 @@ zenml stack register vertex_training_stack \ ``` ### ▶️ Run the Code + Now we're ready. Execute: ```shell @@ -181,6 +211,7 @@ python run.py ``` ### 🧽 Clean up + In order to clean up, delete the remaining ZenML references. ```shell @@ -189,7 +220,9 @@ rm -rf zenml_examples # 📜 Learn more -Our docs regarding the step operator integrations can be found [here](https://docs.zenml.io/advanced-guide/step-operators). +Our docs for the step operator integrations can be +found [here](https://docs.zenml.io/advanced-guide/step-operators). -If you want to learn more about step operators in general or about how to build your own step operator in ZenML +If you want to learn more about step operators in general or about how to build +your own step operator in ZenML check out our [docs](https://docs.zenml.io/extending-zenml/step-operator). diff --git a/examples/vertex_ai_orchestration/README.md b/examples/vertex_ai_orchestration/README.md new file mode 100644 index 0000000000..8e13f3a544 --- /dev/null +++ b/examples/vertex_ai_orchestration/README.md @@ -0,0 +1,155 @@ +# 🏃 Run pipelines in production using Vertex AI + +[Vertex AI Pipelines](https://cloud.google.com/vertex-ai/docs/pipelines/introduction) +is a serverless ML workflow tool running on the Google Cloud Platform. It is +an easy way to quickly run your code in a production-ready, repeatable +cloud orchestrator that requires minimal setup without provisioning and paying +for standby compute. + +## 📄 Pre-requisites + +In order to run on Vertex AI you will need to do quite a bit of configuration +within GCP to create all the required resources. In sum you will need to have a +[GCP container registry](https://cloud.google.com/container-registry/docs), a +[GCP bucket](https://cloud.google.com/storage/docs/creating-buckets), +[CloudSQL database](https://cloud.google.com/sql/docs/mysql/create-instance) and +a [GCP Secret Manager](https://cloud.google.com/secret-manager). Additionally, +the [Vertex AI API](https://cloud.google.com/vertex-ai/?hl=en_GB&_ga=2.241201409.-205697788.1651483076) +needs to be enabled. + +On top of having these resources you will also need to do some permission +tuning. + +First, you will need to authenticate yourself using the +[gcloud cli](https://cloud.google.com/sdk/gcloud/reference/auth/login): + +```shell +gcloud auth login +``` + +Then you will need to +[authorize](https://cloud.google.com/container-registry/docs/advanced-authentication) +your local Docker client to have access the GCP container registry. + +```shell +gcloud auth configure-docker +``` + +You will also need to +[create a service account](https://cloud.google.com/iam/docs/creating-managing-service-accounts). + +This service account will need permissions to run Vertex AI jobs, and access +secrets as Admin. Additionally, your user account will need to have permissions to use the service +account. + +![Grant user access to Service Account](assets/serviceacc3.png) + +For your CloudSQL database it is also recommended to enable SSL authentication. +You will need to create a client certificate and download all three certificates. +Save these, you will need them at a later point. + +## 🥞 Create a GCP Kubeflow Pipelines stack + +Once everything is done on the GCP side, we will need to configure a +stack with all of these components. + +* The **artifact store** to store step outputs in a GCP Bucket. +* The **metadata store** to track metadata inside a MySQL database. +* The Docker images that are created to run your pipeline are stored in GCP + **container registry**. +* The **Vertex orchestrator** is responsible for running your ZenML pipeline + in Vertex AI. +* The **secrets manager** contains the secrets to allow access to the metadata + store. + +When running the upcoming commands, make sure to +replace all the with the correct values from your GCP project. + +```bash +# install CLI +pip install zenml + +# install ZenML integrations +zenml integration install gcp + +# pull example +zenml example pull huggingface +cd zenml_examples/huggingface + +# Create a zenml repository +zenml init + +# In order to create the GCP stack components, we'll need to install one +# additional ZenML integration: +zenml integration install gcp + +# The CONTAINER_REGISTRY_URI will have a format like this: eu.gcr.io/xxx/xxx +zenml container-registry register gcp_registry --flavor=gcp --uri= + +# The DB_HOST_IP is the public IP Address of your Database: xx.xx.xxx.xxx +# The DB_PORT is 3306 by default - set this in case this default does not apply +# The DB_NAME is the name of the database that you have created in GCP as the +# metadata store. The `mysql_secret` will be created once the secrets manager +# is created and the stack is active. +zenml metadata-store register gcp_metadata_store --flavor=mysql --host= --port= --database= --secret=mysql_secret + +# The PATH_TO_YOUR_GCP_BUCKET is the path to your GCP bucket: gs://xxx +zenml artifact-store register gcp_artifact_store --flavor=gcp --path= + +# The orchestrator needs the PROJECT_ID and the GCP_LOCATION in which to +# run the Vertex AI pipeline. Additionally you might need to set the +# WORKLOAD_SERVICE_ACCOUNT to the service account you created with secret +# manager access, it will be in the format: xxx@xxx.iam.gserviceaccount.com +zenml orchestrator register vertex_orch --flavor=vertex --project= --location= + +# For the secrets manager, all we'll need it the gcp PROJECT_ID +zenml secrets-manager register gcp_secrets_manager --flavor=gcp_secrets_manager --project_id= + +# Now we're ready to assemble our stack +zenml stack register gcp_vertex_stack -m gcp_metadata_store -a gcp_artifact_store -o vertex_orch -c gcp_registry -x gcp_secrets_manager --set + +# With the stack up and running, we can now supply the credentials for the +# mysql metadata store. The SSL certificates have to be generated and downloaded +# from within the CloudSQL UI +zenml secret register mysql_secret --schema=mysql --user= --password= \ + --ssl_ca=@ \ + --ssl_cert=@ \ + --ssl_key=@ +``` + +Your stack should look something like this when you're done: + +![Vertex Stack](assets/vertex_stack.png) + +### ▶️ Run the pipeline + +Once your stack is fully set up, you should be good to go. + +```bash +python run.py +``` + +That's it! If everything went as planned this pipeline should now be running in +the cloud! You should be able to access the Vertex AI Pipelines UI with the link +returned to the run logs. It will look something like this: + +![Vertex AI UI](assets/vertex_ai_ui.png) + +### 🧽 Clean up + +Once you're done experimenting, you can stop the port forwarding and delete the +example files by calling: + +```bash +zenml stack down --force +rm -rf zenml_examples +``` + +Additionally, you might have to clean up your cloud resources to avoid running +costs for storage of artifacts, containers, metadata or secrets. + +# 📜 Learn more + +If you want to learn more about orchestrators in general or about how to build +your own orchestrators in ZenML +check out our [docs](https://docs.zenml.io/extending-zenml/orchestrator). diff --git a/tests/unit/integrations/vertex/__init__.py b/examples/vertex_ai_orchestration/__init__.py similarity index 100% rename from tests/unit/integrations/vertex/__init__.py rename to examples/vertex_ai_orchestration/__init__.py diff --git a/examples/vertex_ai_orchestration/assets/vertex_ai_ui.png b/examples/vertex_ai_orchestration/assets/vertex_ai_ui.png new file mode 100644 index 0000000000..e399aaf1cc Binary files /dev/null and b/examples/vertex_ai_orchestration/assets/vertex_ai_ui.png differ diff --git a/examples/vertex_ai_orchestration/assets/vertex_stack.png b/examples/vertex_ai_orchestration/assets/vertex_stack.png new file mode 100644 index 0000000000..6a7ad406b7 Binary files /dev/null and b/examples/vertex_ai_orchestration/assets/vertex_stack.png differ diff --git a/examples/vertex_ai_orchestration/pipelines/__init__.py b/examples/vertex_ai_orchestration/pipelines/__init__.py new file mode 100644 index 0000000000..c68d6ff065 --- /dev/null +++ b/examples/vertex_ai_orchestration/pipelines/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) ZenML GmbH 2022. 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: +# +# https://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. +""" +## Vertex AI examples pipeline + +Very simple pipeline to show a pipeline running on Vertex AI +""" +from .vertex_example_pipeline.vertex_example_pipeline import ( + vertex_example_pipeline, +) + +__all__ = ["vertex_example_pipeline"] diff --git a/examples/vertex_ai_orchestration/pipelines/vertex_example_pipeline/vertex_example_pipeline.py b/examples/vertex_ai_orchestration/pipelines/vertex_example_pipeline/vertex_example_pipeline.py new file mode 100644 index 0000000000..70c670a985 --- /dev/null +++ b/examples/vertex_ai_orchestration/pipelines/vertex_example_pipeline/vertex_example_pipeline.py @@ -0,0 +1,23 @@ +# Copyright (c) ZenML GmbH 2022. 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: +# +# https://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 zenml.pipelines import pipeline + + +@pipeline +def vertex_example_pipeline(first_step, second_step, third_step): + # Link all the steps artifacts together + first_num = first_step() + random_num = second_step() + third_step(first_num, random_num) diff --git a/examples/vertex_ai_orchestration/run.py b/examples/vertex_ai_orchestration/run.py new file mode 100644 index 0000000000..2534901d0f --- /dev/null +++ b/examples/vertex_ai_orchestration/run.py @@ -0,0 +1,27 @@ +# Copyright (c) ZenML GmbH 2022. 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: +# +# https://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 pipelines import vertex_example_pipeline +from steps import get_first_num, get_random_int, subtract_numbers + +if __name__ == "__main__": + # Initialize a new pipeline run + p = vertex_example_pipeline( + first_step=get_first_num(), + second_step=get_random_int(), + third_step=subtract_numbers(), + ) + + # Run the new pipeline + p.run() diff --git a/examples/vertex_ai_orchestration/steps/__init__.py b/examples/vertex_ai_orchestration/steps/__init__.py new file mode 100644 index 0000000000..9e7d38ccb9 --- /dev/null +++ b/examples/vertex_ai_orchestration/steps/__init__.py @@ -0,0 +1,23 @@ +# Copyright (c) ZenML GmbH 2022. 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: +# +# https://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. +""" +## Exemplary Steps + +Three very simple steps +""" +from .first_step.first_step import get_first_num +from .second_step.second_step import get_random_int +from .third_step.third_step import subtract_numbers + +__all__ = ["get_first_num", "get_random_int", "subtract_numbers"] diff --git a/examples/vertex_ai_orchestration/steps/first_step/first_step.py b/examples/vertex_ai_orchestration/steps/first_step/first_step.py new file mode 100644 index 0000000000..d61ade7493 --- /dev/null +++ b/examples/vertex_ai_orchestration/steps/first_step/first_step.py @@ -0,0 +1,7 @@ +from zenml.steps import Output, step + + +@step +def get_first_num() -> Output(first_num=int): + """Returns an integer.""" + return 10 diff --git a/examples/vertex_ai_orchestration/steps/second_step/second_step.py b/examples/vertex_ai_orchestration/steps/second_step/second_step.py new file mode 100644 index 0000000000..c83a1adbbe --- /dev/null +++ b/examples/vertex_ai_orchestration/steps/second_step/second_step.py @@ -0,0 +1,9 @@ +import random + +from zenml.steps import Output, step + + +@step(enable_cache=False) +def get_random_int() -> Output(random_num=int): + """Get a random integer between 0 and 10.""" + return random.randint(0, 10) diff --git a/examples/vertex_ai_orchestration/steps/third_step/third_step.py b/examples/vertex_ai_orchestration/steps/third_step/third_step.py new file mode 100644 index 0000000000..62efd61323 --- /dev/null +++ b/examples/vertex_ai_orchestration/steps/third_step/third_step.py @@ -0,0 +1,7 @@ +from zenml.steps import Output, step + + +@step +def subtract_numbers(first_num: int, random_num: int) -> Output(result=int): + """Subtract random_num from first_num.""" + return first_num - random_num diff --git a/examples/xgboost/training_pipeline.py b/examples/xgboost/training_pipeline.py new file mode 100644 index 0000000000..9f9c2ad7bb --- /dev/null +++ b/examples/xgboost/training_pipeline.py @@ -0,0 +1,27 @@ +# Copyright (c) ZenML GmbH 2020. 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: +# +# https://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 zenml.integrations.constants import XGBOOST +from zenml.pipelines import pipeline + + +@pipeline(enable_cache=True, required_integrations=[XGBOOST]) +def xgboost_pipeline( + data_loader, + trainer, + predictor, +): + """Links all the steps together in a pipeline""" + mat_train, mat_test = data_loader() + model = trainer(mat_train) + predictor(model, mat_test) diff --git a/src/zenml/cli/__init__.py b/src/zenml/cli/__init__.py index 7642fccbc6..347ae6faf9 100644 --- a/src/zenml/cli/__init__.py +++ b/src/zenml/cli/__init__.py @@ -359,7 +359,7 @@ `register` command: ```bash -zenml step-operator register STEP_OPERATOR_NAME --type STEP_OPERATOR_FLAVOR [--STEP_OPERATOR_OPTIONS] +zenml step-operator register STEP_OPERATOR_NAME --flavor STEP_OPERATOR_FLAVOR [--STEP_OPERATOR_OPTIONS] ``` If you want the name of the current step operator, use the `get` command: diff --git a/src/zenml/integrations/__init__.py b/src/zenml/integrations/__init__.py index 65f40d2024..a35166ff42 100644 --- a/src/zenml/integrations/__init__.py +++ b/src/zenml/integrations/__init__.py @@ -21,15 +21,11 @@ from zenml.integrations.airflow import AirflowIntegration # noqa from zenml.integrations.aws import AWSIntegration # noqa from zenml.integrations.azure import AzureIntegration # noqa -from zenml.integrations.azureml import AzureMLIntegration # noqa from zenml.integrations.dash import DashIntegration # noqa from zenml.integrations.evidently import EvidentlyIntegration # noqa from zenml.integrations.facets import FacetsIntegration # noqa from zenml.integrations.feast import FeastIntegration # noqa from zenml.integrations.gcp import GcpIntegration # noqa -from zenml.integrations.gcp_secrets_manager import ( # noqa - GcpSecretManagerIntegration, -) from zenml.integrations.github import GitHubIntegration # noqa from zenml.integrations.graphviz import GraphvizIntegration # noqa from zenml.integrations.huggingface import HuggingfaceIntegration # noqa @@ -43,13 +39,11 @@ PytorchLightningIntegration, ) from zenml.integrations.s3 import S3Integration # noqa -from zenml.integrations.sagemaker import SagemakerIntegration # noqa from zenml.integrations.scipy import ScipyIntegration # noqa from zenml.integrations.seldon import SeldonIntegration # noqa from zenml.integrations.sklearn import SklearnIntegration # noqa from zenml.integrations.slack import SlackIntegration # noqa from zenml.integrations.tensorflow import TensorflowIntegration # noqa -from zenml.integrations.vertex import VertexIntegration # noqa from zenml.integrations.wandb import WandbIntegration # noqa from zenml.integrations.whylogs import WhylogsIntegration # noqa from zenml.integrations.xgboost import XgboostIntegration # noqa diff --git a/src/zenml/integrations/aws/__init__.py b/src/zenml/integrations/aws/__init__.py index 0878f58e8f..33ce37a653 100644 --- a/src/zenml/integrations/aws/__init__.py +++ b/src/zenml/integrations/aws/__init__.py @@ -11,8 +11,13 @@ # 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. -"""AWS integration for ZenML.""" +"""Integrates multiple AWS Tools as Stack Components. +The AWS integration provides a way for our users to manage their secrets +through AWS, a way to use the aws container registry. Additionally, the +Sagemaker integration submodule provides a way to run ZenML steps in +Sagemaker. +""" from typing import List from zenml.enums import StackComponentType @@ -22,13 +27,14 @@ AWS_SECRET_MANAGER_FLAVOR = "aws" AWS_CONTAINER_REGISTRY_FLAVOR = "aws" +AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR = "sagemaker" class AWSIntegration(Integration): """Definition of AWS integration for ZenML.""" NAME = AWS - REQUIREMENTS = ["boto3==1.21.21"] + REQUIREMENTS = ["boto3==1.21.21", "sagemaker==2.82.2"] @classmethod def flavors(cls) -> List[FlavorWrapper]: @@ -47,10 +53,18 @@ def flavors(cls) -> List[FlavorWrapper]: ), FlavorWrapper( name=AWS_CONTAINER_REGISTRY_FLAVOR, - source="zenml.integrations.aws.container_registries.AWSContainerRegistry", + source="zenml.integrations.aws.container_registries" + ".AWSContainerRegistry", type=StackComponentType.CONTAINER_REGISTRY, integration=cls.NAME, ), + FlavorWrapper( + name=AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR, + source="zenml.integrations.aws.step_operators" + ".SagemakerStepOperator", + type=StackComponentType.STEP_OPERATOR, + integration=cls.NAME, + ), ] diff --git a/src/zenml/integrations/sagemaker/step_operators/__init__.py b/src/zenml/integrations/aws/step_operators/__init__.py similarity index 88% rename from src/zenml/integrations/sagemaker/step_operators/__init__.py rename to src/zenml/integrations/aws/step_operators/__init__.py index 05cb5f2a6d..8ca5eaf5b8 100644 --- a/src/zenml/integrations/sagemaker/step_operators/__init__.py +++ b/src/zenml/integrations/aws/step_operators/__init__.py @@ -13,6 +13,6 @@ # permissions and limitations under the License. """Initialization of the Sagemaker Step Operator.""" -from zenml.integrations.sagemaker.step_operators.sagemaker_step_operator import ( # noqa +from zenml.integrations.aws.step_operators.sagemaker_step_operator import ( # noqa SagemakerStepOperator, ) diff --git a/src/zenml/integrations/sagemaker/step_operators/sagemaker_step_operator.py b/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py similarity index 97% rename from src/zenml/integrations/sagemaker/step_operators/sagemaker_step_operator.py rename to src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py index 9844e34d5b..11d1b16cc7 100644 --- a/src/zenml/integrations/sagemaker/step_operators/sagemaker_step_operator.py +++ b/src/zenml/integrations/aws/step_operators/sagemaker_step_operator.py @@ -18,7 +18,7 @@ import sagemaker from zenml.enums import StackComponentType -from zenml.integrations.sagemaker import SAGEMAKER_STEP_OPERATOR_FLAVOR +from zenml.integrations.aws import AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR from zenml.repository import Repository from zenml.stack import Stack, StackValidator from zenml.step_operators import BaseStepOperator @@ -54,7 +54,7 @@ class SagemakerStepOperator(BaseStepOperator): experiment_name: Optional[str] = None # Class Configuration - FLAVOR: ClassVar[str] = SAGEMAKER_STEP_OPERATOR_FLAVOR + FLAVOR: ClassVar[str] = AWS_SAGEMAKER_STEP_OPERATOR_FLAVOR @property def validator(self) -> Optional[StackValidator]: diff --git a/src/zenml/integrations/azure/__init__.py b/src/zenml/integrations/azure/__init__.py index 46b545e0de..c7a357496b 100644 --- a/src/zenml/integrations/azure/__init__.py +++ b/src/zenml/integrations/azure/__init__.py @@ -14,8 +14,10 @@ """Initialization of the ZenML Azure integration. The Azure integration submodule provides a way to run ZenML pipelines in a cloud -environment. Specifically, it allows the use of cloud artifact stores, and an -`io` module to handle file operations on Azure Blob Storage. +environment. Specifically, it allows the use of cloud artifact stores, +and an `io` module to handle file operations on Azure Blob Storage. +The Azure Step Operator integration submodule provides a way to run ZenML steps +in AzureML. """ from typing import List @@ -25,7 +27,8 @@ from zenml.zen_stores.models import FlavorWrapper AZURE_ARTIFACT_STORE_FLAVOR = "azure" -AZURE_SECRETS_MANAGER_FLAVOR = "azure_secrets_manager" +AZURE_SECRETS_MANAGER_FLAVOR = "azure_key_vault" +AZUREML_STEP_OPERATOR_FLAVOR = "azureml" class AzureIntegration(Integration): @@ -36,6 +39,7 @@ class AzureIntegration(Integration): "adlfs==2021.10.0", "azure-keyvault-keys", "azure-identity", + "azureml-core==1.39.0.post1", ] @classmethod @@ -48,16 +52,25 @@ def flavors(cls) -> List[FlavorWrapper]: return [ FlavorWrapper( name=AZURE_ARTIFACT_STORE_FLAVOR, - source="zenml.integrations.azure.artifact_stores.AzureArtifactStore", + source="zenml.integrations.azure.artifact_stores" + ".AzureArtifactStore", type=StackComponentType.ARTIFACT_STORE, integration=cls.NAME, ), FlavorWrapper( name=AZURE_SECRETS_MANAGER_FLAVOR, - source="zenml.integrations.azure.secrets_managers.AzureSecretsManager", + source="zenml.integrations.azure.secrets_managers" + ".AzureSecretsManager", type=StackComponentType.SECRETS_MANAGER, integration=cls.NAME, ), + FlavorWrapper( + name=AZUREML_STEP_OPERATOR_FLAVOR, + source="zenml.integrations.azure.step_operators" + ".AzureMLStepOperator", + type=StackComponentType.STEP_OPERATOR, + integration=cls.NAME, + ), ] diff --git a/src/zenml/integrations/azure/secrets_managers/azure_secrets_manager.py b/src/zenml/integrations/azure/secrets_managers/azure_secrets_manager.py index 324d6b3694..c5e4814cbf 100644 --- a/src/zenml/integrations/azure/secrets_managers/azure_secrets_manager.py +++ b/src/zenml/integrations/azure/secrets_managers/azure_secrets_manager.py @@ -16,9 +16,10 @@ from typing import Any, ClassVar, Dict, List from azure.identity import DefaultAzureCredential -from azure.keyvault.secrets import SecretClient # type: ignore [import] +from azure.keyvault.secrets import SecretClient # type: ignore[import] from zenml.exceptions import SecretExistsError +from zenml.integrations.azure import AZURE_SECRETS_MANAGER_FLAVOR from zenml.logger import get_logger from zenml.secret.base_secret import BaseSecretSchema from zenml.secret.secret_schema_class_registry import SecretSchemaClassRegistry @@ -80,7 +81,7 @@ class AzureSecretsManager(BaseSecretsManager): key_vault_name: str # Class configuration - FLAVOR: ClassVar[str] = "azure_secrets_manager" + FLAVOR: ClassVar[str] = AZURE_SECRETS_MANAGER_FLAVOR CLIENT: ClassVar[Any] = None @classmethod diff --git a/src/zenml/integrations/azureml/step_operators/__init__.py b/src/zenml/integrations/azure/step_operators/__init__.py similarity index 89% rename from src/zenml/integrations/azureml/step_operators/__init__.py rename to src/zenml/integrations/azure/step_operators/__init__.py index 06222cf555..15efed991c 100644 --- a/src/zenml/integrations/azureml/step_operators/__init__.py +++ b/src/zenml/integrations/azure/step_operators/__init__.py @@ -13,6 +13,6 @@ # permissions and limitations under the License. """Initialization of AzureML Step Operator integration.""" -from zenml.integrations.azureml.step_operators.azureml_step_operator import ( # noqa +from zenml.integrations.azure.step_operators.azureml_step_operator import ( # noqa AzureMLStepOperator, ) diff --git a/src/zenml/integrations/azureml/step_operators/azureml_step_operator.py b/src/zenml/integrations/azure/step_operators/azureml_step_operator.py similarity index 99% rename from src/zenml/integrations/azureml/step_operators/azureml_step_operator.py rename to src/zenml/integrations/azure/step_operators/azureml_step_operator.py index 911a8886f5..17dabd305d 100644 --- a/src/zenml/integrations/azureml/step_operators/azureml_step_operator.py +++ b/src/zenml/integrations/azure/step_operators/azureml_step_operator.py @@ -32,7 +32,7 @@ from zenml.config.global_config import GlobalConfiguration from zenml.constants import ENV_ZENML_CONFIG_PATH from zenml.environment import Environment as ZenMLEnvironment -from zenml.integrations.azureml import AZUREML_STEP_OPERATOR_FLAVOR +from zenml.integrations.azure import AZUREML_STEP_OPERATOR_FLAVOR from zenml.io import fileio from zenml.step_operators import BaseStepOperator from zenml.utils.docker_utils import CONTAINER_ZENML_CONFIG_DIR diff --git a/src/zenml/integrations/azureml/__init__.py b/src/zenml/integrations/azureml/__init__.py deleted file mode 100644 index 63ad1bcee5..0000000000 --- a/src/zenml/integrations/azureml/__init__.py +++ /dev/null @@ -1,48 +0,0 @@ -# Copyright (c) ZenML GmbH 2022. 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: -# -# https://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. -"""The AzureML integration provides a way to run ZenML steps in AzureML.""" -from typing import List - -from zenml.enums import StackComponentType -from zenml.integrations.constants import AZUREML -from zenml.integrations.integration import Integration -from zenml.zen_stores.models import FlavorWrapper - -AZUREML_STEP_OPERATOR_FLAVOR = "azureml" - - -class AzureMLIntegration(Integration): - """Definition of AzureML integration for ZenML.""" - - NAME = AZUREML - REQUIREMENTS = ["azureml-core==1.39.0.post1"] - - @classmethod - def flavors(cls) -> List[FlavorWrapper]: - """Declare the stack component flavors for the AzureML integration. - - Returns: - List of stack component flavors for this integration. - """ - return [ - FlavorWrapper( - name=AZUREML_STEP_OPERATOR_FLAVOR, - source="zenml.integrations.azureml.step_operators.AzureMLStepOperator", - type=StackComponentType.STEP_OPERATOR, - integration=cls.NAME, - ) - ] - - -AzureMLIntegration.check_installation() diff --git a/src/zenml/integrations/feast/feature_stores/feast_feature_store.py b/src/zenml/integrations/feast/feature_stores/feast_feature_store.py index a26c5d4548..56f95c93bc 100644 --- a/src/zenml/integrations/feast/feature_stores/feast_feature_store.py +++ b/src/zenml/integrations/feast/feature_stores/feast_feature_store.py @@ -179,7 +179,7 @@ def get_registry(self) -> Registry: Returns: The registry. """ - fs = FeatureStore(repo_path=self.feast_repo) + fs: FeatureStore = FeatureStore(repo_path=self.feast_repo) return fs.registry def get_feast_version(self) -> str: diff --git a/src/zenml/integrations/gcp/__init__.py b/src/zenml/integrations/gcp/__init__.py index 0beb9c9c0a..ea7681e595 100644 --- a/src/zenml/integrations/gcp/__init__.py +++ b/src/zenml/integrations/gcp/__init__.py @@ -17,6 +17,12 @@ environment. Specifically, it allows the use of cloud artifact stores, metadata stores, and an `io` module to handle file operations on Google Cloud Storage (GCS). + +Additionally, the GCP secrets manager integration submodule provides a +way to access the GCP secrets manager from within your ZenML Pipeline runs. + +The Vertex AI integration submodule provides a way to run ZenML pipelines in a +Vertex AI environment. """ from typing import List @@ -27,13 +33,21 @@ from zenml.zen_stores.models import FlavorWrapper GCP_ARTIFACT_STORE_FLAVOR = "gcp" +GCP_SECRETS_MANAGER_FLAVOR = "gcp_secrets_manager" +GCP_VERTEX_ORCHESTRATOR_FLAVOR = "vertex" +GCP_VERTEX_STEP_OPERATOR_FLAVOR = "vertex" class GcpIntegration(Integration): """Definition of Google Cloud Platform integration for ZenML.""" NAME = GCP - REQUIREMENTS = ["gcsfs"] + REQUIREMENTS = [ + "kfp", + "gcsfs", + "google-cloud-secret-manager", + "google-cloud-aiplatform>=1.11.0", + ] @classmethod def flavors(cls) -> List[FlavorWrapper]: @@ -45,10 +59,32 @@ def flavors(cls) -> List[FlavorWrapper]: return [ FlavorWrapper( name=GCP_ARTIFACT_STORE_FLAVOR, - source="zenml.integrations.gcp.artifact_stores.GCPArtifactStore", + source="zenml.integrations.gcp.artifact_stores" + ".GCPArtifactStore", type=StackComponentType.ARTIFACT_STORE, integration=cls.NAME, - ) + ), + FlavorWrapper( + name=GCP_SECRETS_MANAGER_FLAVOR, + source="zenml.integrations.gcp.secrets_manager." + "GCPSecretsManager", + type=StackComponentType.SECRETS_MANAGER, + integration=cls.NAME, + ), + FlavorWrapper( + name=GCP_VERTEX_ORCHESTRATOR_FLAVOR, + source="zenml.integrations.gcp.orchestrators" + ".VertexOrchestrator", + type=StackComponentType.ORCHESTRATOR, + integration=cls.NAME, + ), + FlavorWrapper( + name=GCP_VERTEX_STEP_OPERATOR_FLAVOR, + source="zenml.integrations.gcp.step_operators" + ".VertexStepOperator", + type=StackComponentType.STEP_OPERATOR, + integration=cls.NAME, + ), ] diff --git a/src/zenml/integrations/vertex/constants.py b/src/zenml/integrations/gcp/constants.py similarity index 100% rename from src/zenml/integrations/vertex/constants.py rename to src/zenml/integrations/gcp/constants.py diff --git a/src/zenml/integrations/vertex/google_credentials_mixin.py b/src/zenml/integrations/gcp/google_credentials_mixin.py similarity index 100% rename from src/zenml/integrations/vertex/google_credentials_mixin.py rename to src/zenml/integrations/gcp/google_credentials_mixin.py diff --git a/src/zenml/integrations/vertex/orchestrators/__init__.py b/src/zenml/integrations/gcp/orchestrators/__init__.py similarity index 89% rename from src/zenml/integrations/vertex/orchestrators/__init__.py rename to src/zenml/integrations/gcp/orchestrators/__init__.py index 1a7516a479..c3d1f3cc33 100644 --- a/src/zenml/integrations/vertex/orchestrators/__init__.py +++ b/src/zenml/integrations/gcp/orchestrators/__init__.py @@ -13,6 +13,6 @@ # permissions and limitations under the License. """Initialization for the VertexAI orchestrator.""" -from zenml.integrations.vertex.orchestrators.vertex_orchestrator import ( # noqa +from zenml.integrations.gcp.orchestrators.vertex_orchestrator import ( # noqa VertexOrchestrator, ) diff --git a/src/zenml/integrations/vertex/orchestrators/vertex_entrypoint_configuration.py b/src/zenml/integrations/gcp/orchestrators/vertex_entrypoint_configuration.py similarity index 100% rename from src/zenml/integrations/vertex/orchestrators/vertex_entrypoint_configuration.py rename to src/zenml/integrations/gcp/orchestrators/vertex_entrypoint_configuration.py diff --git a/src/zenml/integrations/vertex/orchestrators/vertex_orchestrator.py b/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py similarity index 81% rename from src/zenml/integrations/vertex/orchestrators/vertex_orchestrator.py rename to src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py index 95748d22ad..300e83a477 100644 --- a/src/zenml/integrations/vertex/orchestrators/vertex_orchestrator.py +++ b/src/zenml/integrations/gcp/orchestrators/vertex_orchestrator.py @@ -40,12 +40,14 @@ from kfp.v2.compiler import Compiler as KFPV2Compiler from zenml.enums import StackComponentType -from zenml.integrations.gcp import GCP_ARTIFACT_STORE_FLAVOR -from zenml.integrations.vertex import VERTEX_ORCHESTRATOR_FLAVOR -from zenml.integrations.vertex.google_credentials_mixin import ( +from zenml.integrations.gcp import ( + GCP_ARTIFACT_STORE_FLAVOR, + GCP_VERTEX_ORCHESTRATOR_FLAVOR, +) +from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsMixin, ) -from zenml.integrations.vertex.orchestrators.vertex_entrypoint_configuration import ( +from zenml.integrations.gcp.orchestrators.vertex_entrypoint_configuration import ( VERTEX_JOB_ID_OPTION, VertexEntrypointConfiguration, ) @@ -66,7 +68,6 @@ from zenml.stack import Stack from zenml.steps import BaseStep - logger = get_logger(__name__) @@ -88,33 +89,40 @@ class VertexOrchestrator(BaseOrchestrator, GoogleCredentialsMixin): Attributes: custom_docker_base_image_name: Name of the Docker image that should be used as the base for the image that will be used to execute each of - the steps. If no custom base image is given, a basic image of the active - ZenML version will be used. **Note**: This image needs to have ZenML - installed, otherwise the pipeline execution will fail. For that reason, - might want to extend the ZenML Docker images found here: - https://hub.docker.com/r/zenmldocker/zenml/ + the steps. If no custom base image is given, a basic image of the + active ZenML version will be used. **Note**: This image needs to + have ZenML installed, otherwise the pipeline execution will fail. + For that reason, you might want to extend the ZenML Docker images found + here: https://hub.docker.com/r/zenmldocker/zenml/ project: GCP project name. If `None`, the project will be inferred from the environment. location: Name of GCP region where the pipeline job will be executed. Vertex AI Pipelines is available in the following regions: - https://cloud.google.com/vertex-ai/docs/general/locations#feature-availability - pipeline_root: a Cloud Storage URI that will be used by the Vertex AI Pipelines. + https://cloud.google.com/vertex-ai/docs/general/locations#feature + -availability + pipeline_root: a Cloud Storage URI that will be used by the Vertex AI + Pipelines. If not provided but the artifact store in the stack used to execute - the pipeline is a `zenml.integrations.gcp.artifact_stores.GCPArtifactStore`, + the pipeline is a + `zenml.integrations.gcp.artifact_stores.GCPArtifactStore`, then a subdirectory of the artifact store will be used. - encryption_spec_key_name: The Cloud KMS resource identifier of the customer + encryption_spec_key_name: The Cloud KMS resource identifier of the + customer managed encryption key used to protect the job. Has the form: - `projects/my-project/locations/my-region/keyRings/my-kr/cryptoKeys/my-key`. - The key needs to be in the same region as where the compute resource - is created. - workload_service_account: the service account for workload run-as account. - Users submitting jobs must have act-as permission on this run-as account. + `projects//locations//keyRings//cryptoKeys/` + . The key needs to be in the same region as where the compute + resource is created. + workload_service_account: the service account for workload run-as + account. Users submitting jobs must have act-as permission on this + run-as account. If not provided, the default service account will be used. - network: the full name of the Compute Engine Network to which the job should + network: the full name of the Compute Engine Network to which the job + should be peered. For example, `projects/12345/global/networks/myVPC` If not provided, the job will not be peered with any network. synchronous: If `True`, running a pipeline using this orchestrator will - block until all steps finished running on Vertex AI Pipelines service. + block until all steps finished running on Vertex AI Pipelines + service. """ custom_docker_base_image_name: Optional[str] = None @@ -129,7 +137,7 @@ class VertexOrchestrator(BaseOrchestrator, GoogleCredentialsMixin): _pipeline_root: str - FLAVOR: ClassVar[str] = VERTEX_ORCHESTRATOR_FLAVOR + FLAVOR: ClassVar[str] = GCP_VERTEX_ORCHESTRATOR_FLAVOR @property def validator(self) -> Optional[StackValidator]: @@ -155,8 +163,9 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: container_registry = stack.container_registry if container_registry and container_registry.is_local: return False, ( - f"The Vertex orchestrator does not support local container registries. " - f"You should replace the component '{container_registry.name}' " + f"The Vertex orchestrator does not support local " + f"container registries. You should replace the component '" + f"{container_registry.name}' " f"{container_registry.TYPE.value} to a remote one." ) @@ -166,12 +175,14 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: if not local_path: continue return False, ( - f"The '{stack_comp.name}' {stack_comp.TYPE.value} is a local " - f"stack component. The Vertex AI Pipelines orchestrator requires " - f"that all the components in the stack used to execute the " - f"pipeline have to be not local, because there is no way for " - f"Vertex to connect to your local machine. You should use a " - f"flavor of {stack_comp.TYPE.value} other than '{stack_comp.FLAVOR}'." + f"The '{stack_comp.name}' {stack_comp.TYPE.value} is a " + f"local stack component. The Vertex AI Pipelines " + f"orchestrator requires that all the components in the " + f"stack used to execute the pipeline have to be not local, " + f"because there is no way for Vertex to connect to your " + f"local machine. You should use a flavor of " + f"{stack_comp.TYPE.value} other than '" + f"{stack_comp.FLAVOR}'." ) # If the `pipeline_root` has not been defined in the orchestrator @@ -182,12 +193,14 @@ def _validate_stack_requirements(stack: "Stack") -> Tuple[bool, str]: and stack.artifact_store.FLAVOR != GCP_ARTIFACT_STORE_FLAVOR ): return False, ( - f"The attribute `pipeline_root` has not been set and it cannot " - f"be generated using the path of the artifact store because " - f"it is not a `zenml.integrations.gcp.artifact_store.GCPArtifactStore`. " - f"To solve this issue, set the `pipeline_root` attribute manually " - f"executing the following command: " - f'`zenml orchestrator update {stack.orchestrator.name} --pipeline_root=""`.' + f"The attribute `pipeline_root` has not been set and it " + f"cannot be generated using the path of the artifact store " + f"because it is not a " + f"`zenml.integrations.gcp.artifact_store.GCPArtifactStore`." + f" To solve this issue, set the `pipeline_root` attribute " + f"manually executing the following command: " + f"`zenml orchestrator update {stack.orchestrator.name} " + f'--pipeline_root=""`.' ) return True, "" @@ -217,7 +230,7 @@ def get_docker_image_name(self, pipeline_name: str) -> str: @property def root_directory(self) -> str: - """Returns path to the root directory for all files concerning this orchestrator. + """Returns path to the root directory for files for this orchestrator. Returns: The path to the root directory for all files concerning this @@ -229,7 +242,7 @@ def root_directory(self) -> str: @property def pipeline_directory(self) -> str: - """Returns path to a directory where kubeflow pipelines files are stored. + """Returns path to directory where kubeflow pipelines files are stored. Returns: Path to the pipeline directory. @@ -301,18 +314,19 @@ def prepare_or_run_pipeline( builds a Docker image that contains the code for the pipeline, all steps the context around these files. - Based on this Docker image a callable is created which builds container_ops - for each step (`_construct_kfp_pipeline`). The function `kfp.components.load_component_from_text` - is used to create the `ContainerOp`, because using the `dsl.ContainerOp` - class directly is deprecated when using the Kubeflow SDK v2. The step - entrypoint command with the entrypoint arguments is the command that will - be executed by the container created using the previously created Docker - image. + Based on this Docker image a callable is created which builds + container_ops for each step (`_construct_kfp_pipeline`). The function + `kfp.components.load_component_from_text` is used to create the + `ContainerOp`, because using the `dsl.ContainerOp` class directly is + deprecated when using the Kubeflow SDK v2. The step entrypoint command + with the entrypoint arguments is the command that will be executed by + the container created using the previously created Docker image. - This callable is then compiled into a JSON file that is used as the intermediary - representation of the Kubeflow pipeline. + This callable is then compiled into a JSON file that is used as the + intermediary representation of the Kubeflow pipeline. - This file then is submitted to the Vertex AI Pipelines service for execution. + This file then is submitted to the Vertex AI Pipelines service for + execution. Args: sorted_steps: List of sorted steps. @@ -322,26 +336,31 @@ class directly is deprecated when using the Kubeflow SDK v2. The step runtime_configuration: The Runtime configuration of the current run. Raises: - ValueError: If the attribute `pipeline_root` is not set and it can be - not generated using the path of the artifact store in the stack - because it is not a `zenml.integrations.gcp.artifact_store.GCPArtifactStore`. + ValueError: If the attribute `pipeline_root` is not set and it + can be not generated using the path of the artifact store in the + stack because it is not a + `zenml.integrations.gcp.artifact_store.GCPArtifactStore`. """ - # If the `pipeline_root` has not been defined in the orchestrator configuration, - # try to create it from the artifact store if it is a `GCPArtifactStore`. + # If the `pipeline_root` has not been defined in the orchestrator + # configuration, + # try to create it from the artifact store if it is a + # `GCPArtifactStore`. if not self.pipeline_root: artifact_store = stack.artifact_store self._pipeline_root = f"{artifact_store.path.rstrip('/')}/vertex_pipeline_root/{pipeline.name}/{runtime_configuration.run_name}" logger.info( - "The attribute `pipeline_root` has not been set in the orchestrator " - "configuration. One has been generated automatically based on the " - "path of the `GCPArtifactStore` artifact store in the stack used " - "to execute the pipeline. The generated `pipeline_root` is `%s`.", + "The attribute `pipeline_root` has not been set in the " + "orchestrator configuration. One has been generated " + "automatically based on the path of the `GCPArtifactStore` " + "artifact store in the stack used to execute the pipeline. " + "The generated `pipeline_root` is `%s`.", self._pipeline_root, ) else: self._pipeline_root = self.pipeline_root - # Build the Docker image that will be used to run the steps of the pipeline. + # Build the Docker image that will be used to run the steps of the + # pipeline. image_name = self.get_docker_image_name(pipeline.name) image_name = get_image_digest(image_name) or image_name @@ -354,9 +373,9 @@ def _construct_kfp_pipeline() -> None: Additionally, this gives each `ContainerOp` information about its direct downstream steps. - If this callable is passed to the `compile()` method of `KFPV2Compiler` - all `dsl.ContainerOp` instances will be automatically added to a singular - `dsl.Pipeline` instance. + If this callable is passed to the `compile()` method of + `KFPV2Compiler` all `dsl.ContainerOp` instances will be + automatically added to a singular `dsl.Pipeline` instance. """ step_name_to_container_op: Dict[str, dsl.ContainerOp] = {} @@ -373,7 +392,8 @@ def _construct_kfp_pipeline() -> None: **{VERTEX_JOB_ID_OPTION: dslv2.PIPELINE_JOB_ID_PLACEHOLDER}, ) - # Create the `ContainerOp` for the step. Using the `dsl.ContainerOp` + # Create the `ContainerOp` for the step. Using the + # `dsl.ContainerOp` # class directly is deprecated when using the Kubeflow SDK v2. container_op = kfp.components.load_component_from_text( f""" @@ -408,7 +428,8 @@ def _construct_kfp_pipeline() -> None: # to generate a JSON representation of the pipeline that can be later # upload to Vertex AI Pipelines service. logger.debug( - "Compiling pipeline using Kubeflow SDK V2 compiler and saving it to `%s`", + "Compiling pipeline using Kubeflow SDK V2 compiler and saving it " + "to `%s`", pipeline_file_path, ) KFPV2Compiler().compile( @@ -417,7 +438,8 @@ def _construct_kfp_pipeline() -> None: pipeline_name=_clean_pipeline_name(pipeline.name), ) - # Using the Google Cloud AIPlatform client, upload and execute the pipeline + # Using the Google Cloud AIPlatform client, upload and execute the + # pipeline # on the Vertex AI Pipelines service. self._upload_and_run_pipeline( pipeline_name=pipeline.name, @@ -448,14 +470,17 @@ def _upload_and_run_pipeline( assert runtime_configuration.run_name job_id = _clean_pipeline_name(runtime_configuration.run_name) - # Warn the user that the scheduling is not available using the Vertex Orchestrator + # Warn the user that the scheduling is not available using the Vertex + # Orchestrator if runtime_configuration.schedule: logger.warning( - "Pipeline scheduling configuration was provided, but Vertex AI Pipelines " + "Pipeline scheduling configuration was provided, but Vertex " + "AI Pipelines " "do not have capabilities for scheduling yet." ) - # Get the credentials that would be used to create the Vertex AI Pipelines + # Get the credentials that would be used to create the Vertex AI + # Pipelines # job. credentials, project_id = self._get_authentication() if self.project and self.project != project_id: @@ -486,7 +511,8 @@ def _upload_and_run_pipeline( ) logger.info( - "Submitting pipeline job with job_id `%s` to Vertex AI Pipelines service.", + "Submitting pipeline job with job_id `%s` to Vertex AI Pipelines " + "service.", job_id, ) @@ -494,14 +520,16 @@ def _upload_and_run_pipeline( try: if self.workload_service_account: logger.info( - "The Vertex AI Pipelines job workload will be executed using `%s` " + "The Vertex AI Pipelines job workload will be executed " + "using `%s` " "service account.", self.workload_service_account, ) if self.network: logger.info( - "The Vertex AI Pipelines job will be peered with `%s` network.", + "The Vertex AI Pipelines job will be peered with `%s` " + "network.", self.network, ) diff --git a/src/zenml/integrations/gcp_secrets_manager/secrets_manager/__init__.py b/src/zenml/integrations/gcp/secrets_manager/__init__.py similarity index 90% rename from src/zenml/integrations/gcp_secrets_manager/secrets_manager/__init__.py rename to src/zenml/integrations/gcp/secrets_manager/__init__.py index c2e97ad98b..5f64e9af43 100644 --- a/src/zenml/integrations/gcp_secrets_manager/secrets_manager/__init__.py +++ b/src/zenml/integrations/gcp/secrets_manager/__init__.py @@ -16,7 +16,7 @@ The GCP Secrets Manager allows your pipeline to directly access the GCP secrets manager and use the secrets within during runtime. """ -from zenml.integrations.gcp_secrets_manager.secrets_manager.gcp_secrets_manager import ( +from zenml.integrations.gcp.secrets_manager.gcp_secrets_manager import ( GCPSecretsManager, ) diff --git a/src/zenml/integrations/gcp_secrets_manager/secrets_manager/gcp_secrets_manager.py b/src/zenml/integrations/gcp/secrets_manager/gcp_secrets_manager.py similarity index 99% rename from src/zenml/integrations/gcp_secrets_manager/secrets_manager/gcp_secrets_manager.py rename to src/zenml/integrations/gcp/secrets_manager/gcp_secrets_manager.py index 628644f4a5..0268d891d7 100644 --- a/src/zenml/integrations/gcp_secrets_manager/secrets_manager/gcp_secrets_manager.py +++ b/src/zenml/integrations/gcp/secrets_manager/gcp_secrets_manager.py @@ -18,7 +18,7 @@ from google.cloud import secretmanager from zenml.exceptions import SecretExistsError -from zenml.integrations.gcp_secrets_manager import GCP_SECRETS_MANAGER_FLAVOR +from zenml.integrations.gcp import GCP_SECRETS_MANAGER_FLAVOR from zenml.logger import get_logger from zenml.secret.base_secret import BaseSecretSchema from zenml.secret.secret_schema_class_registry import SecretSchemaClassRegistry diff --git a/src/zenml/integrations/vertex/step_operators/__init__.py b/src/zenml/integrations/gcp/step_operators/__init__.py similarity index 89% rename from src/zenml/integrations/vertex/step_operators/__init__.py rename to src/zenml/integrations/gcp/step_operators/__init__.py index f2f4c94b6d..d2f31ad4ec 100644 --- a/src/zenml/integrations/vertex/step_operators/__init__.py +++ b/src/zenml/integrations/gcp/step_operators/__init__.py @@ -13,6 +13,6 @@ # permissions and limitations under the License. """Initialization for the VertexAI Step Operator.""" -from zenml.integrations.vertex.step_operators.vertex_step_operator import ( # noqa +from zenml.integrations.gcp.step_operators.vertex_step_operator import ( # noqa VertexStepOperator, ) diff --git a/src/zenml/integrations/vertex/step_operators/vertex_step_operator.py b/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py similarity index 98% rename from src/zenml/integrations/vertex/step_operators/vertex_step_operator.py rename to src/zenml/integrations/gcp/step_operators/vertex_step_operator.py index d927ec749e..83a2727155 100644 --- a/src/zenml/integrations/vertex/step_operators/vertex_step_operator.py +++ b/src/zenml/integrations/gcp/step_operators/vertex_step_operator.py @@ -26,15 +26,15 @@ from zenml import __version__ from zenml.enums import StackComponentType -from zenml.integrations.vertex import VERTEX_STEP_OPERATOR_FLAVOR -from zenml.integrations.vertex.constants import ( +from zenml.integrations.gcp import GCP_VERTEX_STEP_OPERATOR_FLAVOR +from zenml.integrations.gcp.constants import ( CONNECTION_ERROR_RETRY_LIMIT, POLLING_INTERVAL_IN_SECONDS, VERTEX_ENDPOINT_SUFFIX, VERTEX_JOB_STATES_COMPLETED, VERTEX_JOB_STATES_FAILED, ) -from zenml.integrations.vertex.google_credentials_mixin import ( +from zenml.integrations.gcp.google_credentials_mixin import ( GoogleCredentialsMixin, ) from zenml.logger import get_logger @@ -77,7 +77,7 @@ class VertexStepOperator(BaseStepOperator, GoogleCredentialsMixin): encryption_spec_key_name: Optional[str] = None # Class configuration - FLAVOR: ClassVar[str] = VERTEX_STEP_OPERATOR_FLAVOR + FLAVOR: ClassVar[str] = GCP_VERTEX_STEP_OPERATOR_FLAVOR @property def validator(self) -> Optional[StackValidator]: diff --git a/src/zenml/integrations/gcp_secrets_manager/__init__.py b/src/zenml/integrations/gcp_secrets_manager/__init__.py deleted file mode 100644 index 8a11b9e315..0000000000 --- a/src/zenml/integrations/gcp_secrets_manager/__init__.py +++ /dev/null @@ -1,53 +0,0 @@ -# Copyright (c) ZenML GmbH 2021. 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: -# -# https://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. -"""Initialization of the GCP Secrets Manager. - -The GCP secrets manager integration submodule provides a way to access the GCP -secrets manager from within your ZenML Pipeline runs. -""" -from typing import List - -from zenml.enums import StackComponentType -from zenml.integrations.constants import GCP_SECRETS_MANAGER -from zenml.integrations.integration import Integration -from zenml.zen_stores.models import FlavorWrapper - -GCP_SECRETS_MANAGER_FLAVOR = "gcp_secrets_manager" - - -class GcpSecretManagerIntegration(Integration): - """Definition of the Secrets Manager for the GCP integration.""" - - NAME = GCP_SECRETS_MANAGER - REQUIREMENTS = ["google-cloud-secret-manager"] - - @classmethod - def flavors(cls) -> List[FlavorWrapper]: - """Declare the stack component flavors for the GCP integration. - - Returns: - List of stack component flavors for this integration. - """ - return [ - FlavorWrapper( - name=GCP_SECRETS_MANAGER_FLAVOR, - source="zenml.integrations.gcp_secrets_manager.secrets_manager." - "GCPSecretsManager", - type=StackComponentType.SECRETS_MANAGER, - integration=cls.NAME, - ) - ] - - -GcpSecretManagerIntegration.check_installation() diff --git a/src/zenml/integrations/sagemaker/__init__.py b/src/zenml/integrations/sagemaker/__init__.py deleted file mode 100644 index 730c57a7d6..0000000000 --- a/src/zenml/integrations/sagemaker/__init__.py +++ /dev/null @@ -1,54 +0,0 @@ -# Copyright (c) ZenML GmbH 2022. 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: -# -# https://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. -"""Initialization of the Sagemaker integration. - -The Sagemaker integration submodule provides a way to run ZenML steps in -Sagemaker. -""" - - -from typing import List - -from zenml.enums import StackComponentType -from zenml.integrations.constants import SAGEMAKER -from zenml.integrations.integration import Integration -from zenml.zen_stores.models import FlavorWrapper - -SAGEMAKER_STEP_OPERATOR_FLAVOR = "sagemaker" - - -class SagemakerIntegration(Integration): - """Definition of Sagemaker integration for ZenML.""" - - NAME = SAGEMAKER - REQUIREMENTS = ["sagemaker==2.82.2"] - - @classmethod - def flavors(cls) -> List[FlavorWrapper]: - """Declare the stack component flavors for the Sagemaker integration. - - Returns: - List of stack component flavors for this integration. - """ - return [ - FlavorWrapper( - name=SAGEMAKER_STEP_OPERATOR_FLAVOR, - source="zenml.integrations.sagemaker.step_operators.SagemakerStepOperator", - type=StackComponentType.STEP_OPERATOR, - integration=cls.NAME, - ) - ] - - -SagemakerIntegration.check_installation() diff --git a/src/zenml/integrations/vertex/__init__.py b/src/zenml/integrations/vertex/__init__.py deleted file mode 100644 index 30e258af7f..0000000000 --- a/src/zenml/integrations/vertex/__init__.py +++ /dev/null @@ -1,60 +0,0 @@ -# Copyright (c) ZenML GmbH 2021. 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: -# -# https://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. -"""Initialization for the VertexAI integration. - -The Vertex integration submodule provides a way to run ZenML pipelines in -Vertex AI environment. -""" - -from typing import List - -from zenml.enums import StackComponentType -from zenml.integrations.constants import VERTEX -from zenml.integrations.integration import Integration -from zenml.zen_stores.models import FlavorWrapper - -VERTEX_ORCHESTRATOR_FLAVOR = "vertex" -VERTEX_STEP_OPERATOR_FLAVOR = "vertex" - - -class VertexIntegration(Integration): - """Definition of Vertex AI integration for ZenML.""" - - NAME = VERTEX - REQUIREMENTS = ["google-cloud-aiplatform>=1.11.0", "kfp==1.8.9"] - - @classmethod - def flavors(cls) -> List[FlavorWrapper]: - """Declare the stack component flavors for the Vertex integration. - - Returns: - List of stack component flavors for this integration. - """ - return [ - FlavorWrapper( - name=VERTEX_ORCHESTRATOR_FLAVOR, - source="zenml.integrations.vertex.orchestrators.VertexOrchestrator", - type=StackComponentType.ORCHESTRATOR, - integration=cls.NAME, - ), - FlavorWrapper( - name=VERTEX_STEP_OPERATOR_FLAVOR, - source="zenml.integrations.vertex.step_operators.VertexStepOperator", - type=StackComponentType.STEP_OPERATOR, - integration=cls.NAME, - ), - ] - - -VertexIntegration.check_installation() diff --git a/tests/unit/integrations/vertex/orchestrators/__init__.py b/tests/unit/integrations/gcp/orchestrators/__init__.py similarity index 100% rename from tests/unit/integrations/vertex/orchestrators/__init__.py rename to tests/unit/integrations/gcp/orchestrators/__init__.py diff --git a/tests/unit/integrations/vertex/orchestrators/test_vertex_orchestrator.py b/tests/unit/integrations/gcp/orchestrators/test_vertex_orchestrator.py similarity index 97% rename from tests/unit/integrations/vertex/orchestrators/test_vertex_orchestrator.py rename to tests/unit/integrations/gcp/orchestrators/test_vertex_orchestrator.py index 979338ee5c..62db5b4e16 100644 --- a/tests/unit/integrations/vertex/orchestrators/test_vertex_orchestrator.py +++ b/tests/unit/integrations/gcp/orchestrators/test_vertex_orchestrator.py @@ -25,7 +25,7 @@ from zenml.exceptions import StackValidationError from zenml.integrations.azure.artifact_stores import AzureArtifactStore from zenml.integrations.gcp.artifact_stores import GCPArtifactStore -from zenml.integrations.vertex.orchestrators import VertexOrchestrator +from zenml.integrations.gcp.orchestrators import VertexOrchestrator from zenml.metadata_stores import MySQLMetadataStore, SQLiteMetadataStore from zenml.stack import Stack @@ -57,7 +57,7 @@ def test_vertex_orchestrator_stack_validation() -> None: ) local_metadata_store = SQLiteMetadataStore(name="", uri="./metadata.db") - local_artifact_store = LocalArtifactStore(name="", path=".") + local_artifact_store = LocalArtifactStore(name="", path="") mysql_metadata_store = MySQLMetadataStore( name="mysql_metadata_store", username="zenml", diff --git a/tests/unit/integrations/gcp_secrets_manager/__init__.py b/tests/unit/integrations/gcp/secrets_manager/__init__.py similarity index 100% rename from tests/unit/integrations/gcp_secrets_manager/__init__.py rename to tests/unit/integrations/gcp/secrets_manager/__init__.py diff --git a/tests/unit/integrations/gcp_secrets_manager/secrets_manager/test_gcp_secrets_manager.py b/tests/unit/integrations/gcp/secrets_manager/test_gcp_secrets_manager.py similarity index 97% rename from tests/unit/integrations/gcp_secrets_manager/secrets_manager/test_gcp_secrets_manager.py rename to tests/unit/integrations/gcp/secrets_manager/test_gcp_secrets_manager.py index ae41befc30..4f22937f78 100644 --- a/tests/unit/integrations/gcp_secrets_manager/secrets_manager/test_gcp_secrets_manager.py +++ b/tests/unit/integrations/gcp/secrets_manager/test_gcp_secrets_manager.py @@ -3,7 +3,7 @@ import pytest from pydantic import BaseModel -from zenml.integrations.gcp_secrets_manager.secrets_manager.gcp_secrets_manager import ( +from zenml.integrations.gcp.secrets_manager.gcp_secrets_manager import ( prepend_group_name_to_keys, remove_group_name_from_key, ) diff --git a/tests/unit/integrations/gcp_secrets_manager/secrets_manager/__init__.py b/tests/unit/integrations/gcp_secrets_manager/secrets_manager/__init__.py deleted file mode 100644 index e69de29bb2..0000000000