add tunneling support

text_to_image.py

@@ -1,6 +1,8 @@
 import base64
 import logging
 import os
+import hashlib
+
 import requests
 import time
 from io import BytesIO
@@ -11,18 +13,23 @@ from fair import FairClient
 logger = logging.getLogger()
 
 SERVER_ADDRESS = "https://faircompute.com:8000"
-ENDPOINT_ADDRESS = "http://dikobraz.mooo.com:5000"
-TARGET_NODE = "119eccba-2388-43c1-bdb9-02133049604c"
+INFERENCE_NODE = "119eccba-2388-43c1-bdb9-02133049604c"
+TUNNEL_NODE = "c312e6c4788b00c73c287ab0445d3655"
 # SERVER_ADDRESS = "http://localhost:8000"
-# ENDPOINT_ADDRESS = "http://localhost:5000"
-# TARGET_NODE = "ef09913249aa40ecba7d0097f7622855"
+# INFERENCE_NODE = "ef09913249aa40ecba7d0097f7622855"
+# TUNNEL_NODE = "c312e6c4788b00c73c287ab0445d3655"
+
+INFERENCE_DOCKER_IMAGE = "faircompute/diffusers-api-dreamshaper-8"
+TUNNEL_DOCKER_IMAGE = "rapiz1/rathole"
 
-DOCKER_IMAGE = "faircompute/diffusers-api-dreamshaper-8"
+endpoint_client = None
+fair_client = None
 
 
 class EndpointClient:
-    def __init__(self, timeout):
-        response = requests.get(os.path.join(ENDPOINT_ADDRESS, 'healthcheck'), timeout=timeout).json()
+    def __init__(self, server_address, timeout):
+        self.endpoint_address = f'http://{server_address}:5000'
+        response = requests.get(os.path.join(self.endpoint_address, 'healthcheck'), timeout=timeout).json()
         if response['state'] != 'healthy':
             raise Exception("Server is not healthy")
 
@@ -43,7 +50,7 @@ class EndpointClient:
             },
         }
 
-        response = requests.post(ENDPOINT_ADDRESS, json=inputs).json()
+        response = requests.post(self.endpoint_address, json=inputs).json()
         image_data = BytesIO(base64.b64decode(response["image_base64"]))
         image = Image.open(image_data)
 
@@ -54,10 +61,17 @@ class ServerNotReadyException(Exception):
     pass
 
 
-def wait_for_server(retries, timeout=1.0, delay=2.0):
+def create_fair_client():
+    return FairClient(server_address=SERVER_ADDRESS,
+                      user_email=os.getenv('FAIRCOMPUTE_EMAIL', "debug-usr"),
+                      user_password=os.environ.get('FAIRCOMPUTE_PASSWORD', "debug-pwd"))
+
+
+def create_endpoint_client(fc, retries, timeout=1.0, delay=2.0):
+    server_address = next(info['host_address'] for info in fc.get_nodes() if info['node_id'] == TUNNEL_NODE)
     for i in range(retries):
         try:
-            return EndpointClient(timeout=timeout)
+            return EndpointClient(server_address, timeout=timeout)
         except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError, requests.exceptions.Timeout) as e:
             logging.exception(e)
             time.sleep(delay)
@@ -65,27 +79,80 @@ def wait_for_server(retries, timeout=1.0, delay=2.0):
     raise ServerNotReadyException("Failed to start the server")
 
 
-def start_server():
-    # default credentials will work only for local server built in debug mode
-    client = FairClient(server_address=SERVER_ADDRESS,
-                        user_email=os.getenv('FAIRCOMPUTE_EMAIL', "debug-usr"),
-                        user_password=os.environ.get('FAIRCOMPUTE_PASSWORD', "debug-pwd"))
-
-    client.run(node=TARGET_NODE,
-               image=DOCKER_IMAGE,
-               runtime="nvidia",
-               ports=[(5000, 8000)],
-               detach=True)
+def start_tunnel(fc: FairClient):
+    # generate fixed random authentication token based off some secret
+    token = hashlib.sha256(os.environ.get('FAIRCOMPUTE_PASSWORD', "debug-pwd").encode()).hexdigest()
+
+    # start tunnel node
+    server_config = f"""
+[server]
+bind_addr = "0.0.0.0:2333"  # port that rathole listens for clients
+
+[server.services.inference_server]
+token = "{token}"           # token that is used to authenticate the client for the service
+bind_addr = "0.0.0.0:5000"  # port that exposes service to the Internet
+"""
+    with open('server.toml', 'w') as file:
+        file.write(server_config)
+    fc.run(node=TUNNEL_NODE,
+           image=TUNNEL_DOCKER_IMAGE,
+           command=["--server", "/app/config.toml"],
+           volumes=[("./server.toml", "/app/config.toml")],
+           network="host",
+           detach=True)
+
+    server_address = next(info['host_address'] for info in fc.get_nodes() if info['node_id'] == TUNNEL_NODE)
+    client_config = f"""
+[client]
+remote_addr = "{server_address}:2333"       # address of the rathole server
+
+[client.services.inference_server]
+token = "{token}"                           # token that is used to authenticate the client for the service
+local_addr = "127.0.0.1:5001"               # address of the service that needs to be forwarded
+"""
+    with open('client.toml', 'w') as file:
+        file.write(client_config)
+    fc.run(node=INFERENCE_NODE,
+           image=TUNNEL_DOCKER_IMAGE,
+           command=["--client", "/app/config.toml"],
+           volumes=[("./client.toml", "/app/config.toml")],
+           network="host",
+           detach=True)
+
+
+def start_inference_server(fc: FairClient):
+    fc.run(node=INFERENCE_NODE,
+           image=INFERENCE_DOCKER_IMAGE,
+           runtime="nvidia",
+           ports=[(5001, 8000)],
+           detach=True)
+
+
+def start_services(fc):
+    start_tunnel(fc)
+    start_inference_server(fc)
 
 
 def text_to_image(text):
-    try:
-        client = wait_for_server(retries=1)
-    except ServerNotReadyException:
-        start_server()
-        client = wait_for_server(retries=10)
+    global endpoint_client
+    global fair_client
+    if fair_client is None:
+        fair_client = create_fair_client()
 
-    return client.infer(text)
+    try:
+        if endpoint_client is None:  # try connecting to the server
+            endpoint_client = create_endpoint_client(fair_client, 1)
+        else:  # try inference
+            return endpoint_client.infer(text)
+    except (requests.exceptions.ConnectionError, requests.exceptions.HTTPError, requests.exceptions.Timeout):
+        endpoint_client = None
+
+    # start all services
+    if endpoint_client is None:
+        start_services(fair_client)
+        endpoint_client = create_endpoint_client(fair_client, retries=10)
+
+    return endpoint_client.infer(text)
 
 
 if __name__ == "__main__":