From 4e43fbbb4575c33eb9b3f7bbf74d2a016dbd495e Mon Sep 17 00:00:00 2001
From: Dana Benson <danabens@amazon.com>
Date: Fri, 22 Nov 2024 09:55:05 -0800
Subject: [PATCH] add additional setup reminders, delete env link during
 offboard

---
    ml_ops/sm-datazone_import/README.md        | 13 +++++--
 .../offboard-sagemaker-domain.py              | 35 +++++++++++++++++++
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/   ml_ops/sm-datazone_import/README.md b/   ml_ops/sm-datazone_import/README.md
index 5d4b21771b..ad364b1be2 100644
--- a/   ml_ops/sm-datazone_import/README.md	
+++ b/   ml_ops/sm-datazone_import/README.md	
@@ -4,13 +4,22 @@ This example contains python scripts to import an existing SageMaker Domain into
 
 ## Setup
 
-1. Add the Bring-Your-Own-Domain (BYOD) service model
+1. Ensure dependencies are up to date.
+
+Update the CLI. See instructions https://docs.aws.amazon.com/cli/latest/userguide/getting-started-install.html
+`aws --version`
+
+Update boto.
+`pip install --upgrade -r requirements.txt`
+`pip show boto3`
+
+2. 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
 ```
 
-2. Create a federation role
+3. Create a federation role
 
 This role will be used by DataZone to launch the SageMaker Domain. See [BringYourOwnDomainResources.yml](.resources/BringYourOwnDomainResources.yml) for an example.
 
diff --git a/   ml_ops/sm-datazone_import/offboard-sagemaker-domain.py b/   ml_ops/sm-datazone_import/offboard-sagemaker-domain.py
index 525f377406..907cddb604 100644
--- a/   ml_ops/sm-datazone_import/offboard-sagemaker-domain.py	
+++ b/   ml_ops/sm-datazone_import/offboard-sagemaker-domain.py	
@@ -22,6 +22,8 @@ def _offboard_sm_domain(self):
         # Firstly, domain remove tags.
         sm_domain_tags = [
             "AmazonDataZoneDomain",
+            "AmazonDataZoneProject",
+            "AmazonDataZoneEnvironment",
             "AmazonDataZoneDomainAccount",
             "AmazonDataZoneStage",
         ]
@@ -138,6 +140,25 @@ def _select_dz_project(self):
             )
         )
 
+    def _select_dz_environment(self):
+        print("--------------------------------------------------------------------")
+        print("List of DataZone Environments.")
+        print("--------------------------------------------------------------------")
+        dz_env_map = {}
+        for env in self.dz_client.list_environments(
+            domainIdentifier=self.dz_domain_id, projectIdentifier=self.dz_project_id
+        )["items"]:
+            print(f'Name: {env["name"]}')
+            dz_env_map[env["name"]] = env["id"]
+        self.env_name = input("Please provide the name of DataZone environment: ")
+        self.env_id = dz_env_map[self.env_name]
+        print(
+            "Chosen DataZone Environment [{}] with Environment Id [{}]".format(
+                self.env_name, self.env_id
+            )
+        )
+        return self.env_id
+
     def _delete_linked_items(self):
         user_profile_arn = "arn:aws:sagemaker:{}:{}:user-profile/{}/{}".format(
             self.region, self.account_id, self.sm_domain_id, self.sm_user_name
@@ -165,6 +186,18 @@ def _delete_linked_items(self):
             itemIdentifiers=[sm_domain_arn],
         )
 
+    def _delete_env_action_link(self):
+        for action_link in self.dz_client.list_environment_actions(
+            domainIdentifier=self.dz_domain_id, environmentIdentifier=self.env_id
+        )["items"]:
+            name = action_link["name"]
+            if "SageMaker" in name:
+                action_id = action_link["id"]
+                self.dz_client.delete_environment_action(
+                    domainIdentifier=self.dz_domain_id,
+                    environmentIdentifier=self.env_id,
+                    identifier=action_id)
+
     def _print_results(self):
         # Verify no linked entities
         print("--------------------------------------------------------------------")
@@ -182,7 +215,9 @@ def offboard(self):
         self._offboard_sm_users()
         self._select_dz_domain()
         self._select_dz_project()
+        self._select_dz_environment()
         self._delete_linked_items()
+        self._delete_env_action_link()
         self._print_results()