Support govcloud

build/generate_ami_mapping.go

@@ -169,10 +169,10 @@ func SupportedRegions() []string {
 		RegionMESouth1,
 		RegionSAEast1,
 		RegionAFSouth1,
+		RegionUSGovWest1,
+		RegionUSGovEast1,
 		// RegionCNNorthwest1,
 		// RegionCNNorth1,
-		// RegionUSGovWest1,
-		// RegionUSGovEast1,
 	}
 }
 
@@ -200,10 +200,41 @@ func EKSResourceAccountID(region string) string {
 }
 
 func main() {
-	destFile := mustExtractArg()
+	if len(os.Args) > 3 {
+		fmt.Println("usage: go run generate_ami_mapping.go <abs_dest_path> [only-govcloud]")
+		os.Exit(1)
+	}
+
+	destFile := os.Args[1]
+
+	onlyGovcloud := false
+
+	if len(os.Args) == 3 {
+		if os.Args[2] == "only-govcloud" {
+			onlyGovcloud = true
+		}
+	}
+
+	k8sVersionMap := map[string]map[string]map[string]string{}
+
+	if _, err := os.Stat(destFile); !os.IsNotExist(err) {
+		jsonBytes, err := ioutil.ReadFile(destFile)
+		if err != nil {
+			log.Fatal(err.Error())
+		}
+		json.Unmarshal(jsonBytes, &k8sVersionMap)
+	}
+
 	k8sVersion := "1.18"
-	regions := map[string]map[string]string{}
+
+	if k8sVersionMap[k8sVersion] == nil {
+		k8sVersionMap[k8sVersion] = map[string]map[string]string{}
+	}
 	for _, region := range SupportedRegions() {
+		if onlyGovcloud != (region == RegionUSGovEast1 || region == RegionUSGovWest1) {
+			// onlyGovcloud xor (region is us govclouds)
+			continue
+		}
 		fmt.Print(region)
 		sess := session.New(&aws.Config{Region: aws.String(region)})
 		svc := ec2.New(sess)
@@ -215,15 +246,17 @@ func main() {
 		if err != nil {
 			log.Fatal(err.Error())
 		}
-		regions[region] = map[string]string{
+
+		if k8sVersionMap[k8sVersion][region] == nil {
+			k8sVersionMap[k8sVersion][region] = map[string]string{}
+		}
+		k8sVersionMap[k8sVersion][region] = map[string]string{
 			"cpu":         cpuAMI,
 			"accelerated": acceleratedAMI,
 		}
 		fmt.Println(" ✓")
 	}
 
-	k8sVersionMap := map[string]interface{}{}
-	k8sVersionMap[k8sVersion] = regions
 	marshalledBytes, err := json.MarshalIndent(k8sVersionMap, "", "\t")
 	if err != nil {
 		log.Fatal(err.Error())
@@ -288,12 +321,3 @@ func FindImage(ec2api ec2iface.EC2API, ownerAccount, namePattern string) (string
 
 	return *output.Images[0].ImageId, nil
 }
-
-func mustExtractArg() string {
-	if len(os.Args) != 2 {
-		fmt.Println("usage: go run generate_ami_mapping.go <abs_dest_path>")
-		os.Exit(1)
-	}
-
-	return os.Args[1]
-}

dev/create_user.py

@@ -37,14 +37,18 @@
 
 user_name = f"dev-{cluster_name}-{cortex_region}"
 
-iam_client = boto3.client("iam")
+iam_client = boto3.client("iam", region_name=cortex_region)
 
 try:
     iam_client.get_user(UserName=user_name)
 except iam_client.exceptions.NoSuchEntityException:
     iam_client.create_user(UserName=user_name)
 
-policy_arn = f"arn:aws:iam::{account_id}:policy/{user_name}"
+partition = "aws"
+if "us-gov" in cortex_region:
+    partition = "aws-us-gov"
+
+policy_arn = f"arn:{partition}:iam::{account_id}:policy/{user_name}"
 
 try:
     iam_client.get_policy(PolicyArn=policy_arn)

dev/minimum_aws_policy.json

@@ -43,10 +43,10 @@
                 "logs:CreateLogGroup"
             ],
             "Resource": [
-                "arn:aws:ssm:*:$CORTEX_ACCOUNT_ID:parameter/aws/*",
-                "arn:aws:ssm:*::parameter/aws/*",
-                "arn:aws:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/*"
+                "arn:*:ssm:*:$CORTEX_ACCOUNT_ID:parameter/aws/*",
+                "arn:*:ssm:*::parameter/aws/*",
+                "arn:*:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/*"
             ]
         },
         {
@@ -81,12 +81,12 @@
                 "iam:GetRolePolicy"
             ],
             "Resource": [
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:instance-profile/eksctl-*",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/eksctl-*",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/aws-service-role/eks-nodegroup.amazonaws.com/AWSServiceRoleForAmazonEKSNodegroup",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/eksctl-managed-*",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:oidc-provider/*",
-                "arn:aws:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME:*"
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:instance-profile/eksctl-*",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/eksctl-*",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/aws-service-role/eks-nodegroup.amazonaws.com/AWSServiceRoleForAmazonEKSNodegroup",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/eksctl-managed-*",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:oidc-provider/*",
+                "arn:*:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME:*"
             ]
         },
         {
@@ -99,7 +99,7 @@
                 "iam:CreatePolicyVersion",
                 "iam:DeletePolicyVersion"
             ],
-            "Resource": "arn:aws:iam::$CORTEX_ACCOUNT_ID:policy/cortex-*"
+            "Resource": "arn:*:iam::$CORTEX_ACCOUNT_ID:policy/cortex-*"
         },
         {
             "Effect": "Allow",
@@ -123,17 +123,17 @@
         {
             "Effect": "Allow",
             "Action": "sqs:*",
-            "Resource": "arn:aws:sqs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:cx-*"
+            "Resource": "arn:*:sqs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:cx-*"
         },
         {
             "Effect": "Allow",
             "Action": "s3:*",
-            "Resource": "arn:aws:s3:::$CORTEX_CLUSTER_NAME*"
+            "Resource": "arn:*:s3:::$CORTEX_CLUSTER_NAME*"
         },
         {
             "Effect": "Allow",
             "Action": "s3:*",
-            "Resource": "arn:aws:s3:::$CORTEX_CLUSTER_NAME*/*"
+            "Resource": "arn:*:s3:::$CORTEX_CLUSTER_NAME*/*"
         }
     ]
 }

dev/registry.sh

@@ -77,7 +77,19 @@ function registry_login() {
   if [ "$is_registry_logged_in" = "false" ]; then
     blue_echo "Logging in to ECR..."
     aws ecr get-login-password --region $AWS_REGION | docker login --username AWS --password-stdin $registry_push_url
-    aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 790709498068.dkr.ecr.us-west-2.amazonaws.com  # this is for the inferentia device plugin image
+
+    blue_echo "Logging in to 790709498068.dkr.ecr.us-west-2.amazonaws.com for inferentia..."
+    set +e
+    echo "$AWS_REGION" | grep  "us-gov"
+    is_gov_cloud=$?
+    set -e
+    if [ "$is_gov_cloud" == "0" ]; then
+      # set NORMAL_REGION_AWS_ACCESS_KEY_ID and NORMAL_REGION_AWS_SECRET_ACCESS_KEY credentials from a regular AWS account (non govcloud) in your dev/config/env.sh
+      AWS_ACCESS_KEY_ID=$NORMAL_REGION_AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY=$NORMAL_REGION_AWS_SECRET_ACCESS_KEY aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 790709498068.dkr.ecr.us-west-2.amazonaws.com
+    else
+      aws ecr get-login-password --region us-west-2 | docker login --username AWS --password-stdin 790709498068.dkr.ecr.us-west-2.amazonaws.com
+    fi
+
     is_registry_logged_in="true"
     green_echo "Success\n"
   fi

docs/clusters/management/auth.md

@@ -108,10 +108,10 @@ Replace the following placeholders with their respective values in the policy te
                 "logs:CreateLogGroup"
             ],
             "Resource": [
-                "arn:aws:ssm:*:$CORTEX_ACCOUNT_ID:parameter/aws/*",
-                "arn:aws:ssm:*::parameter/aws/*",
-                "arn:aws:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/*"
+                "arn:*:ssm:*:$CORTEX_ACCOUNT_ID:parameter/aws/*",
+                "arn:*:ssm:*::parameter/aws/*",
+                "arn:*:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/*"
             ]
         },
         {
@@ -146,12 +146,12 @@ Replace the following placeholders with their respective values in the policy te
                 "iam:GetRolePolicy"
             ],
             "Resource": [
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:instance-profile/eksctl-*",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/eksctl-*",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/aws-service-role/eks-nodegroup.amazonaws.com/AWSServiceRoleForAmazonEKSNodegroup",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:role/eksctl-managed-*",
-                "arn:aws:iam::$CORTEX_ACCOUNT_ID:oidc-provider/*",
-                "arn:aws:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME:*"
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:instance-profile/eksctl-*",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/eksctl-*",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/aws-service-role/eks-nodegroup.amazonaws.com/AWSServiceRoleForAmazonEKSNodegroup",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:role/eksctl-managed-*",
+                "arn:*:iam::$CORTEX_ACCOUNT_ID:oidc-provider/*",
+                "arn:*:logs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:log-group:$CORTEX_CLUSTER_NAME:*"
             ]
         },
         {
@@ -164,7 +164,7 @@ Replace the following placeholders with their respective values in the policy te
                 "iam:CreatePolicyVersion",
                 "iam:DeletePolicyVersion"
             ],
-            "Resource": "arn:aws:iam::$CORTEX_ACCOUNT_ID:policy/cortex-*"
+            "Resource": "arn:*:iam::$CORTEX_ACCOUNT_ID:policy/cortex-*"
         },
         {
             "Effect": "Allow",
@@ -188,17 +188,17 @@ Replace the following placeholders with their respective values in the policy te
         {
             "Effect": "Allow",
             "Action": "sqs:*",
-            "Resource": "arn:aws:sqs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:cx-*"
+            "Resource": "arn:*:sqs:$CORTEX_REGION:$CORTEX_ACCOUNT_ID:cx-*"
         },
         {
             "Effect": "Allow",
             "Action": "s3:*",
-            "Resource": "arn:aws:s3:::$CORTEX_CLUSTER_NAME*"
+            "Resource": "arn:*:s3:::$CORTEX_CLUSTER_NAME*"
         },
         {
             "Effect": "Allow",
             "Action": "s3:*",
-            "Resource": "arn:aws:s3:::$CORTEX_CLUSTER_NAME*/*"
+            "Resource": "arn:*:s3:::$CORTEX_CLUSTER_NAME*/*"
         }
     ]
 }

manager/generate_eks.py

@@ -22,14 +22,17 @@
 # kubelet config schema:
 # https://github.com/kubernetes/kubernetes/blob/master/staging/src/k8s.io/kubelet/config/v1beta1/types.go
 def default_nodegroup(cluster_config):
+    partition = "aws"
+    if "us-gov" in cluster_config["region"]:
+        partition = "aws-us-gov"
     return {
         "iam": {
             "withAddonPolicies": {"autoScaler": True},
             "attachPolicyARNs": [
-                "arn:aws:iam::aws:policy/AmazonEKSWorkerNodePolicy",
-                "arn:aws:iam::aws:policy/AmazonEKS_CNI_Policy",
-                "arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly",
-                "arn:aws:iam::aws:policy/ElasticLoadBalancingFullAccess",
+                f"arn:{partition}:iam::aws:policy/AmazonEKSWorkerNodePolicy",
+                f"arn:{partition}:iam::aws:policy/AmazonEKS_CNI_Policy",
+                f"arn:{partition}:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly",
+                f"arn:{partition}:iam::aws:policy/ElasticLoadBalancingFullAccess",
                 cluster_config["cortex_policy_arn"],
             ]
             + cluster_config.get("iam_policy_arns", []),

manager/install.sh

@@ -208,7 +208,7 @@ function setup_configmap() {
   kubectl -n=default create configmap 'env-vars' \
     --from-literal='CORTEX_VERSION'=$CORTEX_VERSION \
     --from-literal='CORTEX_REGION'=$CORTEX_REGION \
-    --from-literal='AWS_REGION'=$CORTEX_REGION \
+    --from-literal='AWS_DEFAULT_REGION'=$CORTEX_REGION \
     --from-literal='CORTEX_TELEMETRY_DISABLE'=$CORTEX_TELEMETRY_DISABLE \
     --from-literal='CORTEX_TELEMETRY_SENTRY_DSN'=$CORTEX_TELEMETRY_SENTRY_DSN \
     --from-literal='CORTEX_TELEMETRY_SEGMENT_WRITE_KEY'=$CORTEX_TELEMETRY_SEGMENT_WRITE_KEY \

manager/manifests/ami.json

@@ -41,8 +41,8 @@
 			"cpu": "ami-0f85d2eeb0bea62a7"
 		},
 		"eu-north-1": {
-			"accelerated": "ami-016de826d1d553d1b",
-			"cpu": "ami-05dc6dcd932a8159e"
+			"accelerated": "ami-05bc92c8d24c7a661",
+			"cpu": "ami-09f6f77efde8920b8"
 		},
 		"eu-south-1": {
 			"accelerated": "ami-0ce76bb81c438e3b9",
@@ -76,6 +76,14 @@
 			"accelerated": "ami-0da728680ca5ee572",
 			"cpu": "ami-0576aabae1709e005"
 		},
+		"us-gov-east-1": {
+			"accelerated": "ami-0d0129213a7f16d12",
+			"cpu": "ami-090c13eee66020d9b"
+		},
+		"us-gov-west-1": {
+			"accelerated": "ami-06439eec0732c6baa",
+			"cpu": "ami-04bd686c969f13917"
+		},
 		"us-west-1": {
 			"accelerated": "ami-085fdad868b9007dd",
 			"cpu": "ami-09bec0a8c8d4925a6"

pkg/cortex/serve/cortex_internal/serve/serve.py

@@ -265,7 +265,7 @@ def start_fn():
 
     model_dir = os.getenv("CORTEX_MODEL_DIR")
     cache_dir = os.getenv("CORTEX_CACHE_DIR")
-    region = os.getenv("AWS_REGION")
+    region = os.getenv("AWS_DEFAULT_REGION")
 
     tf_serving_port = os.getenv("CORTEX_TF_BASE_SERVING_PORT", "9000")
     tf_serving_host = os.getenv("CORTEX_TF_SERVING_HOST", "localhost")

pkg/cortex/serve/init/script.py

@@ -121,7 +121,7 @@ def main():
     # get API spec
     spec_path = os.environ["CORTEX_API_SPEC"]
     cache_dir = os.getenv("CORTEX_CACHE_DIR")
-    region = os.getenv("AWS_REGION")  # when it's deployed to AWS
+    region = os.getenv("AWS_DEFAULT_REGION")  # when it's deployed to AWS
 
     with open(spec_path) as json_file:
         api_spec = json.load(json_file)

pkg/cortex/serve/start/async_api.py

@@ -104,7 +104,7 @@ def main():
     workload_path = os.environ["CORTEX_ASYNC_WORKLOAD_PATH"]
     project_dir = os.environ["CORTEX_PROJECT_DIR"]
     readiness_file = os.getenv("CORTEX_READINESS_FILE", "/mnt/workspace/api_readiness.txt")
-    region = os.getenv("AWS_REGION")
+    region = os.getenv("AWS_DEFAULT_REGION")
     queue_url = os.environ["CORTEX_QUEUE_URL"]
     statsd_host = os.getenv("HOST_IP")
     statsd_port = os.getenv("CORTEX_STATSD_PORT", "9125")

pkg/cortex/serve/start/batch.py

@@ -180,7 +180,7 @@ def start():
     tf_serving_port = os.getenv("CORTEX_TF_BASE_SERVING_PORT", "9000")
     tf_serving_host = os.getenv("CORTEX_TF_SERVING_HOST", "localhost")
 
-    region = os.getenv("AWS_REGION")
+    region = os.getenv("AWS_DEFAULT_REGION")
 
     has_multiple_servers = os.getenv("CORTEX_MULTIPLE_TF_SERVERS")
     if has_multiple_servers:

pkg/cortex/serve/start/server_grpc.py

@@ -113,7 +113,7 @@ def init():
 
     model_dir = os.getenv("CORTEX_MODEL_DIR")
     cache_dir = os.getenv("CORTEX_CACHE_DIR")
-    region = os.getenv("AWS_REGION")
+    region = os.getenv("AWS_DEFAULT_REGION")
 
     tf_serving_port = os.getenv("CORTEX_TF_BASE_SERVING_PORT", "9000")
     tf_serving_host = os.getenv("CORTEX_TF_SERVING_HOST", "localhost")

pkg/cortex/serve/start/task.py

@@ -30,7 +30,7 @@
 def start():
     cache_dir = os.environ["CORTEX_CACHE_DIR"]
     project_dir = os.environ["CORTEX_PROJECT_DIR"]
-    region = os.getenv("AWS_REGION")
+    region = os.getenv("AWS_DEFAULT_REGION")
 
     api_spec_path = os.environ["CORTEX_API_SPEC"]
     task_spec_path = os.environ["CORTEX_TASK_SPEC"]

pkg/lib/aws/errors.go

@@ -190,7 +190,7 @@ func ErrorDashboardHeightOutOfRange(height int) error {
 func ErrorRegionNotConfigured() error {
 	return errors.WithStack(&errors.Error{
 		Kind:    ErrRegionNotConfigured,
-		Message: "aws region has not been configured; please set a default region (e.g. `export AWS_REGION=us-west-2`)",
+		Message: "aws region has not been configured; please set a default region (e.g. `export AWS_DEFAULT_REGION=us-west-2`)",
 	})
 }
 

pkg/lib/aws/gen_resource_metadata.py

@@ -18,7 +18,6 @@
 
 # https://docs.aws.amazon.com/general/latest/gr/eks.html
 # China regions don't seem to support these endpoints (yet?)
-# GovCloud is skipped
 REGIONS = [
     "us-east-2",  # Ohio
     "us-east-1",  # N. Virginia
@@ -41,6 +40,8 @@
     "eu-north-1",  # Stockholm
     "me-south-1",  # Bahrain
     "sa-east-1",  # Sao Paulo
+    "us-gov-east-1",  # GovCloud US-East
+    "us-gov-west-1",  # GovCloud US-West
 ]
 
 OUTPUT_FILE_NAME = "resource_metadata.go"