add sagemaker studio jupyter scheduler (#786)

Author: TRNWWZ

template/v2/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json

@@ -8,7 +8,6 @@
       "@amzn/sagemaker-jupyterlab-extensions:spacemenu": true,
       "@amzn/amazon_sagemaker_sql_editor": true,
       "@sagemaker-studio:EmrCluster": true,
-      "@jupyterlab/scheduler": true,
       "@jupyter/collaboration-extension": true
     }
 }

template/v2/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py

@@ -1,5 +1,11 @@
 import os
 
+from sagemaker_studio_jupyter_scheduler.scheduler import (
+    SagemakerEnvironmentManager,
+    SageMakerJobFilesManager,
+    SageMakerUnifiedStudioScheduler,
+)
+
 c.ServerApp.terminado_settings = {"shell_command": ["/bin/bash"]}
 region = os.environ.get("AWS_REGION")
 csp_rule = os.environ.get("JUPYTERSERVER_CSP_RULE")
@@ -27,4 +33,9 @@
     module_location = os.path.dirname(module.__file__)
     c.LanguageServerManager.extra_node_roots = [f"{module_location}/sql-language-server"]
 except:
+
     pass
+
+c.SchedulerApp.scheduler_class = SageMakerUnifiedStudioScheduler
+c.SchedulerApp.environment_manager_class = SagemakerEnvironmentManager
+c.SchedulerApp.job_files_manager_class = SageMakerJobFilesManager

template/v3/dirs/etc/sagemaker-ui/jupyter/lab/settings/page_config.json

@@ -8,7 +8,6 @@
       "@amzn/sagemaker-jupyterlab-extensions:spacemenu": true,
       "@amzn/amazon_sagemaker_sql_editor": true,
       "@sagemaker-studio:EmrCluster": true,
-      "@jupyterlab/scheduler": true,
       "@jupyter/collaboration-extension": true
     }
 }

template/v3/dirs/etc/sagemaker-ui/jupyter/server/jupyter_server_config.py

@@ -1,5 +1,11 @@
 import os
 
+from sagemaker_studio_jupyter_scheduler.scheduler import (
+    SagemakerEnvironmentManager,
+    SageMakerJobFilesManager,
+    SageMakerUnifiedStudioScheduler,
+)
+
 c.ServerApp.terminado_settings = {"shell_command": ["/bin/bash"]}
 region = os.environ.get("AWS_REGION")
 csp_rule = os.environ.get("JUPYTERSERVER_CSP_RULE")
@@ -28,3 +34,7 @@
     c.LanguageServerManager.extra_node_roots = [f"{module_location}/sql-language-server"]
 except:
     pass
+
+c.SchedulerApp.scheduler_class = SageMakerUnifiedStudioScheduler
+c.SchedulerApp.environment_manager_class = SagemakerEnvironmentManager
+c.SchedulerApp.job_files_manager_class = SageMakerJobFilesManager