Initial extensions persistence impl

**Description**

Initial extensions persistence implementation

Author: aws-pangestu

template/v2/dirs/usr/local/bin/merge-extensions-util.py

@@ -0,0 +1,74 @@
+import json
+import os
+import shutil
+import sys
+from datetime import datetime
+from typing import Set
+
+
+def backup_file(file_path):
+    directory, filename = os.path.split(file_path)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    backup_filename = f"{filename}.malformed_{timestamp}.bak"
+    backup_path = os.path.join(directory, backup_filename)
+    shutil.copy2(file_path, backup_path)
+    print(f"A backup of the malformed file has been created: {backup_path}")
+
+
+def load_json_file(file_path):
+    try:
+        with open(file_path, "r") as f:
+            return json.load(f)
+    except json.JSONDecodeError:
+        print(f"Error: {file_path} contains malformed JSON.")
+        if file_path == sys.argv[1]:  # If it's file1
+            backup_file(file_path)
+            print(f"Overwriting {file_path} with content from {sys.argv[2]}")
+            shutil.copy2(sys.argv[2], file_path)
+            with open(file_path, "r") as f:
+                return json.load(f)
+        else:
+            raise
+
+
+def main():
+    if len(sys.argv) != 3:
+        print("Usage: python script.py <file1> <file2>")
+        sys.exit(1)
+
+    file1, file2 = sys.argv[1], sys.argv[2]
+
+    try:
+        data1 = load_json_file(file1)
+        data2 = load_json_file(file2)
+
+        # Construct sets of extensions based on ID
+        data1_exts = {ext["identifier"]["id"]: ext for ext in data1}
+        data2_exts = {ext["identifier"]["id"]: ext for ext in data2}
+        all_exts: Set[str] = set(data1_exts.keys()) | set(data2_exts.keys())
+
+        # Merge extensions
+        merged_exts = []
+        for ext in all_exts:
+            ext_data = data2_exts.get(ext, data1_exts.get(ext, {}))
+            if "relativeLocation" in ext_data:
+                del ext_data["relativeLocation"]
+            merged_exts.append(ext_data)
+
+        # Write the merged data back to file1
+        print("Merged extensions: ", json.dumps(merged_exts, indent=2))
+        with open(file1, "w") as f:
+            json.dump(merged_exts, f, indent=2)
+
+        print(f"Successfully merged extensions and wrote to {file1}")
+
+    except FileNotFoundError as e:
+        print(f"Error: File not found - {e.filename}")
+    except json.JSONDecodeError as e:
+        print(f"Error: JSON decoding failed - {e}")
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+
+
+if __name__ == "__main__":
+    main()

template/v2/dirs/usr/local/bin/start-code-editor

@@ -1,11 +1,11 @@
 #!/bin/bash
 set -e
 
-EFS_MOUNT_POINT="/opt/amazon/sagemaker"
+SERVICE_DIRECTORY="/opt/amazon/sagemaker"
 EBS_MOUNT_POINT="/home/sagemaker-user"
 
 persistent_settings_folder="${EBS_MOUNT_POINT}/sagemaker-code-editor-server-data"
-default_settings_folder="${EFS_MOUNT_POINT}/sagemaker-code-editor-server-data"
+default_settings_folder="${SERVICE_DIRECTORY}/sagemaker-code-editor-server-data"
 
 override_machine_settings() {
   # create a new settings file with preset defaults or merge the defaults into the existing settings file
@@ -39,6 +39,27 @@ copy_user_settings() {
   fi
 }
 
+merge_prepackaged_extensions() {
+  local extensions_metadata_relative_path="/extensions"
+  local default_extensions_metadata_file="${default_settings_folder}/${extensions_metadata_relative_path}/extensions.json"
+  local persistent_extensions_metadata_file="${persistent_settings_folder}/${extensions_metadata_relative_path}/extensions.json"
+  local obsolete_file="${persistent_settings_folder}/${extensions_metadata_relative_path}/.obsolete"
+  if [ ! -f "$persistent_extensions_metadata_file" ]; then
+      # copy settings file to EBS if it doesn't exist in EBS
+      mkdir -p "${persistent_settings_folder}/${extensions_metadata_relative_path}"
+      echo "[]" > $persistent_extensions_metadata_file
+      echo "Created persistent extensions metadata file with default settings at $persistent_extensions_metadata_file"
+    else
+      # if it does exist then merge settings
+      echo "File already exists: ${persistent_extensions_metadata_file}. Overwriting extensions with pre-packaged extensions."
+    fi
+
+  if [ -f "$obsolete_file" ]; then
+    rm $obsolete_file
+  fi
+  python3 /usr/local/bin/merge-extensions-util.py "$persistent_extensions_metadata_file" "$default_extensions_metadata_file"
+}
+
 eval "$(micromamba shell hook --shell=bash)"
 
 # Activate conda environment 'base', which is the default environment for sagemaker-distribution
@@ -49,17 +70,18 @@ if [ -n "$SAGEMAKER_APP_TYPE_LOWERCASE" ]; then
   # SAGEMAKER_APP_TYPE is set, indicating the server is running within a SageMaker app.
   override_machine_settings
   copy_user_settings
+  merge_prepackaged_extensions
   # Configure the base url to be `/<app-type-in-lower-case>/default`.
   sagemaker-code-editor --host 0.0.0.0 --port 8888 \
     --without-connection-token \
     --base-path "/$SAGEMAKER_APP_TYPE_LOWERCASE/default" \
     --server-data-dir $persistent_settings_folder \
-    --extensions-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data/extensions \
+    --extensions-dir $persistent_settings_folder/extensions \
     --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data
 else
   sagemaker-code-editor --host 0.0.0.0 --port 8888 \
     --without-connection-token \
     --server-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data \
     --extension-dir /opt/amazon/sagemaker/sagemaker-code-editor-server-data/extensions \
     --user-data-dir /opt/amazon/sagemaker/sagemaker-code-editor-user-data
-fi
+fi