exo text on start and stop

exo/helpers.py

@@ -6,6 +6,24 @@ DEBUG = int(os.getenv("DEBUG", default="0"))
 DEBUG_DISCOVERY = int(os.getenv("DEBUG_DISCOVERY", default="0"))
 VERSION = "0.0.1"
 
+exo_text = """
+  _____  _____
+ / _ \ \/ / _ \
+|  __/>  < (_) |
+ \___/_/\_\___/
+    """
+
+def print_exo():
+    print(exo_text)
+
+def print_yellow_exo():
+    yellow = "\033[93m"  # ANSI escape code for yellow
+    reset = "\033[0m"    # ANSI escape code to reset color
+    exo = f"""{yellow}
+{exo_text}
+{reset}"""
+    print(exo)
+
 def terminal_link(uri, label=None):
     if label is None: 
         label = uri

exo/orchestration/standard_node.py

@@ -156,10 +156,10 @@ class StandardNode(Node):
         if DEBUG >= 2: print("Connecting to new peers...")
         for peer in self.peers:
             is_connected = await peer.is_connected()
-            if DEBUG >= 2: print(f"Connected to {peer.id()}: {is_connected}")
+            if DEBUG >= 2 and is_connected: print(f"Already connected to {peer.id()}: {is_connected}")
             if not is_connected:
                 await peer.connect()
-                if DEBUG >= 2: print(f"Connected to peer {peer.id()}")
+                if DEBUG >= 0: print(f"Connected to peer {peer.device_capabilities()} ({peer.id()=})")
 
     async def periodic_topology_collection(self, interval: int):
         while True:

main.py

@@ -10,6 +10,7 @@ from exo.networking.grpc.grpc_server import GRPCServer
 from exo.networking.grpc.grpc_discovery import GRPCDiscovery
 from exo.topology.ring_memory_weighted_partitioning_strategy import RingMemoryWeightedPartitioningStrategy
 from exo.api import ChatGPTAPI
+from exo.helpers import print_yellow_exo
 
 # parse args
 parser = argparse.ArgumentParser(description="Initialize GRPC Discovery")
@@ -22,7 +23,8 @@ parser.add_argument("--wait-for-peers", type=int, default=0, help="Number of pee
 parser.add_argument("--chatgpt-api-port", type=int, default=8000, help="ChatGPT API port")
 args = parser.parse_args()
 
-print(f"Starting {platform.system()=} {psutil.virtual_memory()=}")
+print_yellow_exo()
+print(f"Starting exo {platform.system()=} {psutil.virtual_memory()=}")
 if psutil.MACOS:
     from exo.inference.mlx.sharded_inference_engine import MLXDynamicShardInferenceEngine
     inference_engine = MLXDynamicShardInferenceEngine()
@@ -41,6 +43,8 @@ node.on_token.register("main_log").on_next(lambda _, tokens , __: print(inferenc
 async def shutdown(signal, loop):
     """Gracefully shutdown the server and close the asyncio loop."""
     print(f"Received exit signal {signal.name}...")
+    print("Thank you for using exo.")
+    print_yellow_exo()
     server_tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
     [task.cancel() for task in server_tasks]
     print(f"Cancelling {len(server_tasks)} outstanding tasks")