Private Link URL support on rai_core_flask (#2493)

* pl support ckpt * decoupled changes * python lint fixes
microsoft · Jan 19, 2024 · 95d5e05 · 95d5e05
1 parent 946115b
commit 95d5e05
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/rai_core_flask/rai_core_flask/environments/azure_nb_environment.py b/rai_core_flask/rai_core_flask/environments/azure_nb_environment.py
@@ -36,14 +36,20 @@ def __init__(self, service):
             else:
                 instance_name = self.nbvm["instance"]
                 domain_suffix = self.nbvm["domainsuffix"]
+                if service.is_private_link:
+                    url_format = f"{instance_name}.{domain_suffix}\
+                        :{service.port}"
+                else:
+                    url_format = f"{instance_name}-{service.port}\
+                        .{domain_suffix}"
                 self.base_url = \
-                    f"https://{instance_name}-{service.port}.{domain_suffix}"
+                    f"https://{url_format}"
                 self.successfully_detected = True
                 self.nbvm_origins = [
                     f"https://{instance_name}.{domain_suffix}",
-                    f"https://{instance_name}-{service.port}.{domain_suffix}",
+                    f"https://{url_format}",
                     f"wss://{instance_name}.{domain_suffix}",
-                    f"wss://{instance_name}-{service.port}.{domain_suffix}",
+                    f"wss://{url_format}",
                     "https://ml.azure.com",
                     "https://dev.ml.azure.com",
                     "https://int.ml.azure.com"

diff --git a/rai_core_flask/rai_core_flask/flask_helper.py b/rai_core_flask/rai_core_flask/flask_helper.py
@@ -24,13 +24,15 @@
 class FlaskHelper(object):
     """FlaskHelper is a class for common Flask utilities used in dashboards."""
 
-    def __init__(self, ip=None, port=None, with_credentials=False):
+    def __init__(self, ip=None, port=None, with_credentials=False,
+                 is_private_link=False):
         # The name passed to Flask needs to be unique per instance.
         self.app = Flask(uuid.uuid4().hex)
 
         self.port = port
         self.ip = ip
         self.with_credentials = with_credentials
+        self.is_private_link = is_private_link
         # dictionary to store arbitrary state for use by consuming classes
         self.shared_state = {}
         if self.ip is None: