Add the CloudBuild image builder to the AWS one-click stack deployment

zenml-io · Dec 2, 2024 · d41fdb0 · d41fdb0
1 parent 7fad641
commit d41fdb0
Show file tree

Hide file tree

Showing 2 changed files with 184 additions and 113 deletions.
diff --git a/infra/aws/aws-ecr-s3-sagemaker.yaml b/infra/aws/aws-ecr-s3-sagemaker.yaml
@@ -51,11 +51,23 @@ Parameters:
     Description: "The value of the tag to apply to all resources"
     Default: "zenml"
 
+  CodeBuild:
+    Type: String
+    AllowedValues:
+      - true
+      - false
+    Description: |
+      Whether to provision a CodeBuild project as the image builder for the
+      stack. Only supported for ZenML Server versions above 0.70.0.
+    Default: false
+
 Conditions:
   RegisterZenMLStack: !And
     - !Not [ !Equals [ !Ref ZenMLServerURL, "" ] ]
     - !Not [ !Equals [ !Ref ZenMLServerAPIToken, "" ] ]
 
+  RegisterCodeBuild: !Equals [ !Ref CodeBuild, true ]
+
 Resources:
   S3Bucket:
     Type: AWS::S3::Bucket
@@ -73,6 +85,28 @@ Resources:
       Tags:
         - Key: !Ref TagName
           Value: !Sub TagValue
+
+  CodeBuildProject:
+    Condition: RegisterCodeBuild
+    Type: AWS::CodeBuild::Project
+    Properties:
+      Name: !Sub '${ResourceName}'
+      ServiceRole: !GetAtt CodeBuildRole.Arn
+      Artifacts:
+        Type: NO_ARTIFACTS
+      Environment:
+        Type: LINUX_CONTAINER
+        ComputeType: BUILD_GENERAL1_SMALL
+        Image: bentolor/docker-dind-awscli
+        PrivilegedMode: false
+      Source:
+        Type: S3
+        Location: !Sub '${S3Bucket}/codebuild'
+      TimeoutInMinutes: 20
+      LogsConfig:
+        CloudWatchLogs:
+          Status: ENABLED
+          GroupName: !Sub '/aws/codebuild/${ResourceName}'
 
   IAMUser:
     Type: AWS::IAM::User
@@ -176,6 +210,19 @@ Resources:
               - Effect: Allow
                 Action: iam:PassRole
                 Resource: !Sub 'arn:aws:iam::${AWS::AccountId}:role/${ResourceName}-sagemaker'
+        - !If
+          - RegisterCodeBuild
+          - PolicyName: CodeBuildPolicy
+            PolicyDocument:
+              Version: '2012-10-17'
+              Statement:
+                # Allow this role to start and monitor CodeBuild project builds
+                - Effect: Allow
+                  Action:
+                    - 'codebuild:StartBuild'
+                    - 'codebuild:BatchGetBuilds'
+                  Resource: !Sub 'arn:aws:codebuild:${AWS::Region}:${AWS::AccountId}:project/${ResourceName}'
+          - !Ref 'AWS::NoValue'
 
   SageMakerRuntimeRole:
     Type: AWS::IAM::Role
@@ -205,6 +252,53 @@ Resources:
       ManagedPolicyArns: 
         - 'arn:aws:iam::aws:policy/AmazonSageMakerFullAccess'
 
+  CodeBuildRole:
+    Type: AWS::IAM::Role
+    Condition: RegisterCodeBuild
+    Properties:
+      RoleName: !Sub '${ResourceName}-codebuild'
+      AssumeRolePolicyDocument:
+        Version: '2012-10-17'
+        Statement:
+          - Effect: Allow
+            Principal:
+              Service: codebuild.amazonaws.com
+            Action: 'sts:AssumeRole'
+      Policies:
+        - PolicyName: CodeBuildPolicy
+          PolicyDocument:
+            Version: '2012-10-17'
+            Statement:
+              - Effect: Allow
+                Action:
+                  - 'logs:CreateLogGroup'
+                  - 'logs:CreateLogStream'
+                  - 'logs:PutLogEvents'
+                Resource:
+                  - !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/codebuild/${ResourceName}'
+                  - !Sub 'arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/codebuild/${ResourceName}:*'
+              - Effect: Allow
+                Action:
+                  - 's3:GetObject'
+                  - 's3:GetObjectVersion'
+                Resource:
+                  - !Sub '${S3Bucket.Arn}/*'
+              - Effect: Allow
+                Action:
+                  - 'ecr:BatchGetImage'
+                  - 'ecr:DescribeImages'
+                  - 'ecr:BatchCheckLayerAvailability'
+                  - 'ecr:GetDownloadUrlForLayer'
+                  - 'ecr:InitiateLayerUpload'
+                  - 'ecr:UploadLayerPart'
+                  - 'ecr:CompleteLayerUpload'
+                  - 'ecr:PutImage'
+                Resource: !Sub '${ECRRepository.Arn}'
+              - Effect: Allow
+                Action:
+                  - 'ecr:GetAuthorizationToken'
+                Resource: '*'
+
   InvokeZenMLAPIFunction:
     Type: AWS::Serverless::Function
     Condition: RegisterZenMLStack
@@ -308,63 +402,77 @@ Resources:
     Properties:
       ServiceToken: !GetAtt InvokeZenMLAPIFunction.Arn
       ServiceTimeout: 300
-      Payload: !Sub |
-        {
-          "name": "${AWS::StackName}",
-          "description": "Deployed by AWS CloudFormation stack ${AWS::StackName} in the ${AWS::AccountId} account and ${AWS::Region} region.",
-          "labels": {
-            "zenml:provider": "aws",
-            "zenml:deployment": "cloud-formation"
-          },
-          "service_connectors": [
+      Payload: !Join
+        - ''
+        - - !Sub |
             {
-              "type": "aws",
-              "auth_method": "iam-role",
-              "configuration": {
-                "aws_access_key_id": "${IAMUserAccessKey}",
-                "aws_secret_access_key": "${IAMUserAccessKey.SecretAccessKey}",
-                "role_arn": "${StackAccessRole.Arn}",
-                "region": "${AWS::Region}"
-              }
-            }
-          ],
-          "components": {
-            "artifact_store": [{
-              "flavor": "s3",
-              "service_connector_index": 0,
-              "configuration": {
-                "path": "s3://${S3Bucket}"
-              }
-            }],
-            "container_registry":[{
-              "flavor": "aws",
-              "service_connector_index": 0,
-              "configuration": {
-                "uri": "${AWS::AccountId}.dkr.ecr.${AWS::Region}.amazonaws.com",
-                "default_repository": "${ECRRepository}"
-              }
-            }],
-            "orchestrator": [{
-              "flavor": "sagemaker",
-              "service_connector_index": 0,
-              "configuration": {
-                "execution_role": "${SageMakerRuntimeRole.Arn}",
-                "output_data_s3_uri": "s3://${S3Bucket}/sagemaker"
+              "name": "${AWS::StackName}",
+              "description": "Deployed by AWS CloudFormation stack ${AWS::StackName} in the ${AWS::AccountId} account and ${AWS::Region} region.",
+              "labels": {
+                "zenml:provider": "aws",
+                "zenml:deployment": "cloud-formation"
+              },
+              "service_connectors": [
+                {
+                  "type": "aws",
+                  "auth_method": "iam-role",
+                  "configuration": {
+                    "aws_access_key_id": "${IAMUserAccessKey}",
+                    "aws_secret_access_key": "${IAMUserAccessKey.SecretAccessKey}",
+                    "role_arn": "${StackAccessRole.Arn}",
+                    "region": "${AWS::Region}"
+                  }
+                }
+              ],
+              "components": {
+                "artifact_store": [{
+                  "flavor": "s3",
+                  "service_connector_index": 0,
+                  "configuration": {
+                    "path": "s3://${S3Bucket}"
+                  }
+                }],
+                "container_registry":[{
+                  "flavor": "aws",
+                  "service_connector_index": 0,
+                  "configuration": {
+                    "uri": "${AWS::AccountId}.dkr.ecr.${AWS::Region}.amazonaws.com",
+                    "default_repository": "${ECRRepository}"
+                  }
+                }],
+                "orchestrator": [{
+                  "flavor": "sagemaker",
+                  "service_connector_index": 0,
+                  "configuration": {
+                    "execution_role": "${SageMakerRuntimeRole.Arn}",
+                    "output_data_s3_uri": "s3://${S3Bucket}/sagemaker"
+                  }
+                }],
+                "step_operator": [{
+                  "flavor": "sagemaker",
+                  "service_connector_index": 0,
+                  "configuration": {
+                    "role": "${SageMakerRuntimeRole.Arn}",
+                    "bucket": "${S3Bucket}"
+                  }
+                }],
+          - !If
+            - RegisterCodeBuild
+            - !Sub |
+                "image_builder": [{
+                  "flavor": "aws",
+                  "service_connector_index": 0,
+                  "configuration": {
+                    "code_build_project": "${CodeBuildProject}"
+                  }
+                }]
+            - |
+                "image_builder": [{
+                  "flavor": "local"
+                }]
+          - |
               }
-            }],
-            "step_operator": [{
-              "flavor": "sagemaker",
-              "service_connector_index": 0,
-              "configuration": {
-                "role": "${SageMakerRuntimeRole.Arn}",
-                "bucket": "${S3Bucket}"
               }
-            }],
-            "image_builder": [{
-              "flavor": "local"
-            }]
-          }
-        }
 
 Outputs:
   AWSRegion:
@@ -387,62 +495,3 @@ Outputs:
     Description: "SageMaker execution IAM Role ARN"
     Value: !GetAtt SageMakerRuntimeRole.Arn
 
-  ZenMLStack:
-    Description: "ZenML Stack JSON (can be imported with `zenml stack import`)"
-    Value: !Sub |
-      {
-        "name": "${AWS::StackName}",
-        "description": "Deployed by AWS CloudFormation stack ${AWS::StackName} in the ${AWS::AccountId} account and ${AWS::Region} region.",
-        "labels": {
-          "zenml:provider": "aws",
-          "zenml:deployment": "aws-cloud-formation"
-        },
-        "service_connectors": [
-          {
-            "type": "aws",
-            "auth_method": "iam-role",
-            "configuration": {
-              "aws_access_key_id": "${IAMUserAccessKey}",
-              "aws_secret_access_key": "${IAMUserAccessKey.SecretAccessKey}",
-              "role_arn": "${StackAccessRole.Arn}",
-              "region": "${AWS::Region}"
-            }
-          }
-        ],
-        "components": {
-          "artifact_store": [{
-            "flavor": "s3",
-            "service_connector_index": 0,
-            "configuration": {
-              "path": "s3://${S3Bucket}"
-            }
-          }],
-          "container_registry":[{
-            "flavor": "aws",
-            "service_connector_index": 0,
-            "configuration": {
-              "uri": "${AWS::AccountId}.dkr.ecr.${AWS::Region}.amazonaws.com",
-              "default_repository": "${ECRRepository}"
-            }
-          }],
-          "orchestrator": [{
-            "flavor": "sagemaker",
-            "service_connector_index": 0,
-            "configuration": {
-              "execution_role": "${SageMakerRuntimeRole.Arn}",
-              "output_data_s3_uri": "s3://${S3Bucket}/sagemaker"
-            }
-          }],
-          "step_operator": [{
-            "flavor": "sagemaker",
-            "service_connector_index": 0,
-            "configuration": {
-              "role": "${SageMakerRuntimeRole.Arn}",
-              "bucket": "${S3Bucket}"
-            }
-          }],
-          "image_builder": [{
-            "flavor": "local"
-          }]
-        }
-      }
diff --git a/src/zenml/stack_deployments/aws_stack_deployment.py b/src/zenml/stack_deployments/aws_stack_deployment.py
@@ -73,6 +73,7 @@ def instructions(cls) -> str:
 - An ECR repository registered as a [ZenML container registry](https://docs.zenml.io/stack-components/container-registries/aws).
 - Sagemaker registered as a [ZenML orchestrator](https://docs.zenml.io/stack-components/orchestrators/sagemaker)
 as well as a [ZenML step operator](https://docs.zenml.io/stack-components/step-operators/sagemaker).
+- A CodeBuild project registered as a [ZenML image builder](https://docs.zenml.io/stack-components/image-builder/aws).
 - An IAM user and IAM role with the minimum necessary permissions to access the
 above resources.
 - An AWS access key used to give access to ZenML to connect to the above
@@ -158,6 +159,26 @@ def permissions(cls) -> Dict[str, List[str]]:
                 "ecr:PutImage",
                 "ecr:GetAuthorizationToken",
             ],
+            "CloudBuild (Client)": [
+                "codebuild:CreateProject",
+                "codebuild:BatchGetBuilds",
+            ],
+            "CloudBuild (Service)": [
+                "s3:GetObject",
+                "s3:GetObjectVersion",
+                "logs:CreateLogGroup",
+                "logs:CreateLogStream",
+                "logs:PutLogEvents",
+                "ecr:BatchGetImage",
+                "ecr:DescribeImages",
+                "ecr:BatchCheckLayerAvailability",
+                "ecr:GetDownloadUrlForLayer",
+                "ecr:InitiateLayerUpload",
+                "ecr:UploadLayerPart",
+                "ecr:CompleteLayerUpload",
+                "ecr:PutImage",
+                "ecr:GetAuthorizationToken",
+            ],
             "SageMaker (Client)": [
                 "sagemaker:CreatePipeline",
                 "sagemaker:StartPipelineExecution",
@@ -243,6 +264,7 @@ def get_deployment_config(
             param_ResourceName=f"zenml-{random_str(6).lower()}",
             param_ZenMLServerURL=self.zenml_server_url,
             param_ZenMLServerAPIToken=self.zenml_server_api_token,
+            param_CodeBuild="true",
         )
         # Encode the parameters as URL query parameters
         query_params = "&".join([f"{k}={v}" for k, v in params.items()])