Merge cb64b62 into 0c1571d

awslabs · Jan 26, 2023 · c6b70c6 · c6b70c6
2 parents 0c1571d + cb64b62
commit c6b70c6
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 781 deletions.
diff --git a/main/solution/post-deployment/config/environment-files/bootstrap.sh b/main/solution/post-deployment/config/environment-files/bootstrap.sh
@@ -45,20 +45,17 @@ update_jupyter_config() {
     import subprocess
     import os
     from notebook.services.sessions.sessionmanager import SessionManager as BaseSessionManager
-    from notebook.services.kernels.kernelmanager import MappingKernelManager
+    from notebook.services.kernels.kernelmanager import AsyncMappingKernelManager
 
-    class CustomSessionManager(BaseSessionManager):
-        def __init__(self, *args, **kwargs):
-            """Override default __init__() method"""
+    class SessionManager(BaseSessionManager):
+        def list_sessions(self, *args, **kwargs):
+            """Override default list_sessions() method"""
             self.mount_studies()
-            result = super(CustomSessionManager, self).__init__(*args, **kwargs)
+            result = super(SessionManager, self).list_sessions(*args, **kwargs)
             return result
-
         def mount_studies(self):
-            """Execute mount_s3.sh if mounting hasn't been done yet"""
-            path = "~/studies"
-            isExist = os.path.exists(path)
-            if not isExist:
+            """Execute mount_s3.sh if it hasn't already been run"""
+            if not hasattr(self, 'studies_mounted'):
                 mounting_result = subprocess.run(
                     "mount_s3.sh",
                     stdout=subprocess.PIPE, stderr=subprocess.STDOUT
@@ -68,9 +65,8 @@ update_jupyter_config() {
                     for line in mounting_result.stdout.decode("utf-8").split("\n"):
                         if line: # Skip empty lines
                             self.log.info(line)
-                            
-    c.NotebookApp.session_manager_class = CustomSessionManager
-    c.NotebookApp.kernel_manager_class = AsyncMappingKernelManager
+                self.studies_mounted = True
+    c.NotebookApp.session_manager_class = SessionManager
 EOF
 }
 
@@ -174,10 +170,7 @@ case "$(env_type)" in
             sudo yum --disablerepo=* localinstall -y *.rpm
             echo "Finish installing fuse"
         fi
-        cat "${FILES_DIR}/offline-packages/sagemaker/jupyter_notebook_config.py" > "/home/ec2-user/.jupyter/jupyter_notebook_config.py"
         update_jupyter_config "/home/ec2-user/.jupyter/jupyter_notebook_config.py"
-        # TODO: Automate running mount_s3 script. update_jupyter_config does not seem to work on JupyterLabv3
-        echo "Setting up a one-time mount script."
         if [ $OS_VERSION = '2' ]
         then
             systemctl restart jupyter-server