From ea82736b89b8321a8f70f7872bbbdbd4f51283da Mon Sep 17 00:00:00 2001 From: Danidite <1099335751qq@gmail.com> Date: Wed, 21 Aug 2024 05:15:39 +0000 Subject: [PATCH] Migrated to clean branch. --- .gitignore | 5 +- app.py | 27 ++++ .../client/cli/aws_lambda_cli/iam_policy.json | 79 ++++++++++ caribou/deployment/client/cli/cli.py | 143 ++++++++++++++++++ .../common/deploy/deployment_packager.py | 65 ++++++++ 5 files changed, 318 insertions(+), 1 deletion(-) create mode 100644 app.py create mode 100644 caribou/deployment/client/cli/aws_lambda_cli/iam_policy.json diff --git a/.gitignore b/.gitignore index 61485ad0..cbd86895 100644 --- a/.gitignore +++ b/.gitignore @@ -96,4 +96,7 @@ benchmarks/solver_benchmarks/images/ # Benchmark results benchmarks/solver_benchmarks/results/ -Dockerfile \ No newline at end of file +Dockerfile + +# Framework deployment packages +.caribou/ \ No newline at end of file diff --git a/app.py b/app.py new file mode 100644 index 00000000..6cf4f48e --- /dev/null +++ b/app.py @@ -0,0 +1,27 @@ +from typing import Any + +import json +import logging +from caribou.endpoint.client import Client + + +# Configure logging +logger = logging.getLogger() +logger.setLevel(logging.INFO) # Set the logging level + +def caribou_cli(event: dict[str, Any], context: dict[str, Any]) -> dict[str, Any]: + print(f"Received event: {event}") + if event.get("action", None) == "run": + workflow_id = event.get("workflow_id", None) + argument = event.get("argument", None) + print("Request to run workflow") + + if workflow_id: + client = Client(workflow_id) + print(f"Running workflow {workflow_id}") + if argument: + client.run(argument) + else: + client.run() + + return {"status": 200} \ No newline at end of file diff --git a/caribou/deployment/client/cli/aws_lambda_cli/iam_policy.json b/caribou/deployment/client/cli/aws_lambda_cli/iam_policy.json new file mode 100644 index 00000000..bfaf0347 --- /dev/null +++ b/caribou/deployment/client/cli/aws_lambda_cli/iam_policy.json @@ -0,0 +1,79 @@ +{ + "aws": { + "Version": "2012-10-17", + "Statement": [ + { + "Action": [ + "logs:CreateLogGroup", + "logs:CreateLogStream", + "logs:PutLogEvents" + ], + "Resource": "arn:aws:logs:*:*:*", + "Effect": "Allow" + }, + { + "Action": [ + "sns:Publish" + ], + "Resource": "arn:aws:sns:*:*:*", + "Effect": "Allow" + }, + { + "Action": [ + "dynamodb:GetItem", + "dynamodb:UpdateItem" + ], + "Resource": "arn:aws:dynamodb:*:*:*", + "Effect": "Allow" + }, + { + "Action": [ + "s3:GetObject", + "s3:PutObject" + ], + "Resource": "arn:aws:s3:::*", + "Effect": "Allow" + }, + { + "Sid": "Statement1", + "Effect": "Allow", + "Action": [ + "iam:AttachRolePolicy", + "iam:CreateRole", + "iam:CreatePolicy", + "iam:PutRolePolicy", + "ecr:CreateRepository", + "ecr:GetAuthorizationToken", + "ecr:InitiateLayerUpload", + "ecr:UploadLayerPart", + "ecr:CompleteLayerUpload", + "ecr:BatchGetImage", + "ecr:BatchCheckLayerAvailability", + "ecr:DescribeImages", + "ecr:DescribeRepositories", + "ecr:GetDownloadUrlForLayer", + "ecr:ListImages", + "ecr:PutImage", + "ecr:SetRepositoryPolicy", + "ecr:GetRepositoryPolicy", + "ecr:DeleteRepository", + "lambda:GetFunction", + "lambda:AddPermission", + "pricing:ListPriceLists", + "pricing:GetPriceListFileUrl", + "logs:FilterLogEvents" + ], + "Resource": "*" + }, + { + "Sid": "Statement2", + "Effect": "Allow", + "Action": [ + "iam:GetRole", + "iam:PassRole" + ], + "Resource": "arn:aws:iam:::role/*" + } + ] + } +} \ No newline at end of file diff --git a/caribou/deployment/client/cli/cli.py b/caribou/deployment/client/cli/cli.py index 5c573cde..3bb1d788 100644 --- a/caribou/deployment/client/cli/cli.py +++ b/caribou/deployment/client/cli/cli.py @@ -1,8 +1,16 @@ +import json import os +from pathlib import Path +import subprocess +import sys +import tempfile from typing import Optional +from unittest.mock import MagicMock +import zipfile import click +from caribou.common.models.remote_client.aws_remote_client import AWSRemoteClient from caribou.common.setup.setup_tables import main as setup_tables_func from caribou.data_collector.components.carbon.carbon_collector import CarbonCollector from caribou.data_collector.components.performance.performance_collector import PerformanceCollector @@ -12,11 +20,15 @@ from caribou.deployment.client.cli.new_workflow import create_new_workflow_directory from caribou.deployment.common.config.config import Config from caribou.deployment.common.deploy.deployer import Deployer +from caribou.deployment.common.deploy.deployment_packager import DeploymentPackager +from caribou.deployment.common.deploy.models.resource import Resource from caribou.deployment.common.factories.deployer_factory import DeployerFactory from caribou.endpoint.client import Client from caribou.monitors.deployment_manager import DeploymentManager from caribou.monitors.deployment_migrator import DeploymentMigrator from caribou.syncers.log_syncer import LogSyncer +import boto3 + @click.group() @@ -128,5 +140,136 @@ def remove(workflow_id: str) -> None: client = Client(workflow_id) client.remove() +region_name = "us-east-2" # Test region + +@cli.command("deploy_framework", help="Deploy the framework to Lambda.") +@click.pass_context +def deploy_framework(ctx: click.Context) -> None: + # factory: DeployerFactory = ctx.obj["factory"] + # config: Config = factory.create_config_obj() + # deployer: Deployer = factory.create_deployer(config=config) + # deployer.deploy([config.home_region]) + lambda_trust_policy = { + "Version": "2012-10-17", + "Statement": [ + { + "Sid": "", + "Effect": "Allow", + "Principal": { + "Service": ["lambda.amazonaws.com", "states.amazonaws.com"] + }, + "Action": "sts:AssumeRole", + } + ], + } + + + + aws_remote_client = AWSRemoteClient(region_name) + + project_dir = ctx.obj["project_dir"] + handler = "app.caribou_cli" + function_name = "caribou_cli" + runtime = "python:3.12" + timeout = 600 + memory_size = 3008 + iam_policy_name = "caribou_deployment_policy" + + deployment_packager_config: MagicMock = MagicMock(spec=Config) + deployment_packager_config.workflow_version = "1.0.0" + deployment_packager: DeploymentPackager = DeploymentPackager(deployment_packager_config) + + # # Read the iam_policies_content from the file + # with open("caribou/deployment/client/cli/aws_lambda_cli/iam_policy.json", "r") as file: + # iam_policies_content = file.read() + # iam_policies_content = json.dumps(json.loads(iam_policies_content)["aws"]) + + # # Delete role if exists, then create a new role + # # Delete a role + # if aws_remote_client.resource_exists(Resource(iam_policy_name, "iam_role")): # For iam role + # print(f"Deleting role {iam_policy_name}") + # aws_remote_client.remove_role(iam_policy_name) + + # # Create a role + # role_arn = aws_remote_client.create_role("caribou_deployment_policy", iam_policies_content, lambda_trust_policy) + + # print(f"Role ARN: {role_arn}") + + role_arn = "arn:aws:iam::226414417076:role/caribou_deployment_policy" + + # # Delete function if exists. + # if aws_remote_client.resource_exists(Resource(function_name, "function")): # For lambda function + # aws_remote_client.remove_function(function_name) + + # Create lambda function + ## First zip the code content + deployment_packager_config.workflow_name = function_name + + print(f"Creating deployment package for {function_name}") + zip_path = deployment_packager.create_framework_package(project_dir) + + with open(zip_path, 'rb') as f: + zip_contents = f.read() + + deploy_to_aws(function_name, handler, runtime, role_arn, timeout, memory_size, zip_contents) + +def deploy_to_aws( + function_name: str, handler: str, runtime: str, role_arn: str, timeout: int, memory_size: int, zip_contents: bytes +): + aws_remote_client = AWSRemoteClient(region_name) + + with tempfile.TemporaryDirectory() as tmpdirname: + # Step 1: Unzip the ZIP file + zip_path = os.path.join(tmpdirname, "code.zip") + with open(zip_path, "wb") as f_zip: + f_zip.write(zip_contents) + with zipfile.ZipFile(zip_path, "r") as zip_ref: + zip_ref.extractall(tmpdirname) + + # Step 2: Create a Dockerfile in the temporary directory + dockerfile_content = generate_framework_dockerfile(handler, runtime) + with open(os.path.join(tmpdirname, "Dockerfile"), "w", encoding="utf-8") as f_dockerfile: + f_dockerfile.write(dockerfile_content) + + # Step 3: Build the Docker Image + image_name = f"{function_name.lower()}:latest" + aws_remote_client._build_docker_image(tmpdirname, image_name) + + # Step 4: Upload the Image to ECR + image_uri = aws_remote_client._upload_image_to_ecr(image_name) + + create_lambda_function(function_name, image_uri, role_arn, timeout, memory_size) + +def generate_framework_dockerfile(handler: str, runtime: str) -> str: + return f""" + FROM public.ecr.aws/lambda/{runtime} + + RUN pip3 install poetry + + COPY app.py ./ + COPY caribou ./caribou + COPY caribou-go ./caribou-go + COPY pyproject.toml ./ + + RUN poetry install --no-dev + CMD ["{handler}"] + """ + +def create_lambda_function(function_name: str, image_uri: str, role: str, timeout: int, memory_size: int) -> None: + lambda_client = boto3.client("lambda", region_name=region_name) + + try: + lambda_client.create_function( + FunctionName=function_name, + Role=role, + Code={"ImageUri": image_uri}, + PackageType="Image", + Timeout=timeout, + MemorySize=memory_size, + ) + except lambda_client.exceptions.ResourceConflictException: + print(f"Lambda function {function_name} already exists") + pass + __version__ = MULTI_X_SERVERLESS_VERSION diff --git a/caribou/deployment/common/deploy/deployment_packager.py b/caribou/deployment/common/deploy/deployment_packager.py index b8f61e72..15e42ff5 100644 --- a/caribou/deployment/common/deploy/deployment_packager.py +++ b/caribou/deployment/common/deploy/deployment_packager.py @@ -8,6 +8,7 @@ import subprocess import sys import tempfile +import time import zipfile from typing import Optional @@ -69,6 +70,70 @@ def _create_deployment_package(self, project_dir: str, python_version: str) -> s self._add_requirements_file(z, requirements_filename) return package_filename + def create_framework_package(self, project_dir: str) -> str: + filename = "caribou_framework_cli.zip" + package_filename = os.path.join(project_dir, ".caribou", "deployment-packages", filename) + self._create_deployment_package_dir(package_filename) + if os.path.exists(package_filename): + return package_filename + + # Remove existing framework package if it exists + if os.path.exists(package_filename): + os.remove(package_filename) + + # Wait for the file to be removed + time.sleep(1) + + with zipfile.ZipFile(package_filename, "w", zipfile.ZIP_DEFLATED) as z: + self._add_framework_deployment_files(z, project_dir) + self._add_framework_files(z, project_dir) + self._add_framework_go_files(z, project_dir) + return package_filename + + def _add_framework_deployment_files(self, zip_file: zipfile.ZipFile, project_dir: str) -> None: + for root, _, files in os.walk(project_dir): + for filename in files: + if filename.endswith(".pyc"): + continue + + full_path = os.path.join(root, filename) + if (full_path == os.path.join(project_dir, "app.py") or + full_path.startswith(os.path.join(project_dir, "src")) or + filename == "pyproject.toml" + # or filename == "poetry.lock" + ): + zip_path = full_path[len(project_dir) + 1 :] + zip_file.write(full_path, zip_path) + + def _add_framework_files(self, zip_file: zipfile.ZipFile, project_dir: str) -> None: + framework_dir = os.path.join(project_dir, "caribou") + + for root, _, files in os.walk(project_dir): + for filename in files: + if not filename.endswith(".py"): # Only add .py files + continue + + full_path = os.path.join(root, filename) + if full_path.startswith(framework_dir): + zip_path = full_path[len(project_dir) + 1 :] + zip_file.write(full_path, zip_path) + + def _add_framework_go_files(self, zip_file: zipfile.ZipFile, project_dir: str) -> None: + framework_go_dir = os.path.join(project_dir, "caribou-go") + allowed_extensions = ['.go', '.py', '.sum', '.mod', '.sh', '.so'] + + for root, _, files in os.walk(project_dir): + for filename in files: + if not any(filename.endswith(ext) for ext in allowed_extensions): # Check if file has an allowed extension + continue + + full_path = os.path.join(root, filename) + if full_path.startswith( + framework_go_dir + ): + zip_path = full_path[len(project_dir) + 1 :] + zip_file.write(full_path, zip_path) + def _ensure_requirements_filename_complete(self, requirements_filename: str) -> None: with open(requirements_filename, "r", encoding="utf-8") as file: requirements = file.read().splitlines()