cli_utils.py 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556
  1. import functools
  2. from typing import Union
  3. import click
  4. def bool_cast(string: str) -> Union[bool, str]:
  5. """Cast a string to a boolean if possible, otherwise return the string."""
  6. if string.lower() == "true" or string == "1":
  7. return True
  8. elif string.lower() == "false" or string == "0":
  9. return False
  10. else:
  11. return string
  12. class BoolOrStringParam(click.ParamType):
  13. """A click parameter that can be either a boolean or a string."""
  14. name = "BOOL | TEXT"
  15. def convert(self, value, param, ctx):
  16. if isinstance(value, bool):
  17. return value
  18. else:
  19. return bool_cast(value)
  20. def add_common_job_options(func):
  21. """Decorator for adding CLI flags shared by all `ray job` commands."""
  22. @click.option(
  23. "--verify",
  24. default=True,
  25. show_default=True,
  26. type=BoolOrStringParam(),
  27. help=(
  28. "Boolean indication to verify the server's TLS certificate or a path to"
  29. " a file or directory of trusted certificates."
  30. ),
  31. )
  32. @click.option(
  33. "--headers",
  34. required=False,
  35. type=str,
  36. default=None,
  37. help=(
  38. "Used to pass headers through http/s to the Ray Cluster."
  39. 'please follow JSON formatting formatting {"key": "value"}'
  40. ),
  41. )
  42. @functools.wraps(func)
  43. def wrapper(*args, **kwargs):
  44. return func(*args, **kwargs)
  45. return wrapper