Building a multi-tenant Kubeflow environment on Amazon EKS using Amazon Cognito and ADFS

• An AWS account
• Amazon Certificate Manager (ACM) certificate and a valid domain
• Knowledge about Amazon Cognito, Identity providers and Federation and Kubernetes/Amazon EKS
• eksctl
• Minimal knowledge in other AWS services like Amazon Route 53, and AWS Lambda
• Kubeflow basic administration knowledge
• ADFS Claims configured to support the Amazon Cognito Configuration
• Terraform v1.1.5

1. User opens the custom URL for the Kubeflow ALB
2. Route 53 resolves the FQDN to the ALB
3. ALB redirects the request to the Kubeflow containers running on EKS
4. IstioGateway intercepts the request
5. IstioGateway redirects the request to Amazon Cognito
6. Amazon Cognito App client is configured with a custom Domain and valid certificate from Amazon Certificate Manager
7. Amazon Cognito redirects the request to the configured identify provider (IdP) (ADFS in this example)
8. User logs in and ADFS redirects to Amazon Cognito with the SAML assertion
9. Amazon Cognito has the Pre-Token generation event configured to trigger the Lambda function
10. “Profile Manager” Lambda function is triggered with the groups that the user belongs in the AD as part of the event
11. The Lambda function, based on the groups received on the event will connect to the Amazon EKS and generate the necessary role and serviceRoleBindings
12. Amazon Cognito redirect back to the Kubeflow UI

1. We need to create to own a valid domain and create the necessary certificates to be used in later steps by the Amazon Cognito Domain URL and the ALB for Kubeflow (Step One – AWS ACM)
2. As part of the integration, Kubeflow doesn’t support the automatic creation of namespaces and user groups. To enable this functionality, we will create a new Lambda function (we called it Profile Manager) that will be triggered as part of the login flow by Amazon Cognito. The steps and code to support this implementation can be found below (Step Two – Namespace Automation)
3. Once the Microsoft ADFS is configured, we can now proceed to configure Amazon Cognito as the Identity Provider for Kubeflow. The workflow details and configuration to achieve this can be found in the Amazon Cognito IdP Integration section. (Step Three – Amazon Cognito IdP Integration Flow)
4. After the IdP is set up, we can proceed with the creation of the Amazon EKS Cluster. (Step Four – Amazon EKS Cluster Creation)
5. After the Cluster is created and the nodes are in service. we can proceed with the Kubeflow installation. (Step Five – Kubeflow Integration)
6. As a last step, and as we create the LoadBalancer with the Kubeflow deployed, we need to get the URL of the LoadBalancer, create the DNS record in Amazon Route 53 (we assume the certificate for this alias exists and is valid), and configure the Callback URL in Amazon Cognito. This lets the Login workflow know where to redirect (Kubeflow Entrypoint) after the authentication is successful. (Step Six – Configure CallBack on Amazon Cognito)
7. As all the steps are now complete, we can log in to Kubeflow using the DNS entry created in the previous step. (Step Seven – Log in to Kubeflow)

1. Make sure you own a domain and it’s registered in Route 53 (we have provided a sample domain below).
   1. 
2. Go to the AWS Certificate Manager console and select Request a Certificate in the left pane.
   1. 
3. Add the following domains to the certificate:
   1. <mydomain>
   2. kubeflow.<mydomain>
      1. Needed for the ALB in a later step
   3. auth.<mydomain>
      1. Needed by Cognito to validate the sub-domain afterwards
   4. cognito.auth.<mydomain>
      1. Needed for the Custom Domain in Amazon Cognito
   5. 
4. Press “Request” and wait a few minutes for the certificate to be validated in Route 53.
5. Once the certificate is valid, you should see something similar to the image below.
   1. 

• User tries to log in to Kubeflow.
• The Amazon Cognito flow will start and contact the ADFS to validate the user credentials.
• An Amazon Cognito trigger for “Pre-Auth token” will trigger the Lambda function KF_profile_Manager
• This Lambda function will:
  • Create user owner profile with the name “profile-<username>” as this is mandatory for Kubeflow. This profile has no resource permissions, hence nothing can be created.
  • Check if the group namespaces exist and/or create them.
  • Based on groups from the ADFS, named with a certain prefix, will create the rolebindings and the serviceRolebindings to add the user to the respective groups namespaces.

• Lambda function code can be found in GitHub
• Go to the AWS Lambda Console and create a new Function with Python3.9 as interpreter
• Timeout should be increased to 30 seconds and Memory to 256 MB
• handler: index.lambda_handler
• The function expects the following environment variables:
  • The Lambda Environment Variable ClusterName, needs to be updated to the correct EKS cluster name
  • Configure the following environment variables:
    • clusterName = <clusterName>
      cognitoGroupField = “custom:groups”
• Additionally, we need to make sure that Amazon Cognito has permissions to invoke the Lambda function.
  • Go to the AWS Lambda function, and select the Configuration Tab
  • Select Permissions from the left hand menu
  • And from the Resource-based policy, click on Add Permissions
• 
• The Lambda function is now ready to be used by Amazon Cognito to create the namespaces and the respective group membership in Kubeflow.

• Role Binding Function
  • Role Binding grants the permissions defined in a role to a user or set of users. It holds a list of subjects (users, groups, or service accounts), and a reference to the role being granted.
  • Role Binding grants permissions within a specific namespace whereas a ClusterRoleBinding grants that access cluster-wide.
• Service Role Binding
  • Istio needs a service role binding to allow the user to have permissions in a certain namespace
  • A ServiceRoleBinding specification includes two parts:
    • The roleRef field that refers to a ServiceRole object in the same namespace.
    • A list of subjects that are assigned the roles.

1. Create a User Pool
   1. Add the following attributes:
      1. Select Username and allow validated email as secondary login
      2. In the attribute section select email as mandatory
      3. In the custom attributes create the following:
         1. Name: groups
         2. Type: String
         3. Min Length: 1
         4. Max Length: 2048
      4. 
2. Create an APP Client
   1. Give a name to the App Client
   2. Enable Security
   3. Select Generate client secret
   4. Press Create app client
   5. 
3. In the triggers section, choose a “Pre Token Generation” trigger to point the Lambda function that will manage the profiles in Kubeflow (Lambda function in this section is the pre-condition to select this option)
   1. 
4. Create a SAML Identity Provider with the following mappings:
   1. To create the SAML Provider, you will need to use the metadata file from your Microsoft ADFS environment (or alternatively, use a URL in which you can download this file)
   2. 
   3. Edit the attribute mappings and add the following:
      1. https://aws.amazon.com/SAML/Attributes/validation → Email Verified
      2. https://aws.amazon.com/SAML/Attributes/Role → custom:groups
      3. http://schemas.xmlsoap.org/ws/2005/05/identity/claims/emailaddress → emailaddress
   4. 
5. Go back to the APP client and select the following options:
   1. Enable Identity providers
      1. ADFS
   2. Under Oauth 2.0, enable “Authorization code grant” under Allowed OAuth Flows, and then select
      1. Email
   3. Callback URL is the Route 53 Domain domain selected (Domain that was used in the second step for the certificate creation). This will also be used in a later step and Alias created for the kubeflow ALB
      1. IE: https://kubeflow.<mydomain>/oauth2/idpresponse
   4. Sign Out URL is optional
      1. This can be built based on the following template “https://<CognitoCustomDomainURL>/logout?client_id=<APPID>&logout_uri=https://<kubeflowURL>/logout”
      2. CognitoCustomDomainURL is the url from the customDomain Creation
      3. APPID is the Cognito APPID, identified on the top of the screenshot below
      4. kubeflowURL is the Alias created for the ALB of Kubeflow
   5. 
6. As a final step, we need to update the ADFS trust relation with the Amazon Cognito ARN in the Microsoft ADFS server
   1. Go to your ADFS server
   2. Go to the Federation Console and open the Relaying Party trust
   3. 
   4. Make sure that the identifier is correct
   5. The region and the Amazon Cognito Pool ID should be updated to point to the one just created in this section
   6. 
   7. Make sure also, that the Endpoints have the correct endpoint for your domain
   8. It should point to https://cognito.auth.<domain>/saml2/idpresponse
   9. 

1. Create a YAML file using the following command:

   1. cat << EOF > cluster.yaml
      ---
      apiVersion: eksctl.io/v1alpha5
      kind: ClusterConfig
      
      metadata:
        name: <myClusterName>
        version: 1.18
        region: us-east-1
      
      vpc:
        subnets:
          private:
            us-east-1a: { id: xxx }
            us-east-1b: { id: xxx }
          public:
            us-east-1a: { id: xxx }
            us-east-1b: { id: xxx }
      
      managedNodeGroups:
        - name: kubeflow-mng
          desiredCapacity: 3
          instanceType: m5.large
      EOF

2. Run the eksctl command to create the Amazon EKS cluster
   • eksctl create cluster -f cluster.yaml
3. Wait until the cluster and nodes are created (usually takes around 15 minutes) before proceeding to the next step.

#!/bin/bash
export CONFIG_URI="https://raw.githubusercontent.com/kubeflow/manifests/v1.2-branch/kfdef/kfctl_aws_cognito.v1.2.0.yaml"
# 3. Set an environment variable for your AWS cluster name.
export AWS_CLUSTER_NAME=kubeflowworkshop-eksctl

# 4. Create the directory you want to store deployment, this has to be ${AWS_CLUSTER_NAME}
mkdir ${AWS_CLUSTER_NAME} && cd ${AWS_CLUSTER_NAME}

# 5. Download your configuration files, so that you can customize the configuration before deploying Kubeflow.
wget -O kfctl_aws.yaml $CONFIG_URI

./install_kubeflow.sh

apiVersion: kfdef.apps.kubeflow.org/v1 kind: KfDef metadata:   namespace: kubeflow spec:   ...   plugins:   - kind: KfAwsPlugin     metadata:       creationTimestamp: null       name: aws     spec:       auth:         cognito:           certArn: <valid ACM certificate >           cognitoAppClientId: <CognitoAppClientID>           cognitoUserPoolArn: <CognitoUserPoolArn>           cognitoUserPoolDomain: <CognitoUserPoolDomain>       enablePodIamPolicy: false       region: <Region>       roles:       - <EksClusteRole>   repos:   ...

• • • CertARN:
      • Certificate arn from the Certificate created in the section
      • arn:aws:acm:<region>:<accountID>:certificate/<acm cert id>
    • cognitoAppClientId:
      • App Client ID from Amazon Cognito
    • cognitoUserPoolArn:
      • Arn from the Amazon Cognito User Pool
      • arn:aws:cognito-idp:<region>:<accountID>:userpool/<PoolID>
    • cognitoUserPoolDomain:
      • Domain used for the Custom Domain in cognito
      • IE: cognito.auth.<mydomain>
    • region: <Region>
    • EKSClusterRole:
      • This is the role of the worker nodes from the Amazon EKS cluster
      • Can be found by issuing the following command:

aws iam list-roles \
    | jq -r ".Roles[] \
    | select(.RoleName \
    | startswith(\"eksctl-$AWS_CLUSTER_NAME\") and contains(\"NodeInstanceRole\")) \
    .RoleName"

export AWS_CLUSTER_NAME=eksworkshop-eksctl
export KF_NAME=${AWS_CLUSTER_NAME}

export BASE_DIR=/Users/<user>/environment
export KF_DIR=${BASE_DIR}/${KF_NAME}

export CONFIG_FILE=${KF_DIR}/kfctl_aws.yaml <- This is the file configured before

CD ${BASE_DIR}
kfctl apply -V -f kfctl_aws.yaml

kubectl get pods -A

1. 1. Get the URL for the callbackkubectl get ingress istio-ingress -n istio-system
      

1. 2. Create a Route 53 CNAME for the ELB
      • This should match the certificate created on AWS Certificate Manager.
      • Copy the “ADDRESS” from the terminal and go to Amazon Cognito App Client Settings to update the callback URL
      • ALB address should be an alias to kubeflow.<mydomain>

Step Seven – Log in to Kubeflow

Now that we have finished the steps, we can test by logging into our Kubeflow environment. In a new browser tab, type in the address bar the Route 53 CNAME that was created for ELB (in our example above, this is kubeflow.<mydomain).

During the login process, Kubeflow will redirect us to Cognito, which will then redirect to the ADFS Login Page. After filling the credentials, ADFS will redirect back to Amazon Cognito.

Before Amazon Cognito redirects us to the Kubeflow URL, the Lambda function will trigger (“Pre Token Generation”) and create the namespaces in Amazon EKS and add our membership.

Finally, we will end in the Kubeflow Main Page with access to the shared namespaces(groups), allowing team collaboration on resources.

Central Dashboard

Cleaning up

Cleanup is the exact reverse steps:

Conclusion

Kubeflow is a great engine for ML workflows on Kubernetes. It’s simple, portable, and scalable. In this post we walked you through how to extend your Kubeflow environment to achieve multi-tenancy using Amazon Cognito and Microsoft ADFS as an Identity source provider. While we used Microsoft IdP as the IdP in this post, any other third-party solution that can integrate with Amazon Cognito can be used, IE: OpenSSO.