metagraph-stack.yaml

AWSTemplateFormatVersion: '2010-09-09'
Description: 'Deploys an AWS Batch-based system for Metagraph query execution.'

Parameters:
  MetagraphAmiId:
    Type: AWS::EC2::Image::Id
    Description: AMI ID to use for Metagraph Batch instances
  NotificationEmail:
    Type: String
    Description: Email address to subscribe to the SNS topic

Resources:

### Roles and profiles
  ECSInstanceRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - ec2.amazonaws.com
            Action:
              - "sts:AssumeRole"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AmazonEC2ContainerServiceforEC2Role
      Policies:
        - PolicyName: MountS3ReadAccess
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - s3:ListBucket
                  - s3:GetObject
                  - s3:PutObject
                Resource:
                  - arn:aws:s3:::metagraph
                  - arn:aws:s3:::metagraph/*
                  - !Sub arn:aws:s3:::${MetagraphBucket}
                  - !Sub arn:aws:s3:::${MetagraphBucket}/*

  ECSInstanceProfile:
    Type: AWS::IAM::InstanceProfile
    Properties:
      Roles:
        - !Ref ECSInstanceRole

  MetagraphLambdaRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: LambdaSubmitBatchJobsPolicy
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 'batch:SubmitJob'
                  - 'batch:DescribeJobs'
                Resource: "*"
              - Effect: Allow
                Action:
                  - 's3:ListBucket'
                  - 's3:GetObject'
                  - 's3:PutObject'
                Resource:
                  - 'arn:aws:s3:::metagraph*'
                  - !Sub 'arn:aws:s3:::${MetagraphBucket}*'
              - Effect: Allow
                Action:
                  - sns:Publish
                Resource: !Ref MetagraphNotificationTopic

  MetagraphStepFunctionsRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: states.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: StepFunctionsNestedExecution
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - states:StartExecution
                Resource: "*"
        - PolicyName: StepFunctionsBatchAccess
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - batch:SubmitJob
                  - batch:DescribeJobs
                Resource: "*"
        - PolicyName: StepFunctionsInvokeLambda
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: lambda:InvokeFunction
                Resource:
                  - !GetAtt MetagraphTaskScheduler.Arn
                  - !GetAtt MetagraphNotifier.Arn
        - PolicyName: StepFunctionsEventBridgePermissions
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - events:PutRule
                  - events:PutTargets
                  - events:DescribeRule
                  - events:DeleteRule
                  - events:RemoveTargets
                Resource: "*"

  MetagraphBatchJobRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service: ecs-tasks.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: MetagraphJobS3Access
          PolicyDocument:
            Version: '2012-10-17'
            Statement:
              - Effect: Allow
                Action:
                  - 's3:GetObject'
                  - 's3:PutObject'
                  - 's3:ListBucket'
                Resource:
                  - 'arn:aws:s3:::metagraph*'
                  - !Sub 'arn:aws:s3:::${MetagraphBucket}'
                  - !Sub 'arn:aws:s3:::${MetagraphBucket}/*'

### Main stack bucket
  MetagraphBucket:
    Type: AWS::S3::Bucket

### SNS topic for notifications
  MetagraphNotificationTopic:
    Type: AWS::SNS::Topic

  MetagraphNotificationSubscription:
    Type: AWS::SNS::Subscription
    Properties:
      TopicArn: !Ref MetagraphNotificationTopic
      Protocol: email
      Endpoint: !Ref NotificationEmail

### Launch template for container service with s3 bucket mounting
  MetagraphLaunchTemplate:
    Type: AWS::EC2::LaunchTemplate
    Properties:
      LaunchTemplateData:
        IamInstanceProfile:
          Name: !Ref ECSInstanceProfile
        ImageId: !Ref MetagraphAmiId
        UserData:
          Fn::Base64: !Sub |
            MIME-Version: 1.0
            Content-Type: multipart/mixed; boundary="==BOUNDARY=="

            --==BOUNDARY==
            Content-Type: text/cloud-config; charset="us-ascii"

            cloud_final_modules:
              - [scripts-user, always]

            --==BOUNDARY==
            Content-Type: text/x-shellscript; charset="us-ascii"

            #!/bin/bash
            set -e

            mkdir -p /mnt/data
            mkdir -p /mnt/bucket

            mount-s3 metagraph /mnt/data
            mount-s3 ${MetagraphBucket} /mnt/bucket

            --==BOUNDARY==--


  MetagraphBatchComputeEnvironment:
    Type: AWS::Batch::ComputeEnvironment
    Properties:
      Type: MANAGED
      State: ENABLED
      ComputeResources:
        Type: EC2
        MinvCpus: 0
        MaxvCpus: 512
        InstanceTypes:
          - r6i.4xlarge # 16 vCPU, 128 GiB RAM
        Subnets:
          - subnet-064a41e7e9f7821a8 # eu-central-2a
          - subnet-0247650ef94e8aa42 # eu-central-2b
          - subnet-02ebd4969f85757c4 # eu-central-2c
        SecurityGroupIds:
          - sg-08aa6d6a56d63b8cd
        InstanceRole: !GetAtt ECSInstanceProfile.Arn
        LaunchTemplate:
          LaunchTemplateId: !Ref MetagraphLaunchTemplate
          Version: !GetAtt MetagraphLaunchTemplate.LatestVersionNumber

  MetagraphBatchJobQueue:
    Type: AWS::Batch::JobQueue
    Properties:
      State: ENABLED
      Priority: 1
      ComputeEnvironmentOrder:
        - Order: 1
          ComputeEnvironment: !Ref MetagraphBatchComputeEnvironment

  MetagraphMergeJobDefinition:
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: container
      RetryStrategy:
        Attempts: 1
      ContainerProperties:
        Image: public.ecr.aws/amazonlinux/amazonlinux:2023
        Vcpus: 4
        Memory: 8192
        Volumes:
          - Name: metagraph-bucket
            Host:
              SourcePath: /mnt/bucket
        MountPoints:
          - SourceVolume: metagraph-bucket
            ContainerPath: /mnt/bucket
        Command:
          - python3
          - -c
          - |
              import os
              import shutil
              import itertools as it
              from contextlib import ExitStack
              from concurrent.futures import ThreadPoolExecutor

              BUCKET_MOUNT = '/mnt/bucket'
              job_id = os.environ["JOB_ID"]
              mode = os.environ["QUERY_MODE"]

              if mode not in ['labels', 'matches']:
                  raise ValueError(f"Invalid query mode: {mode}")

              job_path = os.path.join(BUCKET_MOUNT, job_id)
              merged_key = f"{job_id}-result.txt"

              with ExitStack() as stack:
                  merged_file = stack.enter_context(open(merged_key, 'w'))

                  print(f"Merging results from {job_path}...")
                  with ThreadPoolExecutor() as executor:
                      file_paths = [os.path.join(job_path, f) for f in os.listdir(job_path)]
                      files = list(executor.map(lambda path: stack.enter_context(open(path, 'r')), file_paths))
                  print(f"Found {len(files)} files to merge...")

                  for Idx in it.count(0):
                      if Idx and Idx % 1000 == 0:
                          print(f"Processed {Idx} lines...")
                      with ThreadPoolExecutor() as executor:
                          lines = list(executor.map(lambda f: f.readline() or None, files))

                      if not lines or not all(lines):
                          break

                      Name = ''
                      Matches = []

                      for line in lines:
                          idx, name, *matches = line.strip().split('\t')
                          if not Name:
                              Name = name
                          assert name == Name and int(idx) == Idx

                          if mode == 'labels':
                              Matches += matches
                          elif mode == 'matches':
                              for match in matches:
                                  label, count = match.split(':')
                                  Matches.append((int(count), label))

                      if mode == 'labels':
                          content = f"{Idx}\t{Name}\t{':'.join(Matches)}\n"
                      elif mode == 'matches':
                          match_str = '\t'.join(f'{label}:{count}' for count, label in reversed(sorted(Matches)))
                          content = f"{Idx}\t{Name}\t{match_str}\n"

                      merged_file.write(content)

              print(f"Copying {merged_key} to the s3 bucket...")
              with open(merged_key, 'rb') as src, open(os.path.join(BUCKET_MOUNT, merged_key), 'wb') as dst:
                dst.write(src.read())
        Environment:
          - Name: PYTHONUNBUFFERED
            Value: "1"
        JobRoleArn: !GetAtt MetagraphBatchJobRole.Arn

  MetagraphBatchJobDefinition:
    Type: AWS::Batch::JobDefinition
    Properties:
      Type: container
      RetryStrategy:
        Attempts: 1
      ContainerProperties:
        Image: public.ecr.aws/amazonlinux/amazonlinux:2023
        Vcpus: 4
        Memory: 32768
        Volumes:
          - Name: metagraph-data
            Host:
              SourcePath: /mnt/data
          - Name: metagraph-bucket
            Host:
              SourcePath: /mnt/bucket
          - Name: metagraph-bin
            Host:
              SourcePath: /opt/metagraph
          - Name: libdeflate
            Host:
              SourcePath: /usr/local/lib64
          - Name: lib64
            Host:
              SourcePath: /usr/lib64
        MountPoints:
          - SourceVolume: metagraph-data
            ContainerPath: /mnt/data
          - SourceVolume: metagraph-bucket
            ContainerPath: /mnt/bucket
          - SourceVolume: metagraph-bin
            ContainerPath: /opt/metagraph
          - SourceVolume: libdeflate
            ContainerPath: /usr/local/lib64
          - SourceVolume: lib64
            ContainerPath: /usr/lib64
        Command:
          - "/bin/bash"
          - "-c"
          - |
            set -e

            export PATH="/opt/metagraph/metagraph/build:$PATH"
            export LD_LIBRARY_PATH="/usr/local/lib64:/usr/lib64:$LD_LIBRARY_PATH"

            metagraph query \
              -i "/mnt/data/$DATASET/$GRAPH_FILE" \
              -a "/mnt/data/$DATASET/$ANNO_FILE" \
              --query-mode "$QUERY_MODE" \
              --num-top-labels "$NUM_TOP_LABELS" \
              --min-kmers-fraction-label "$MIN_KMERS_FRACTION_LABEL" \
              --min-kmers-fraction-graph "$MIN_KMERS_FRACTION_GRAPH" \
              "/mnt/bucket/$QUERY" > result.txt

            mkdir -p "/mnt/bucket/$JOB_ID"
            cp result.txt "/mnt/bucket/$JOB_ID/$AWS_BATCH_JOB_ID.txt"
        JobRoleArn: !GetAtt MetagraphBatchJobRole.Arn

  MetagraphTaskScheduler:
    Type: AWS::Lambda::Function
    Properties:
      Role: !GetAtt MetagraphLambdaRole.Arn
      Runtime: python3.13
      Handler: index.lambda_handler
      Timeout: 300
      Code:
        ZipFile: !Sub |
          import boto3
          import re
          from collections import defaultdict

          s3 = boto3.client('s3')

          def list_indices(bucket, prefix):
              prefix = prefix.strip('/').rstrip('/')
              prefix += '/' if prefix else ''
              directories = defaultdict(dict)
              paginator = s3.get_paginator("list_objects_v2")
              for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
                  for element in page.get("Contents", []):
                      key = element['Key']
                      is_graph = key.endswith('.dbg')
                      is_anno = key.endswith('.annodbg')
                      if is_graph or is_anno:
                          splt = key.split('/')
                          directories['/'.join(splt[:-1])[len(prefix):]]['graph' if is_graph else 'anno'] = splt[-1]
              return directories

          def file_size(bucket, key): # bytes
              return s3.head_object(Bucket=bucket, Key=key)["ContentLength"]

          def lambda_handler(event, context):
              prefix = event["query"]["index_prefix"].rstrip('/')
              bucket = 'metagraph'

              all_directories = list_indices(bucket, prefix)
              regex_pattern = event["query"].get("index_filter", ".*")
              matching_dirs = [
                  (path, files)
                  for path, files in all_directories.items()
                  if re.match(regex_pattern, path)
              ]

              job_id = event["arn"].split(":")[-1] if "arn" in event else context.aws_request_id
              query_mode = event["query"].get("query_mode", "labels")

              print(f"Found {len(all_directories)} directories, {len(matching_dirs)} matched '{regex_pattern}'")

              query_mode = event["query"].get("query_mode", "labels")
              query_filename = event["query"]["query_filename"]
              query_path = f"queries/{query_filename}"
              num_top_labels = str(event["query"].get("num_top_labels", "inf"))
              min_kmers_fraction_label = str(event["query"].get("min_kmers_fraction_label", 0.7))
              min_kmers_fraction_graph = str(event["query"].get("min_kmers_fraction_graph", 0.0))

              tasks = []
              for relative_path, files in matching_dirs:
                  if "graph" not in files or "anno" not in files:
                      print(f"Skipping {relative_path}, missing required files: {list(files.keys())}")
                      continue

                  graph_file = files["graph"]
                  anno_file = files["anno"]
                  graph_size = file_size(bucket, f"{prefix}/{relative_path}/{graph_file}")
                  anno_size = file_size(bucket, f"{prefix}/{relative_path}/{anno_file}")
                  memory_mib = (graph_size + anno_size) // 2**20 + 8*1024

                  task_name = relative_path.replace("/", "_")
                  tasks.append({
                      "jobName": f"Query-{job_id}-{task_name}"[:128],
                      "dataset": f"{prefix}/{relative_path}",
                      "graph_file": graph_file,
                      "anno_file": anno_file,
                      "memory_mib": memory_mib,
                      "job_id": job_id,
                      "query_mode": query_mode,
                      "query": query_path,
                      "num_top_labels": num_top_labels,
                      "min_kmers_fraction_label": min_kmers_fraction_label,
                      "min_kmers_fraction_graph": min_kmers_fraction_graph,
                  })

              return {
                  "job_id": job_id,
                  "query_mode": query_mode,
                  "num_top_labels": num_top_labels,
                  "tasks": tasks
              }

  MetagraphNotifier:
    Type: AWS::Lambda::Function
    Properties:
      Role: !GetAtt MetagraphLambdaRole.Arn
      Runtime: python3.13
      Handler: index.lambda_handler
      Timeout: 60
      Code:
        ZipFile: !Sub |
          import boto3, os
          from botocore.config import Config
          from urllib.parse import urlparse

          sns = boto3.client("sns")
          s3 = boto3.client("s3", region_name='eu-central-2', endpoint_url='https://s3.eu-central-2.amazonaws.com')

          def lambda_handler(event, context):
              tasks = event["tasks"]

              fails = sum(task["status"] == "FAILED" for task in tasks)
              total = len(tasks)

              job_id = event["job_id"]
              mode = event["query_mode"]

              presigned_url = s3.generate_presigned_url(
                  ClientMethod='get_object',
                  Params={'Bucket': '${MetagraphBucket}', 'Key': f'{job_id}-result.txt'},
                  ExpiresIn=7*24*60*60  # 7 days
              )

              print(f"Sending notification to topic ${MetagraphNotificationTopic}")
              sns.publish(
                  TopicArn="${MetagraphNotificationTopic}",
                  Message=(
                      f"Your Metagraph query {job_id} has completed.\n\n"
                      f"Query mode: {mode}\n"
                      f"Failed tasks: {fails} out of {total}\n\n"
                      f"Download your results here (valid for 7 days):\n\n{presigned_url}"
                  ),
                  Subject="Your Metagraph query has completed"
              )

  MetagraphStateMachine:
    Type: AWS::StepFunctions::StateMachine
    Properties:
      RoleArn: !GetAtt MetagraphStepFunctionsRole.Arn
      DefinitionString: !Sub |
        {
          "StartAt": "InvokeScheduler",
          "States": {
            "InvokeScheduler": {
              "Type": "Task",
              "Resource": "arn:aws:states:::lambda:invoke",
              "OutputPath": "$.Payload",
              "Parameters": {
                "FunctionName": "${MetagraphTaskScheduler.Arn}",
                "Payload": {
                  "arn.$": "$$.Execution.Id",
                  "query.$": "$"
                }
              },
              "Next": "RunQueries"
            },
            "RunQueries": {
              "Type": "Map",
              "ItemsPath": "$.tasks",
              "MaxConcurrency": 150,
              "ItemBatcher": {
                "MaxItemsPerBatch": 1
              },
              "ItemProcessor": {
                "ProcessorConfig": {
                  "Mode": "DISTRIBUTED",
                  "ExecutionType": "STANDARD"
                },
                "StartAt": "SubmitQuery",
                "States": {
                  "SubmitQuery": {
                    "Type": "Task",
                    "Resource": "arn:aws:states:::batch:submitJob.sync",
                    "Parameters": {
                      "JobName.$": "$.Items[0].jobName",
                      "JobQueue": "${MetagraphBatchJobQueue}",
                      "JobDefinition": "${MetagraphBatchJobDefinition}",
                      "ContainerOverrides": {
                        "Memory.$": "$.Items[0].memory_mib",
                        "Environment": [
                          { "Name": "DATASET", "Value.$": "$.Items[0].dataset" },
                          { "Name": "GRAPH_FILE", "Value.$": "$.Items[0].graph_file" },
                          { "Name": "ANNO_FILE", "Value.$": "$.Items[0].anno_file" },
                          { "Name": "QUERY", "Value.$": "$.Items[0].query" },
                          { "Name": "JOB_ID", "Value.$": "$.Items[0].job_id" },
                          { "Name": "QUERY_MODE", "Value.$": "$.Items[0].query_mode" },
                          { "Name": "NUM_TOP_LABELS", "Value.$": "$.Items[0].num_top_labels" },
                          { "Name": "MIN_KMERS_FRACTION_LABEL", "Value.$": "$.Items[0].min_kmers_fraction_label" },
                          { "Name": "MIN_KMERS_FRACTION_GRAPH", "Value.$": "$.Items[0].min_kmers_fraction_graph" },
                          { "Name": "MANAGED_BY_AWS", "Value": "STARTED_BY_STEP_FUNCTIONS" }
                        ]
                      }
                    },
                    "Retry": [
                      {
                        "ErrorEquals": ["Batch.TooManyRequestsException", "States.TaskFailed"],
                        "IntervalSeconds": 2,
                        "BackoffRate": 2.0,
                        "MaxAttempts": 3
                      }
                    ],
                    "Catch": [
                      {
                        "ErrorEquals": ["States.ALL"],
                        "Next": "MarkFailed"
                      }
                    ],
                    "Next": "MarkSucceeded"
                  },
                  "MarkSucceeded": {
                    "Type": "Pass",
                    "ResultPath": "$",
                    "Parameters": {
                      "jobName.$": "$.JobName",
                      "status": "SUCCEEDED"
                    },
                    "End": true
                  },
                  "MarkFailed": {
                    "Type": "Pass",
                    "ResultPath": "$",
                    "Parameters": {
                      "jobName.$": "$.JobName",
                      "status": "FAILED"
                    },
                    "End": true
                  }
                }
              },
              "ResultPath": "$.tasks",
              "Next": "GatherResults"
            },
            "GatherResults": {
              "Type": "Task",
              "Resource": "arn:aws:states:::batch:submitJob.sync",
              "Parameters": {
                "JobName.$": "States.Format('Merge-{}', $.job_id)",
                "JobQueue": "${MetagraphBatchJobQueue}",
                "JobDefinition": "${MetagraphMergeJobDefinition}",
                "ContainerOverrides": {
                  "Environment": [
                    { "Name": "MANAGED_BY_AWS", "Value": "STARTED_BY_STEP_FUNCTIONS" },
                    { "Name": "JOB_ID", "Value.$": "$.job_id" },
                    { "Name": "QUERY_MODE", "Value.$": "$.query_mode" },
                    { "Name": "PYTHONUNBUFFERED", "Value": "1" }
                  ]
                }
              },
              "ResultPath": null,
              "Next": "NotifyUser"
            },
            "NotifyUser": {
              "Type": "Task",
              "Resource": "arn:aws:states:::lambda:invoke",
              "Parameters": {
                "FunctionName": "${MetagraphNotifier.Arn}",
                "Payload.$": "$"
              },
              "End": true
            }
          }
        }

Outputs:
  BucketName:
    Description: 'S3 Bucket for storing queries and their results.'
    Value: !Ref MetagraphBucket
  StateMachineArn:
    Description: 'ARN of the Step Functions state machine to trigger queries.'
    Value: !Ref MetagraphStateMachine