Jelajahi Sumber

restore the cursor to the terminal on exit from CLI

Daniel Newman 6 bulan lalu
induk
melakukan
729669c92a
1 mengubah file dengan 10 tambahan dan 0 penghapusan
  1. 10 0
      exo/main.py

+ 10 - 0
exo/main.py

@@ -9,6 +9,7 @@ import sys
 import time
 import traceback
 import uuid
+import atexit
 from exo.networking.manual.manual_discovery import ManualDiscovery
 from exo.networking.manual.network_topology_config import NetworkTopology
 from exo.orchestration.standard_node import StandardNode
@@ -29,6 +30,8 @@ from exo.models import build_base_shard, get_repo
 from exo.viz.topology_viz import TopologyViz
 from exo.download.hf.hf_helpers import has_hf_home_read_access, has_hf_home_write_access, get_hf_home, move_models_to_hf
 
+
+
 # parse args
 parser = argparse.ArgumentParser(description="Initialize GRPC Discovery")
 parser.add_argument("command", nargs="?", choices=["run"], help="Command to run")
@@ -215,6 +218,13 @@ async def main():
     except Exception as e:
       print(f"Error moving models to .cache/huggingface: {e}")
 
+  def restore_cursor():
+    if platform.system() != "Windows":
+        os.system("tput cnorm")  # Show cursor
+
+  # Restore the cursor when the program exits
+  atexit.register(restore_cursor)
+
   # Use a more direct approach to handle signals
   def handle_exit():
     asyncio.ensure_future(shutdown(signal.SIGTERM, loop, node.server))