gpu_status.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114
  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. "--columns",
  42. type=int,
  43. default=-1,
  44. help="The number of columns the node information is formatted in. By default, this is set by the compactness style (default, --compact, --ultracompact). Changing this value might break formatting."
  45. )
  46. parser.add_argument(
  47. "--pbar",
  48. type=int,
  49. default=-1,
  50. help="The length of the progress bar. By default, this is set by the compactness style (default, --compact, --ultracompact). Changing this value might break formatting."
  51. )
  52. parser.add_argument(
  53. "-uc",
  54. "--ultracompact",
  55. action="store_true",
  56. help="Flag; if set, a very compact overview will be displayed."
  57. )
  58. parser.add_argument(
  59. "-u",
  60. "--users",
  61. action="store_true",
  62. help="Flag; if set, the current users of each GPU will be displayed."
  63. )
  64. parser.add_argument(
  65. "-f",
  66. "--filter",
  67. type=str,
  68. nargs="*",
  69. default=None,
  70. help="The node names of nodes to filter for."
  71. )
  72. args = parser.parse_args()
  73. response = requests.get(
  74. f"{args.protocol}://{args.server}:{args.port}/api/clients/"
  75. )
  76. if args.ultracompact:
  77. renderer = Renderer(
  78. columns=3 if args.columns == -1 else args.columns,
  79. progress_bar_width=20 if args.pbar == -1 else args.pbar,
  80. gpu_inner_spacings=False,
  81. gpu_outer_spacings=False,
  82. node_names=args.filter,
  83. display_power=False,
  84. display_users=args.users,
  85. dict_type="shortened",
  86. )
  87. elif args.compact:
  88. renderer = Renderer(
  89. columns=2 if args.columns == -1 else args.columns,
  90. progress_bar_width=40 if args.pbar == -1 else args.pbar,
  91. gpu_inner_spacings=True,
  92. gpu_outer_spacings=False,
  93. node_names=args.filter,
  94. display_power=True,
  95. display_users=args.users,
  96. )
  97. else:
  98. renderer = Renderer(
  99. columns=1 if args.columns == -1 else args.columns,
  100. progress_bar_width=50 if args.pbar == -1 else args.pbar,
  101. gpu_inner_spacings=True,
  102. gpu_outer_spacings=True,
  103. node_names=args.filter,
  104. display_users=args.users,
  105. )
  106. print(renderer.render_info_dict(response.json()))