fix: add domain execution role fallback, additional domain tags (#4785)

Author: danabens

ml_ops/sm-datazone_import/README.md

@@ -4,7 +4,7 @@ This example contains python scripts to import an existing SageMaker Domain into
 
 ## Setup
 
-1. Add the byod service model
+1. Add the Bring-Your-Own-Domain (BYOD) service model
 
 ```bash
 aws configure add-model --service-model file://resources/datazone-linkedtypes-2018-05-10.normal.json --service-name datazone-byod
@@ -33,4 +33,9 @@ python import-sagemaker-domain.py \
     --region REGION \
     --federation-role ARN_OF_FEDERATION_ROLE \
     --account-id ACCOUNTID
-```
+```
+
+### Additional Configuration
+
+- SageMaker execution roles need DataZone API permissions in order for the Assets UI to function. See [DataZoneUserPolicy.json](./resources/DataZoneUserPolicy.json) for an example.
+- Ensure the DataZone Domain trusts SageMaker. In the AWS DataZone console navigate to Domain details and select the "Trusted services".

ml_ops/sm-datazone_import/import-sagemaker-domain.py

@@ -60,6 +60,9 @@ def _choose_sm_domain(self):
                 self.sm_domain_name, self.sm_domain_id
             )
         )
+        sm_domain = self.sm_client.describe_domain(DomainId=self.sm_domain_id)
+        self.auth_mode = sm_domain["AuthMode"]
+        self.default_execution_role = sm_domain["DefaultUserSettings"]["ExecutionRole"]
         return self.sm_domain_id
 
     def _choose_dz_domain(self):
@@ -107,15 +110,18 @@ def _choose_dz_project(self):
         )
         return self.dz_project_id
 
-    def _tag_dm_domain(self):
+    def _tag_sm_domain(self):
         # [3.5 Tagging] Before getting started on byod-e2e, ensure that CX has the SM domain and ExecutionRole's tagged accordingly.
         # 	1. Tag the SM domain by admin (DZ DomainId and tag the stage, and domainAccountId)
         # 	2. Tag the execution role with DZ domainId and projectId
 
+        # TODO - remove project/env tags once front end behavior is fixed
         domain_tag = {"Key": "AmazonDataZoneDomain", "Value": self.dz_domain_id}
+        project_tag = {"Key": "AmazonDataZoneProject", "Value": self.dz_project_id}
+        env_tag = {"Key": "AmazonDataZoneEnvironment", "Value": self.env_id}
         account_tag = {"Key": "AmazonDataZoneDomainAccount", "Value": self.account_id}
         stage_tag = {"Key": "AmazonDataZoneStage", "Value": self.stage}
-        sm_domain_tags = [domain_tag, account_tag, stage_tag]
+        sm_domain_tags = [domain_tag, project_tag, env_tag, account_tag, stage_tag]
         sm_domain_arn = "arn:aws:sagemaker:{}:{}:domain/{}".format(
             self.region, self.account_id, self.sm_domain_id
         )
@@ -156,9 +162,18 @@ def _map_users(self):
         self.sm_user_info = {}
         self.sm_user_info["name"] = sm_user_name
         self.sm_user_info["arn"] = sm_user_profile_full["UserProfileArn"]
-        self.sm_user_info["exec_role_arn"] = sm_user_profile_full["UserSettings"][
-            "ExecutionRole"
-        ]
+        exec_role = None
+        if "UserSettings" in sm_user_profile_full:
+            user_settings = sm_user_profile_full["UserSettings"]
+            if "ExecutionRole" in user_settings:
+                exec_role = user_settings["ExecutionRole"]
+
+        if exec_role is None:
+            print(f'User {sm_user_name} has no execution role set, using default from domain.')
+            exec_role = self.default_execution_role
+
+        self.sm_user_info["exec_role_arn"] = exec_role
+
         self.sm_user_info["id"] = sm_user_profile_full[
             "HomeEfsFileSystemUid"
         ]  # e.g. d-7d4uvydb9rcy
@@ -170,7 +185,7 @@ def _map_users(self):
 
         # get role name from arn "arn:aws:iam::047923724610:role/service-role/AmazonSageMaker-ExecutionRole-20241008T155288"
         role_name = self.sm_user_info["exec_role_arn"][
-            self.sm_user_info["exec_role_arn"].rfind("/") + 1:
+            self.sm_user_info["exec_role_arn"].rfind("/") + 1 :
         ]  # rfind searches from back of string
         self.iam_client.tag_role(RoleName=role_name, Tags=sm_exec_role_tags)
         print(
@@ -360,7 +375,7 @@ def _link_domain(self):
             {
                 "itemIdentifier": f"arn:aws:sagemaker:{self.region}:{self.account_id}:domain/{self.sm_domain_id}",
                 "itemType": "SAGEMAKER_DOMAIN",
-                "configuration": {"AuthMode": "NonSSO"},
+                "configuration": {"AuthMode": self.auth_mode},
                 "connectedEntities": [
                     {
                         "connectedEntityIdentifier": self.env_id,
@@ -464,10 +479,10 @@ def import_interactive(self):
         self._choose_sm_domain()
         self._choose_dz_domain()
         self._choose_dz_project()
-        self._tag_dm_domain()
+        self._configure_environment()
+        self._tag_sm_domain()
         self._map_users()
         self._configure_blueprint()
-        self._configure_environment()
         self._add_environment_action()
         self._associate_fed_role()
         self._link_domain()

ml_ops/sm-datazone_import/resources/DatazoneUserPolicy.json

@@ -0,0 +1,65 @@
+{
+    "Statement": [
+        {
+            "Condition": {
+                "StringEquals": {
+                    "iam:PassedToService": "datazone.amazonaws.com"
+                }
+            },
+            "Action": [
+                "iam:PassRole"
+            ],
+            "Resource": [
+                "arn:aws:iam::*:role/AmazonDataZone*",
+                "arn:aws:iam::*:role/service-role/AmazonDataZone*"
+            ],
+            "Effect": "Allow"
+        },
+        {
+            "Action": [
+                "datazone:AcceptSubscriptionRequest",
+                "datazone:CreateSubscriptionRequest",
+                "datazone:CreateDataSource",
+                "datazone:DeleteSubscriptionRequest",
+                "datazone:Search",
+                "datazone:SearchListings",
+                "datazone:ListSubscriptionRequests",
+                "datazone:ListAssetRevisions",
+                "datazone:GetListing",
+                "datazone:GetAsset",
+                "datazone:GetEnvironment",
+                "datazone:GetProject",
+                "datazone:RejectSubscriptionRequest",
+                "datazone:CancelSubscription",
+                "datazone:ListSubscriptions",
+                "datazone:ListDataSources",
+                "datazone:ListDataSourceRuns",
+                "datazone:ListDataSourceRunActivities",
+                "datazone:GetDataSource",
+                "datazone:GetDataSourceRun",
+                "datazone:StartDataSourceRun",
+                "datazone:UpdateDataSource",
+                "datazone:DeleteDataSource",
+                "datazone:ListEnvironments",
+                "datazone:CreateListingChangeSet",
+                "datazone:RevokeSubscription",
+                "datazone:GetGlossaryTerm",
+                "datazone:CreateAsset",
+                "datazone:CreateAssetRevision",
+                "datazone:GetUserProfile",
+                "datazone:GetUserProfile",
+                "datazone:ListSubscriptionGrants",
+                "datazone:CreateSubscriptionGrant",
+                "datazone:GetSubscriptionGrant",
+                "datazone:DeleteAsset",
+                "datazone:DeleteDataSource",
+                "datazone:GetFormType",
+                "datazone:ListGroupsForUser",
+                "datazone:ListProjectMemberships",
+                "datazone:DeleteListing"
+            ],
+            "Resource": "*",
+            "Effect": "Allow"
+        }
+    ]
+}