Migrated to clean branch.

.gitignore

@@ -96,4 +96,7 @@ benchmarks/solver_benchmarks/images/
 # Benchmark results
 benchmarks/solver_benchmarks/results/
 
-Dockerfile
+Dockerfile
+
+# Framework deployment packages
+.caribou/

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}

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/*"
+      }
+     ]
+  }
+}

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

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()