Skip to content

Commit

Permalink
NeptuneML integration (#48)
Browse files Browse the repository at this point in the history
- 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.
  • Loading branch information
austinkline authored Dec 29, 2020
1 parent 15ae7c9 commit 6b82d45
Show file tree
Hide file tree
Showing 22 changed files with 629 additions and 31 deletions.
4 changes: 1 addition & 3 deletions .github/workflows/integration.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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: |
Expand Down Expand Up @@ -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: |
Expand Down
2 changes: 1 addition & 1 deletion MANIFEST.in
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 30 additions & 0 deletions THIRD_PARTY_LICENSES.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
4 changes: 3 additions & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
8 changes: 8 additions & 0 deletions setupbase.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
2 changes: 1 addition & 1 deletion src/graph_notebook/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@
SPDX-License-Identifier: Apache-2.0
"""

__version__ = '2.0.2'
__version__ = '2.0.3'
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down
25 changes: 21 additions & 4 deletions src/graph_notebook/authentication/iam_headers.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,10 @@
import datetime
import hashlib
import hmac
import json
import logging
import urllib


logging.basicConfig()
logger = logging.getLogger("graph_magic")

Expand Down Expand Up @@ -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)

Expand All @@ -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 = []

Expand All @@ -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
Expand Down
18 changes: 18 additions & 0 deletions src/graph_notebook/magics/graph_magic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Loading

0 comments on commit 6b82d45

Please sign in to comment.