Bug fixes and improvements (#1)

* Bug fixes and improvements * Remove binary * Fix the nested fields for the stack component and stack data resource * Sort components in the stack data resource output * Add better error messages regarding missing or expired credentials * Add docs references to terraform modules * Add retry mechanism for service connector validation * Fully migrate to SDK v2 and properly handle 404 HTTP errors * Mark updating component type/flavor as a re-create op * Updated and tested all complete examples
zenml-io · Nov 11, 2024 · 1870e8a · 1870e8a
1 parent 2d0ebfd
commit 1870e8a
Show file tree

Hide file tree

Showing 30 changed files with 1,711 additions and 645 deletions.
diff --git a/README.md b/README.md
@@ -44,6 +44,8 @@ terraform {
 
 ### Authentication
 
+#### Service Account API Key
+
 Configure the provider with your ZenML server URL and API key:
 
 ```hcl
@@ -78,8 +80,30 @@ zenml service-account create <MYSERVICEACCOUNTNAME>
 
 This command will print out the ZENML_API_KEY that you can use with this provider.
 
+#### API Token
+
+Alternatively, you can use an API token for authentication:
+
+```hcl
+provider "zenml" {
+  server_url = "https://your-zenml-server.com"
+  api_token  = "your-api-token"
+}
+```
+
+You can also use environment variables:
+```bash
+export ZENML_SERVER_URL="https://your-zenml-server.com"
+export ZENML_API_TOKEN="your-api-token"
+```
+
 ### Example Usage
 
+> **Hint:** The ZenML Terraform provider is being heavily used in all our Terraform modules. Their code is available on GitHub and can be used as a reference:
+> - [zenml-stack/aws](https://github.com/zenml-io/terraform-aws-zenml-stack)
+> - [zenml-stack/gcp](https://github.com/zenml-io/terraform-gcp-zenml-stack)
+> - [zenml-stack/azure](https://github.com/zenml-io/terraform-azure-zenml-stack)
+
 Here's a basic example of creating a stack with components:
 
 ```hcl
@@ -90,18 +114,9 @@ resource "zenml_service_connector" "gcp" {
   auth_method = "service-account"
   # workspace defaults to "default" if not specified
   
-  resource_types = [
-    "artifact-store",
-    "container-registry",
-    "orchestrator"
-  ]
-  
   configuration = {
     project_id = "my-project"
     location   = "us-central1"
-  }
-  
-  secrets = {
     service_account_json = file("service-account.json")
   }
   

diff --git a/docs/data-sources/server.md b/docs/data-sources/server.md
@@ -0,0 +1,50 @@
+---
+page_title: "zenml_server Data Source - terraform-provider-zenml"
+subcategory: ""
+description: |-
+  Data source for retrieving information about the ZenML server.
+---
+
+# zenml_server (Data Source)
+
+Use this data source to retrieve information about the ZenML server.
+
+## Example Usage
+
+```hcl
+data "zenml_server" "server_info" {
+}
+
+output "version" {
+  value = data.zenml_server.server_info.version
+}
+
+output "dashboard_url" {
+  value = data.zenml_server.server_info.dashboard_url
+}
+```
+
+## Argument Reference
+
+The `zenml_server` data source does not take any arguments.
+
+## Attributes Reference
+
+The following attributes are exported:
+
+* `id` - The ID of the server.
+* `name` - The name of the server.
+* `version` - The version of the server.
+* `deployment_type` - The deployment type of the server.
+* `auth_scheme` - The authentication scheme of the server.
+* `server_url` - The URL where the server API is hosted.
+* `dashboard_url` - The URL where the server's dashboard is hosted.
+* `metadata` - A map of metadata associated with the server.
+
+## Import
+
+The server info can be imported using the `id`, e.g.
+
+```
+$ terraform import zenml_server.server_info 12345678-1234-1234-1234-123456789012
+```
diff --git a/docs/data-sources/service_connector.md b/docs/data-sources/service_connector.md
@@ -38,7 +38,8 @@ In addition to all arguments above, the following attributes are exported:
 * `name` - The name of the service connector.
 * `type` - The type of the service connector (e.g., "gcp", "aws", "azure", etc.).
 * `auth_method` - The authentication method used by the service connector.
-* `resource_types` - A list of resource types supported by this service connector.
+* `resource_type` - The type of resource the service connector is connected to (e.g., "s3-bucket", "docker-registry", etc.).
+* `resource_id` - The ID of the resource the service connector is connected to.
 * `configuration` - A map of configuration key-value pairs for the service connector. Sensitive values are not included.
 * `workspace` - The workspace ID this service connector belongs to.
 * `labels` - A map of labels associated with this service connector.

diff --git a/docs/data-sources/stack_component.md b/docs/data-sources/stack_component.md
@@ -39,6 +39,9 @@ In addition to all arguments above, the following attributes are exported:
 * `flavor` - The flavor of the stack component (e.g., "local", "gcp", "aws", etc.).
 * `configuration` - A map of configuration key-value pairs for the stack component.
 * `workspace` - The workspace ID this stack component belongs to.
+* `labels` - A map of labels associated with this stack component.
+* `connector` - The ID of the service connector associated with this stack component.
+* `connector_resource_id` - The ID of the resource the service connector is connected to.
 
 ## Import
 

diff --git a/docs/index.md b/docs/index.md
@@ -113,10 +113,23 @@ zenml service-account create <MYSERVICEACCOUNTNAME>
 
 This command will print out the ZENML_API_KEY that you can use with this provider.
 
+#### API Token
+
+Alternatively, you can use an API token for authentication, but this is not recommended for production use because API
+tokens can expire:
+
+```hcl
+provider "zenml" {
+  server_url = "https://your-zenml-server.com"
+  api_token  = "your-api-token"
+}
+```
+
 ## Provider Arguments
 
 * `server_url` - (Optional) The URL of your ZenML server. Can be set with the `ZENML_SERVER_URL` environment variable.
 * `api_key` - (Optional) Your ZenML API key. Can be set with the `ZENML_API_KEY` environment variable.
+* `api_token` - (Optional) Your ZenML API token. Can be set with the `ZENML_API_TOKEN` environment variable.
 
 ## Resources
 
@@ -126,6 +139,7 @@ This command will print out the ZENML_API_KEY that you can use with this provide
 
 ## Data Sources
 
+* [zenml_server](data-sources/server.md) - Retrieve information about the ZenML server
 * [zenml_service_connector](data-sources/service_connector.md) - Retrieve information about a service connector
 * [zenml_stack_component](data-sources/stack_component.md) - Retrieve information about a stack component
 * [zenml_stack](data-sources/stack.md) - Retrieve information about a stack
diff --git a/examples/complete-aws/main.tf b/examples/complete-aws/main.tf
@@ -1,12 +1,12 @@
-# examples/complete-aws/main.tf
 terraform {
   required_providers {
+    aws = {
+      source  = "hashicorp/aws"
+      version = "~> 4.0"
+    }
     zenml = {
       source = "zenml-io/zenml"
     }
-    aws = {
-      source = "hashicorp/aws"
-    }
   }
 }
 
@@ -19,61 +19,150 @@ provider "aws" {
   region = var.region
 }
 
-# Create AWS resources if needed
+data "aws_caller_identity" "current" {}
+data "aws_region" "current" {}
+
+# Create S3 bucket for ZenML artifacts
 resource "aws_s3_bucket" "artifacts" {
-  bucket = "${var.name_prefix}-zenml-artifacts-${var.environment}"
+  bucket = "${data.aws_caller_identity.current.account_id}-zenml-artifacts-${var.environment}"
 }
 
-resource "aws_s3_bucket_versioning" "artifacts" {
-  bucket = aws_s3_bucket.artifacts.id
-  versioning_configuration {
-    status = "Enabled"
-  }
-}
+# Create ECR repository for ZenML containers
 
 resource "aws_ecr_repository" "containers" {
-  name = "${var.name_prefix}-zenml-containers-${var.environment}"
+  name = "zenml-containers-${var.environment}"
 }
 
-# IAM Role for ZenML
-resource "aws_iam_role" "zenml" {
-  name = "${var.name_prefix}-zenml-role-${var.environment}"
+# Create IAM user and role with required permissions and keys
 
+resource "aws_iam_user" "iam_user" {
+  name = "zenml-${var.environment}"
+}
+
+resource "aws_iam_user_policy" "assume_role_policy" {
+  name = "AssumeRole"
+  user = aws_iam_user.iam_user.name
+
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Effect = "Allow"
+        Action = "sts:AssumeRole"
+        Resource = "*"
+      }
+    ]
+  })
+}
+
+resource "aws_iam_access_key" "iam_user_access_key" {
+  user = aws_iam_user.iam_user.name
+}
+
+resource "aws_iam_role" "zenml" {
+  name               = "zenml-${var.environment}"
   assume_role_policy = jsonencode({
     Version = "2012-10-17"
     Statement = [
       {
+        Effect = "Allow"
+        Principal = {
+          AWS = aws_iam_user.iam_user.arn
+        }
         Action = "sts:AssumeRole"
+      },
+      {
         Effect = "Allow"
         Principal = {
-          Service = [
-            "sagemaker.amazonaws.com",
-            "lambda.amazonaws.com"
-          ]
+          Service = "sagemaker.amazonaws.com"
         }
+        Action = "sts:AssumeRole"
       }
     ]
   })
 }
 
+
+resource "aws_iam_role_policy" "s3_policy" {
+  name = "S3Policy"
+  role = aws_iam_role.zenml.id
+
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Effect = "Allow"
+        Action = [
+          "s3:ListBucket",
+          "s3:GetObject",
+          "s3:PutObject",
+          "s3:DeleteObject",
+          "s3:GetBucketVersioning"
+        ]
+        Resource = [
+          aws_s3_bucket.artifacts.arn,
+          "${aws_s3_bucket.artifacts.arn}/*"
+        ]
+      }
+    ]
+  })
+}
+
+resource "aws_iam_role_policy" "ecr_policy" {
+  name = "ECRPolicy"
+  role = aws_iam_role.zenml.id
+
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Effect = "Allow"
+        Action = [
+          "ecr:DescribeRegistry",
+          "ecr:BatchGetImage",
+          "ecr:DescribeImages",
+          "ecr:BatchCheckLayerAvailability",
+          "ecr:GetDownloadUrlForLayer",
+          "ecr:InitiateLayerUpload",
+          "ecr:UploadLayerPart",
+          "ecr:CompleteLayerUpload",
+          "ecr:PutImage"
+        ]
+        Resource = aws_ecr_repository.containers.arn
+      },
+      {
+        Effect = "Allow"
+        Action = "ecr:GetAuthorizationToken"
+        Resource = "*"
+      },
+      {
+        Effect = "Allow"
+        Action = [
+          "ecr:DescribeRepositories",
+          "ecr:ListRepositories"
+        ]
+        Resource = "arn:aws:ecr:${data.aws_region.current.name}:${data.aws_caller_identity.current.account_id}:repository/*"
+      }
+    ]
+  })
+}
+
+resource aws_iam_role_policy_attachment "sagemaker_policy" {
+  policy_arn = "arn:aws:iam::aws:policy/AmazonSageMakerFullAccess"
+  role = aws_iam_role.zenml.name
+}
+
 # ZenML Service Connector for AWS
 resource "zenml_service_connector" "aws" {
   name           = "aws-${var.environment}"
   type           = "aws"
   auth_method    = "iam-role"
-  user           = var.user_id
-  workspace      = var.workspace_id
-
-  resource_types = [
-    "artifact-store",
-    "container-registry",
-    "orchestrator",
-    "step-operator"
-  ]
 
   configuration = {
     region   = var.region
     role_arn = aws_iam_role.zenml.arn
+    aws_access_key_id = aws_iam_access_key.iam_user_access_key.id
+    aws_secret_access_key = aws_iam_access_key.iam_user_access_key.secret
   }
 
   labels = {
@@ -87,8 +176,6 @@ resource "zenml_stack_component" "artifact_store" {
   name      = "s3-${var.environment}"
   type      = "artifact_store"
   flavor    = "s3"
-  user      = var.user_id
-  workspace = var.workspace_id
 
   configuration = {
     path = "s3://${aws_s3_bucket.artifacts.bucket}/artifacts"
@@ -106,11 +193,10 @@ resource "zenml_stack_component" "container_registry" {
   name      = "ecr-${var.environment}"
   type      = "container_registry"
   flavor    = "aws"
-  user      = var.user_id
-  workspace = var.workspace_id
 
   configuration = {
-    uri = aws_ecr_repository.containers.repository_url
+    uri = regex("^([^/]+)/?", aws_ecr_repository.containers.repository_url)[0]
+    default_repository = "${aws_ecr_repository.containers.name}"
   }
 
   connector_id = zenml_service_connector.aws.id
@@ -125,11 +211,11 @@ resource "zenml_stack_component" "orchestrator" {
   name      = "sagemaker-${var.environment}"
   type      = "orchestrator"
   flavor    = "sagemaker"
-  user      = var.user_id
-  workspace = var.workspace_id
 
   configuration = {
-    role_arn = aws_iam_role.zenml.arn
+      region = data.aws_region.current.name
+      execution_role = aws_iam_role.zenml.arn
+      output_data_s3_uri = "s3://${aws_s3_bucket.artifacts.bucket}/sagemaker"
   }
 
   connector_id = zenml_service_connector.aws.id

diff --git a/examples/complete-aws/outputs.tf b/examples/complete-aws/outputs.tf
@@ -1,4 +1,3 @@
-# examples/complete-aws/outputs.tf
 output "stack_id" {
   description = "ID of the created ZenML stack"
   value       = zenml_stack.aws_stack.id