feat: Add mlflow-ai-gw-image module and update mlflow to 2.17.0

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## UNRELEASED
 
 ### **Added**
+- added `mlflow-ai-gw-image` module
 
 ### **Changed**
 

README.md

@@ -54,10 +54,11 @@ End-to-end example use-cases built using modules in this repository.
 
 ### Mlflow Modules
 
-| Type                                                                    | Description                                                                                                                                                                                       |
-|-------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| [Mlflow Image Module](modules/mlflow/mlflow-image/README.md)            | Creates Mlflow Docker container image and pushes the image to Elastic Container Registry                                                                                                          |
-| [Mlflow on AWS Fargate Module](modules/mlflow/mlflow-fargate/README.md) | Runs Mlflow container on AWS Fargate in a load-balanced Elastic Container Service. Supports Elastic File System and Relational Database Store for metadata persistence, and S3 for artifact store |
+| Type                                                                          | Description                                                                                                                                                                                       |
+|-------------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| [Mlflow Image Module](modules/mlflow/mlflow-image/README.md)                  | Creates Mlflow Tracing Server Docker image and pushes the image to Elastic Container Registry                                                                                                     |
+| [Mlflow on AWS Fargate Module](modules/mlflow/mlflow-fargate/README.md)       | Runs Mlflow container on AWS Fargate in a load-balanced Elastic Container Service. Supports Elastic File System and Relational Database Store for metadata persistence, and S3 for artifact store |
+| [Mlflow AI Gateway Image Module](modules/mlflow/mlflow-ai-gw-image/README.md) | Creates Mlflow AI Gateway Docker image and pushes the image to Elastic Container Registry                                                                                                         |
 
 ### FMOps/LLMOps Modules
 

manifests/mlflow-tracking/deployment.yaml

@@ -12,6 +12,8 @@ groups:
     path: manifests/mlflow-tracking/images-modules.yaml
   - name: mlflow
     path: manifests/mlflow-tracking/mlflow-modules.yaml
+  - name: mlflow-ai-gw
+    path: manifests/mlflow-tracking/images-ai-gw-modules.yaml
 targetAccountMappings:
   - alias: primary
     accountId:

manifests/mlflow-tracking/images-ai-gw-modules.yaml

@@ -0,0 +1,10 @@
+name: mlflow-ai-gw-image
+path: modules/mlflow/mlflow-ai-gw-image/
+targetAccount: primary
+parameters:
+  - name: ecr-repository-name
+    valueFrom:
+      moduleMetadata:
+        group: storage
+        name: ecr-mlflow-ai-gw
+        key: EcrRepositoryName

manifests/mlflow-tracking/storage-modules.yaml

@@ -11,6 +11,19 @@ parameters:
   - name: removal-policy
     value: DESTROY
 ---
+name: ecr-mlflow-ai-gw
+path: git::https://github.com/awslabs/idf-modules.git//modules/storage/ecr?ref=release/1.12.0&depth=1
+targetAccount: primary
+parameters:
+  - name: image-tag-mutability
+    value: MUTABLE
+  - name: image-scan-on-push
+    value: True
+  - name: encryption
+    value: KMS_MANAGED
+  - name: removal-policy
+    value: DESTROY
+---
 name: buckets
 path: git::https://github.com/awslabs/idf-modules.git//modules/storage/buckets?ref=release/1.12.0&depth=1
 targetAccount: primary

modules/mlflow/mlflow-ai-gw-image/README.md

@@ -0,0 +1,39 @@
+# Mlflow AI Gateway Image
+
+## Description
+
+This module creates a [MLFlow AI Gateway](https://mlflow.org/docs/latest/llms/index.html#id2) server container image and pushes to the specified Elastic Container Repository.
+
+## Inputs/Outputs
+
+### Input Parameters
+
+#### Required
+
+- `ecr-repository-name`: The name of the ECR repository to push the image to.
+
+### Sample manifest declaration
+
+```yaml
+name: mlflow-image
+path: modules/mlflow/mlflow-image
+parameters:
+  - name: ecr-repository-name
+    valueFrom:
+      moduleMetadata:
+        group: storage
+        name: ecr-mlflow-ai-gw
+        key: EcrRepositoryName
+```
+
+### Module Metadata Outputs
+
+- `MlflowAIGWImageUri`: Mlflow AI Gateway image URI
+
+#### Output Example
+
+```json
+{
+  "MlflowAIGWImageUri": "xxxxxxxxxxxx.dkr.ecr.us-east-1.amazonaws.com/ecr-mlflow-ai-gw:latest"
+}
+```

modules/mlflow/mlflow-ai-gw-image/deployspec.yaml

@@ -0,0 +1,21 @@
+publishGenericEnvVariables: true
+deploy:
+  phases:
+    build:
+      commands:
+        - aws ecr describe-repositories --repository-names ${SEEDFARMER_PARAMETER_ECR_REPOSITORY_NAME} || aws ecr create-repository --repository-name ${SEEDFARMER_PARAMETER_ECR_REPOSITORY_NAME} --image-scanning-configuration scanOnPush=true
+        - export COMMIT_HASH=$(echo $CODEBUILD_RESOLVED_SOURCE_VERSION | cut -c 1-7)
+        - export IMAGE_TAG=${COMMIT_HASH:=latest}
+        - export REPOSITORY_URI=$AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com/${SEEDFARMER_PARAMETER_ECR_REPOSITORY_NAME}
+        - aws ecr get-login-password --region $AWS_DEFAULT_REGION | docker login --username AWS --password-stdin $AWS_ACCOUNT_ID.dkr.ecr.$AWS_DEFAULT_REGION.amazonaws.com
+        - echo Building the Docker image...
+        - cd src/ && docker build -t $REPOSITORY_URI:latest .
+        - docker tag $REPOSITORY_URI:latest $REPOSITORY_URI:$IMAGE_TAG
+        - docker push $REPOSITORY_URI:$IMAGE_TAG
+        - seedfarmer metadata add -k ImageUri -v $REPOSITORY_URI:$IMAGE_TAG
+destroy:
+  phases:
+    build:
+      commands:
+        - aws ecr delete-repository --repository-name ${SEEDFARMER_PARAMETER_ECR_REPOSITORY_NAME} --force
+# build_type: BUILD_GENERAL1_LARGE

modules/mlflow/mlflow-ai-gw-image/modulestack.yaml

@@ -0,0 +1,42 @@
+AWSTemplateFormatVersion: 2010-09-09
+Description: This stack deploys a Module specific IAM permissions
+
+Parameters:
+  # DeploymentName:
+  #   Type: String
+  #   Description: The name of the deployment
+  # ModuleName:
+  #   Type: String
+  #   Description: The name of the Module
+  RoleName:
+    Type: String
+    Description: The name of the IAM Role
+  ECRRepositoryName:
+    Type: String
+    Description: The name of the ECR repository
+
+Resources:
+  Policy:
+    Type: "AWS::IAM::Policy"
+    Properties:
+      PolicyDocument:
+        Statement:
+          - Effect: Allow
+            Action:
+              - "ecr:Describe*"
+              - "ecr:Get*"
+              - "ecr:List*"
+            Resource: "*"
+          - Action:
+              - "ecr:Create*"
+              - "ecr:Delete*"
+              - "ecr:*LayerUpload"
+              - "ecr:UploadLayerPart"
+              - "ecr:Batch*"
+              - "ecr:Put*"
+            Effect: Allow
+            Resource:
+              - !Sub "arn:${AWS::Partition}:ecr:${AWS::Region}:${AWS::AccountId}:repository/${ECRRepositoryName}"
+        Version: 2012-10-17
+      PolicyName: "modulespecific-policy"
+      Roles: [!Ref RoleName]

modules/mlflow/mlflow-ai-gw-image/src/Dockerfile

@@ -0,0 +1,15 @@
+FROM public.ecr.aws/docker/library/python:3.10.15
+
+RUN pip install \
+    'mlflow[gateway]' \
+    mlflow==2.17.0 && \
+    mkdir /mlflow/
+
+EXPOSE 7000
+
+COPY config.yaml ./mlflow/config.yaml
+
+CMD mlflow gateway start \
+    --host 0.0.0.0 \
+    --port 7000 \
+    --config-path ./mlflow/config.yaml;

modules/mlflow/mlflow-ai-gw-image/src/config.yaml

@@ -0,0 +1,13 @@
+endpoints:
+  - name: chat
+    endpoint_type: llm/v1/chat
+    model:
+      provider: bedrock
+      name: anthropic.claude-3-5-sonnet-20240620-v1:0
+      config:
+        aws_config:
+          aws_region: 'AWS_REGION'
+          aws_role_arn: 'AWS_ROLE_ARN'
+      limit:
+        renewal_period: minute
+        calls: 10

modules/mlflow/mlflow-image/src/Dockerfile

@@ -1,7 +1,7 @@
 FROM python:3.10.12
 
 RUN pip install \
-    mlflow==2.16.0 \
+    mlflow==2.17.0 \
     pymysql==1.1.1 \
     boto3==1.34.45 && \
     mkdir /mlflow/