Update NeptuneML pre-trained model resources for CN regions (aws#226)

Co-authored-by: Michael Chin <[email protected]>
sam-orca · Nov 30, 2021 · 88c9e49 · 88c9e49
1 parent b91ed46
commit 88c9e49
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 14 deletions.
diff --git a/src/graph_notebook/notebooks/04-Machine-Learning/neptune-ml-pretrained-model-config.json b/src/graph_notebook/notebooks/04-Machine-Learning/neptune-ml-pretrained-model-config.json
@@ -1,12 +1,19 @@
 {
   "models": {
-      "node_classification": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/node-classification/model.tar.gz",
-        "node_regression": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/node-regression/model.tar.gz",
-        "link_prediction": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/link-prediction/model.tar.gz",
-       "edge_regression": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/edge-regression/model.tar.gz",
-       "edge_classification": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/edge-classification/model.tar.gz"
-    },
-    "container_images": {
+    "node_classification": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/node-classification/model.tar.gz",
+    "node_regression": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/node-regression/model.tar.gz",
+    "link_prediction": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/link-prediction/model.tar.gz",
+    "edge_regression": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/edge-regression/model.tar.gz",
+    "edge_classification": "s3://aws-neptune-customer-samples/v2/neptune-ml/pretrained-models/v2/edge-classification/model.tar.gz"
+  },
+  "models_cn": {
+    "node_classification": "s3://aws-neptune-customer-samples-cn-northwest-1/v2/neptune-ml/pretrained-models/v2/node-classification/model.tar.gz",
+    "node_regression": "s3://aws-neptune-customer-samples-cn-northwest-1/v2/neptune-ml/pretrained-models/v2/node-regression/model.tar.gz",
+    "link_prediction": "s3://aws-neptune-customer-samples-cn-northwest-1/v2/neptune-ml/pretrained-models/v2/link-prediction/model.tar.gz",
+    "edge_regression": "s3://aws-neptune-customer-samples-cn-northwest-1/v2/neptune-ml/pretrained-models/v2/edge-regression/model.tar.gz",
+    "edge_classification": "s3://aws-neptune-customer-samples-cn-northwest-1/v2/neptune-ml/pretrained-models/v2/edge-classification/model.tar.gz"
+  },
+  "container_images": {
     "us-west-1":"891482049861.dkr.ecr.us-west-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
     "us-west-2":"891482049861.dkr.ecr.us-west-2.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
     "us-east-1":"891482049861.dkr.ecr.us-east-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
@@ -25,8 +32,8 @@
     "eu-west-3":"891482049861.dkr.ecr.eu-west-3.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
     "me-south-1":"931515848886.dkr.ecr.me-south-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
     "sa-east-1":"891482049861.dkr.ecr.sa-east-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
-    "cn-north-1":"639043989634.dkr.ecr.sa-east-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
-    "cn-northwest-1":"639043989634.dkr.ecr.sa-east-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
+    "cn-north-1":"639043989634.dkr.ecr.cn-north-1.amazonaws.com.cn/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
+    "cn-northwest-1":"639043989634.dkr.ecr.cn-northwest-1.amazonaws.com.cn/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3",
     "us-gov-west-1":"106643864031.dkr.ecr.sa-east-1.amazonaws.com/graphlytics-pytorch-inference:v1-1.6.0-cpu-py3"
-    }
+  }
 }
diff --git a/src/graph_notebook/notebooks/04-Machine-Learning/neptune_ml_utils.py b/src/graph_notebook/notebooks/04-Machine-Learning/neptune_ml_utils.py
@@ -183,7 +183,7 @@ def prepare_movielens_data(s3_bucket_uri: str):
 
 
 def setup_pretrained_endpoints(s3_bucket_uri: str, setup_node_classification: bool,
-                               setup_node_regression: bool, setup_link_prediction: bool, \
+                               setup_node_regression: bool, setup_link_prediction: bool,
                                setup_edge_classification: bool, setup_edge_regression: bool):
     delete_pretrained_data(setup_node_classification,
                            setup_node_regression, setup_link_prediction,
@@ -488,9 +488,12 @@ class PretrainedModels:
     def __init__(self):
         with open('./neptune-ml-pretrained-model-config.json') as f:
             config = json.load(f)
-            self.PRETRAINED_MODEL = config['models']
-            self.PYTORCH_CPU_CONTAINER_IMAGE = config['container_images'][boto3.session.Session(
-            ).region_name]
+            region_name = boto3.session.Session().region_name
+            if region_name in ['cn-north-1', 'cn-northwest-1']:
+                self.PRETRAINED_MODEL = config['models_cn']
+            else:
+                self.PRETRAINED_MODEL = config['models']
+            self.PYTORCH_CPU_CONTAINER_IMAGE = config['container_images'][region_name]
 
     def __run_create_model(self, sm_client,
                            name,