gpu_status.py 2.2 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798
  1. import argparse
  2. import requests
  3. from renderer import Renderer
  4. parser = argparse.ArgumentParser(description="Prints the status of server nodes and their GPUs.")
  5. parser.add_argument(
  6. "watch",
  7. type=str,
  8. nargs="?",
  9. choices=("watch",),
  10. help="If the word 'watch' is supplied as first argument, the script will be run in watch mode. Supplying this argument in any other position will have no effect."
  11. )
  12. parser.add_argument(
  13. "-s",
  14. "--server",
  15. type=str,
  16. default="deimos",
  17. help="The name of the server to query."
  18. )
  19. parser.add_argument(
  20. "-p",
  21. "--port",
  22. type=str,
  23. default="8091",
  24. help="The port of the server to query.",
  25. )
  26. parser.add_argument(
  27. "-t",
  28. "--protocol",
  29. type=str,
  30. default="http",
  31. choices=("http", "https"),
  32. help="The protocol to use for the query.",
  33. )
  34. parser.add_argument(
  35. "-c",
  36. "--compact",
  37. action="store_true",
  38. help="Flag; if set, a compact overview will be displayed."
  39. )
  40. parser.add_argument(
  41. "-uc",
  42. "--ultracompact",
  43. action="store_true",
  44. help="Flag; if set, a very compact overview will be displayed."
  45. )
  46. parser.add_argument(
  47. "-u",
  48. "--users",
  49. action="store_true",
  50. help="Flag; if set, the current users of each GPU will be displayed."
  51. )
  52. parser.add_argument(
  53. "-f",
  54. "--filter",
  55. type=str,
  56. nargs="*",
  57. default=None,
  58. help="The node names of nodes to filter for."
  59. )
  60. args = parser.parse_args()
  61. response = requests.get(
  62. f"{args.protocol}://{args.server}:{args.port}/api/clients/"
  63. )
  64. if args.ultracompact:
  65. renderer = Renderer(
  66. columns=3,
  67. progress_bar_width=30,
  68. use_space_lines=False,
  69. node_names=args.filter,
  70. display_power=False,
  71. display_users=args.users,
  72. )
  73. elif args.compact:
  74. renderer = Renderer(
  75. columns=2,
  76. progress_bar_width=40,
  77. use_space_lines=False,
  78. node_names=args.filter,
  79. display_power=True,
  80. display_users=args.users,
  81. )
  82. else:
  83. renderer = Renderer(
  84. columns=1,
  85. progress_bar_width=50,
  86. use_space_lines=True,
  87. node_names=args.filter,
  88. display_users=args.users,
  89. )
  90. print(renderer.render_info_dict(response.json()))