NeptuneML integration (#48)

- Integration with NeptuneML feature set in AWS Neptune
- Add helper library to perform Sigv4 signing for `%neptune_ml export ...`, we will move our other signing at a later date.
- Swap how credentials are obtained for `ROLE` iam credentials provider such that it uses a botocore session now instead of calling the ec2 metadata service. This should make the module more usable outside of Sagemaker.

New Line magics:
- `%neptune_ml export status`
- `%neptune_ml dataprocessing start`
- `%neptune_ml dataprocessing status`
- `%neptune_ml training start`
- `%neptune_ml training status`
- `%neptune_ml endpoint create`
- `%neptune_ml endpoint status`

New Cell magics:
- `%%neptune_ml export start`
- `%%neptune_ml dataprocessing start`
- `%%neptune_ml training start`
- `%%neptune_ml endpoint create`

NOTE: If a cell magic is used, its line inputs for specifying parts of the command will be ignore such as `--job-id` as a line-param.

Inject variable as cell input:
Currently this will only work for our new cell magic commands details above. You can now specify a variable to use as the cell input received by our `neptune_ml` magics using the syntax ${var_name}. For example...

```
# in one notebook cell:
foo = {'foo', 'bar'}

# in another notebook cell:
%%neptune_ml export start

${foo}
```

NOTE: The above will only work if it is the sole content of the cell body. You cannot inline multiple variables at this time.

Author: austinkline

.github/workflows/integration.yml

@@ -38,8 +38,7 @@ jobs:
       - uses: actions/checkout@v2
       - name: Install dependencies
         run: |
-          python -m pip install --upgrade pip
-          pip install flake8 pytest
+          pip install pytest
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
       - name: Install
         run: |
@@ -75,7 +74,6 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install flake8 pytest
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
       - name: Install
         run: |

MANIFEST.in

@@ -4,7 +4,7 @@ include webpack.config.js
 
 # Javascript files
 recursive-include src/graph_notebook/widgets *
-prune **/node_modules
+prune src/graph_notebook/widgets/node_modules
 prune coverage
 
 # Patterns to exclude from any directory

THIRD_PARTY_LICENSES.txt

@@ -2591,6 +2591,36 @@ SOFTWARE.
 
 ------
 
+** requests-aws4auth 1.0.1; version 1.0.1 --
+https://pypi.org/project/requests-aws4auth/
+The MIT License (MIT)
+
+Copyright (c) 2015 Sam Washington
+
+The MIT License (MIT)
+
+Copyright (c) 2015 Sam Washington
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+------
+
 ** @types/webpack-env; version 1.15.2 --
 https://github.com/DefinitelyTyped/DefinitelyTyped
 Copyright (c) Microsoft Corporation. All rights reserved.

setup.py

@@ -76,7 +76,9 @@ def get_version():
         'notebook',
         'jupyter-contrib-nbextensions',
         'widgetsnbextension',
-        'jupyter>=1.0.0'
+        'jupyter>=1.0.0',
+        'requests-aws4auth==1.0.1',
+        'botocore>=1.19.37'
     ],
     package_data={
         'graph_notebook': ['graph_notebook/widgets/nbextensions/static/*.js',

setupbase.py

@@ -382,6 +382,14 @@ def run(self):
             if should_build:
                 run(npm_cmd + ['run', build_cmd], cwd=node_package)
 
+            # ensure that __init__.py files are added to generated directories, otherwise it will not be packaged with
+            # package distribution to pypi
+            dirs_from_node_path = ['nbextension', pjoin('nbextension', 'static'), 'lib', 'labextension']
+            for init_path in dirs_from_node_path:
+                full_path = pjoin(node_package, init_path, '__init__.py')
+                with open(full_path, 'w+'):
+                    pass
+
     return NPM
 
 

src/graph_notebook/__init__.py

@@ -3,4 +3,4 @@
 SPDX-License-Identifier: Apache-2.0
 """
 
-__version__ = '2.0.2'
+__version__ = '2.0.3'

src/graph_notebook/authentication/iam_credentials_provider/ec2_metadata_credentials_provider.py

@@ -3,13 +3,14 @@
 SPDX-License-Identifier: Apache-2.0
 """
 
+import botocore.session
 import requests
 
+
 from graph_notebook.authentication.iam_credentials_provider.credentials_provider import CredentialsProviderBase, \
     Credentials
 
 region_url = 'http://169.254.169.254/latest/meta-data/placement/availability-zone'
-iam_url = 'http://169.254.169.254/latest/meta-data/iam/security-credentials/neptune-db'
 
 
 class MetadataCredentialsProvider(CredentialsProviderBase):
@@ -20,10 +21,6 @@ def __init__(self):
         self.region = region
 
     def get_iam_credentials(self) -> Credentials:
-        res = requests.get(iam_url)
-        if res.status_code != 200:
-            raise Exception(f'unable to get iam credentials {res.content}')
-
-        js = res.json()
-        creds = Credentials(key=js['AccessKeyId'], secret=js['SecretAccessKey'], token=js['Token'], region=self.region)
-        return creds
+        session = botocore.session.get_session()
+        creds = session.get_credentials()
+        return Credentials(key=creds.access_key, secret=creds.secret_key, token=creds.token, region=self.region)

src/graph_notebook/authentication/iam_credentials_provider/env_credentials_provider.py

@@ -28,7 +28,7 @@ def load_iam_credentials(self):
         self.loaded = True
         return
 
-    def get_iam_credentials(self) -> Credentials:
+    def get_iam_credentials(self, service=None) -> Credentials:
         if not self.loaded:
             self.load_iam_credentials()
 

src/graph_notebook/authentication/iam_headers.py

@@ -6,10 +6,10 @@
 import datetime
 import hashlib
 import hmac
+import json
 import logging
 import urllib
 
-
 logging.basicConfig()
 logger = logging.getLogger("graph_magic")
 
@@ -69,6 +69,19 @@ def get_canonical_uri_and_payload(query_type, query):
     elif query_type == "system":
         canonical_uri = "/system/"
         payload = query
+
+    elif query_type.startswith("ml"):
+        canonical_uri = f'/{query_type}'
+        payload = query
+
+    elif query_type.startswith("ml/dataprocessing"):
+        canonical_uri = f'/{query_type}'
+        payload = query
+
+    elif query_type.startswith("ml/endpoints"):
+        canonical_uri = f'/{query_type}'
+        payload = query
+
     else:
         raise ValueError('query_type %s is not valid' % query_type)
 
@@ -85,7 +98,8 @@ def normalize_query_string(query):
     return normalized
 
 
-def make_signed_request(method, query_type, query, host, port, signing_access_key, signing_secret, signing_region, use_ssl=False, signing_token='', additional_headers=None):
+def make_signed_request(method, query_type, query, host, port, signing_access_key, signing_secret, signing_region,
+                        use_ssl=False, signing_token='', additional_headers=None):
     if additional_headers is None:
         additional_headers = []
 
@@ -103,8 +117,11 @@ def make_signed_request(method, query_type, query, host, port, signing_access_ke
     # get canonical_uri and payload
     canonical_uri, payload = get_canonical_uri_and_payload(query_type, query)
 
-    request_parameters = urllib.parse.urlencode(payload, quote_via=urllib.parse.quote)
-    request_parameters = request_parameters.replace('%27', '%22')
+    if 'content-type' in additional_headers and additional_headers['content-type'] == 'application/json':
+        request_parameters = payload if type(payload) is str else json.dumps(payload)
+    else:
+        request_parameters = urllib.parse.urlencode(payload, quote_via=urllib.parse.quote)
+        request_parameters = request_parameters.replace('%27', '%22')
     t = datetime.datetime.utcnow()
     amz_date = t.strftime('%Y%m%dT%H%M%SZ')
     date_stamp = t.strftime('%Y%m%d')  # Date w/o time, used in credential scope

src/graph_notebook/magics/graph_magic.py

@@ -23,6 +23,7 @@
 import graph_notebook
 from graph_notebook.configuration.generate_config import generate_default_config, DEFAULT_CONFIG_LOCATION
 from graph_notebook.decorators.decorators import display_exceptions
+from graph_notebook.magics.ml import neptune_ml_magic_handler, generate_neptune_ml_parser
 from graph_notebook.network import SPARQLNetwork
 from graph_notebook.network.gremlin.GremlinNetwork import parse_pattern_list_str, GremlinNetwork
 from graph_notebook.sparql.table import get_rows_and_columns
@@ -994,3 +995,20 @@ def graph_notebook_vis_options(self, line='', cell=''):
         else:
             options_dict = json.loads(cell)
             self.graph_notebook_vis_options = vis_options_merge(self.graph_notebook_vis_options, options_dict)
+
+    @line_cell_magic
+    @display_exceptions
+    @needs_local_scope
+    def neptune_ml(self, line, cell='', local_ns: dict = None):
+        parser = generate_neptune_ml_parser()
+        args = parser.parse_args(line.split())
+        logger.info(f'received call to neptune_ml with details: {args.__dict__}, cell={cell}, local_ns={local_ns}')
+        request_generator = create_request_generator(self.graph_notebook_config.auth_mode,
+                                                     self.graph_notebook_config.iam_credentials_provider_type)
+        main_output = widgets.Output()
+        display(main_output)
+        res = neptune_ml_magic_handler(args, request_generator, self.graph_notebook_config, main_output, cell, local_ns)
+        message = json.dumps(res, indent=2) if type(res) is dict else res
+        store_to_ns(args.store_to, res, local_ns)
+        with main_output:
+            print(message)