Updated NeptuneML notebooks (new instances)

.github/workflows/unit.yml

@@ -18,7 +18,7 @@ jobs:
       - name: Set up Node.js 20.x
         uses: actions/setup-node@v3
         with:
-          node-version: '20.18.3'
+          node-version: '20.19.0'
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip

Dockerfile

@@ -11,7 +11,7 @@ ENV pipargs=""
 ENV WORKING_DIR="/root"
 ENV NOTEBOOK_DIR="${WORKING_DIR}/notebooks"
 ENV EXAMPLE_NOTEBOOK_DIR="${NOTEBOOK_DIR}/Example Notebooks"
-ENV NODE_VERSION=20.18.3
+ENV NODE_VERSION=20.19.0
 ENV PYTHON_VERSION=3.10
 ENV GRAPH_NOTEBOOK_AUTH_MODE="DEFAULT"
 ENV GRAPH_NOTEBOOK_HOST="neptune.cluster-XXXXXXXXXXXX.us-east-1.neptune.amazonaws.com"

additional-databases/sagemaker/neptune-notebook-cloudformation/neptune-workbench-stack.yaml

@@ -16,26 +16,16 @@ Parameters:
     - ml.t3.large
     - ml.t3.medium
     - ml.t3.xlarge
-    - ml.m4.xlarge
-    - ml.m4.2xlarge
-    - ml.m4.4xlarge
-    - ml.m4.10xlarge
-    - ml.m4.16xlarge
     - ml.m5.12xlarge
     - ml.m5.24xlarge
     - ml.m5.2xlarge
     - ml.m5.4xlarge
     - ml.m5.xlarge
-    - ml.p2.16xlarge
-    - ml.p2.8xlarge
-    - ml.p2.xlarge
-    - ml.p3.16xlarge
-    - ml.p3.2xlarge
-    - ml.p3.8xlarge
-    - ml.c4.2xlarge
-    - ml.c4.4xlarge
-    - ml.c4.8xlarge
-    - ml.c4.xlarge
+    - ml.g5.12xlarge
+    - ml.g5.8xlarge
+    - ml.g5.xlarge
+    - ml.g5.2xlarge
+    - ml.g5.4xlarge
     - ml.c5.18xlarge
     - ml.c5.2xlarge
     - ml.c5.4xlarge

src/graph_notebook/notebooks/03-Neptune-ML/01-Gremlin/02-Introduction-to-Node-Classification-Gremlin.ipynb

@@ -456,7 +456,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name}\n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/01-Gremlin/03-Introduction-to-Node-Regression-Gremlin.ipynb

@@ -453,7 +453,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/01-Gremlin/04-Introduction-to-Link-Prediction-Gremlin.ipynb

@@ -427,7 +427,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/01-Gremlin/05-Introduction-to-Edge-Classification-Gremlin.ipynb

@@ -455,7 +455,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/01-Gremlin/06-Introduction-to-Edge-Regression-Gremlin.ipynb

@@ -454,7 +454,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/02-SPARQL/Neptune-ML-01-Introduction-to-Object-Classification-SPARQL.ipynb

@@ -317,7 +317,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name}\n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training \"\"\""
    ]
   },

src/graph_notebook/notebooks/03-Neptune-ML/02-SPARQL/Neptune-ML-02-Introduction-to-Object-Regression-SPARQL.ipynb

@@ -355,7 +355,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/02-SPARQL/Neptune-ML-03-Introduction-to-Link-Prediction-SPARQL.ipynb

@@ -399,7 +399,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name} \n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training\n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/02-SPARQL/neptune_ml_sparql_utils.py

@@ -5,6 +5,42 @@
 import pandas as pd
 import numpy as np
 import pickle
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """Restrict unpickling to safe types to prevent arbitrary code execution."""
+
+    ALLOWED_CLASSES = {
+        ('builtins', 'set'),
+        ('builtins', 'frozenset'),
+        ('collections', 'OrderedDict'),
+        ('collections', 'defaultdict'),
+        ('numpy', 'ndarray'),
+        ('numpy', 'dtype'),
+        ('numpy.core.multiarray', '_reconstruct'),
+        ('numpy', 'core'),
+    }
+
+    def find_class(self, module, name):
+        if (module, name) in self.ALLOWED_CLASSES:
+            return super().find_class(module, name)
+        if module == 'builtins' and name in ('dict', 'list', 'tuple', 'str', 'int', 'float', 'bool', 'bytes'):
+            return super().find_class(module, name)
+        if module.startswith('numpy') and name in (
+            'ndarray', 'dtype', '_reconstruct', 'scalar', 'array',
+        ):
+            return super().find_class(module, name)
+        if module.startswith('sklearn'):
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f"Deserialization of {module}.{name} is blocked for security reasons"
+        )
+
+
+def safe_pickle_load(filepath):
+    """Load a pickle file using the restricted unpickler."""
+    with open(filepath, "rb") as f:
+        return RestrictedUnpickler(f).load()
 import os
 import requests
 import json
@@ -209,15 +245,14 @@ def get_node_to_idx_mapping(training_job_name: str = None, dataprocessing_job_na
             return
         S3Downloader.download(os.path.join(job_s3_output, filename), model_artifacts_location)
 
-    with open(os.path.join(model_artifacts_location, filename), "rb") as f:
-        mapping = pickle.load(f)[mapping_key]
-        if vertex_label is not None:
-            if vertex_label in mapping:
-                mapping = mapping[vertex_label]
-            else:
-                print("Mapping for vertex label: {} not found.".format(vertex_label))
-                print("valid vertex labels which have vertices mapped to embeddings: {} ".format(list(mapping.keys())))
-                print("Returning mapping for all valid vertex labels")
+    mapping = safe_pickle_load(os.path.join(model_artifacts_location, filename))[mapping_key]
+    if vertex_label is not None:
+        if vertex_label in mapping:
+            mapping = mapping[vertex_label]
+        else:
+            print("Mapping for vertex label: {} not found.".format(vertex_label))
+            print("valid vertex labels which have vertices mapped to embeddings: {} ".format(list(mapping.keys())))
+            print("Returning mapping for all valid vertex labels")
 
     return mapping
 

src/graph_notebook/notebooks/03-Neptune-ML/03-Sample-Applications/03-Real-Time-Fraud-Detection-Using-Inductive-Inference.ipynb

@@ -484,7 +484,7 @@
     "training_params=f\"\"\"\n",
     "--job-id {training_job_name} \n",
     "--data-processing-id {training_job_name}\n",
-    "--instance-type ml.p3.2xlarge\n",
+    "--instance-type ml.g5.2xlarge\n",
     "--s3-output-uri {str(s3_bucket_uri)}/training \n",
     "--max-hpo-number 2\n",
     "--max-hpo-parallel 2 \"\"\""

src/graph_notebook/notebooks/03-Neptune-ML/03-Sample-Applications/04-Telco-Networks/neptune_ml_utils.py

@@ -2,6 +2,42 @@
 import pandas as pd
 import numpy as np
 import pickle
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """Restrict unpickling to safe types to prevent arbitrary code execution."""
+
+    ALLOWED_CLASSES = {
+        ('builtins', 'set'),
+        ('builtins', 'frozenset'),
+        ('collections', 'OrderedDict'),
+        ('collections', 'defaultdict'),
+        ('numpy', 'ndarray'),
+        ('numpy', 'dtype'),
+        ('numpy.core.multiarray', '_reconstruct'),
+        ('numpy', 'core'),
+    }
+
+    def find_class(self, module, name):
+        if (module, name) in self.ALLOWED_CLASSES:
+            return super().find_class(module, name)
+        if module == 'builtins' and name in ('dict', 'list', 'tuple', 'str', 'int', 'float', 'bool', 'bytes'):
+            return super().find_class(module, name)
+        if module.startswith('numpy') and name in (
+            'ndarray', 'dtype', '_reconstruct', 'scalar', 'array',
+        ):
+            return super().find_class(module, name)
+        if module.startswith('sklearn'):
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f"Deserialization of {module}.{name} is blocked for security reasons"
+        )
+
+
+def safe_pickle_load(filepath):
+    """Load a pickle file using the restricted unpickler."""
+    with open(filepath, "rb") as f:
+        return RestrictedUnpickler(f).load()
 import os
 import requests
 import json
@@ -241,15 +277,14 @@ def get_node_to_idx_mapping(training_job_name: str = None, dataprocessing_job_na
             return
         S3Downloader.download(os.path.join(job_s3_output, filename), model_artifacts_location)
 
-    with open(os.path.join(model_artifacts_location, filename), "rb") as f:
-        mapping = pickle.load(f)[mapping_key]
-        if vertex_label is not None:
-            if vertex_label in mapping:
-                mapping = mapping[vertex_label]
-            else:
-                print("Mapping for vertex label: {} not found.".format(vertex_label))
-                print("valid vertex labels which have vertices mapped to embeddings: {} ".format(list(mapping.keys())))
-                print("Returning mapping for all valid vertex labels")
+    mapping = safe_pickle_load(os.path.join(model_artifacts_location, filename))[mapping_key]
+    if vertex_label is not None:
+        if vertex_label in mapping:
+            mapping = mapping[vertex_label]
+        else:
+            print("Mapping for vertex label: {} not found.".format(vertex_label))
+            print("valid vertex labels which have vertices mapped to embeddings: {} ".format(list(mapping.keys())))
+            print("Returning mapping for all valid vertex labels")
 
     return mapping
 

src/graph_notebook/notebooks/03-Neptune-ML/neptune_ml_utils.py

@@ -2,6 +2,42 @@
 import pandas as pd
 import numpy as np
 import pickle
+
+
+class RestrictedUnpickler(pickle.Unpickler):
+    """Restrict unpickling to safe types to prevent arbitrary code execution."""
+
+    ALLOWED_CLASSES = {
+        ('builtins', 'set'),
+        ('builtins', 'frozenset'),
+        ('collections', 'OrderedDict'),
+        ('collections', 'defaultdict'),
+        ('numpy', 'ndarray'),
+        ('numpy', 'dtype'),
+        ('numpy.core.multiarray', '_reconstruct'),
+        ('numpy', 'core'),
+    }
+
+    def find_class(self, module, name):
+        if (module, name) in self.ALLOWED_CLASSES:
+            return super().find_class(module, name)
+        if module == 'builtins' and name in ('dict', 'list', 'tuple', 'str', 'int', 'float', 'bool', 'bytes'):
+            return super().find_class(module, name)
+        if module.startswith('numpy') and name in (
+            'ndarray', 'dtype', '_reconstruct', 'scalar', 'array',
+        ):
+            return super().find_class(module, name)
+        if module.startswith('sklearn'):
+            return super().find_class(module, name)
+        raise pickle.UnpicklingError(
+            f"Deserialization of {module}.{name} is blocked for security reasons"
+        )
+
+
+def safe_pickle_load(filepath):
+    """Load a pickle file using the restricted unpickler."""
+    with open(filepath, "rb") as f:
+        return RestrictedUnpickler(f).load()
 import os
 import requests
 import json
@@ -241,15 +277,14 @@ def get_node_to_idx_mapping(training_job_name: str = None, dataprocessing_job_na
             return
         S3Downloader.download(os.path.join(job_s3_output, filename), model_artifacts_location)
 
-    with open(os.path.join(model_artifacts_location, filename), "rb") as f:
-        mapping = pickle.load(f)[mapping_key]
-        if vertex_label is not None:
-            if vertex_label in mapping:
-                mapping = mapping[vertex_label]
-            else:
-                print("Mapping for vertex label: {} not found.".format(vertex_label))
-                print("valid vertex labels which have vertices mapped to embeddings: {} ".format(list(mapping.keys())))
-                print("Returning mapping for all valid vertex labels")
+    mapping = safe_pickle_load(os.path.join(model_artifacts_location, filename))[mapping_key]
+    if vertex_label is not None:
+        if vertex_label in mapping:
+            mapping = mapping[vertex_label]
+        else:
+            print("Mapping for vertex label: {} not found.".format(vertex_label))
+            print("valid vertex labels which have vertices mapped to embeddings: {} ".format(list(mapping.keys())))
+            print("Returning mapping for all valid vertex labels")
 
     return mapping