wandb_require.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. """Feature Flags Module.
  2. This module implements a feature flag system for the wandb library to require experimental features
  3. and notify the user when features have been deprecated.
  4. Example:
  5. import wandb
  6. wandb.require("wandb-service@beta")
  7. wandb.require("incremental-artifacts@beta")
  8. """
  9. from __future__ import annotations
  10. from collections.abc import Iterable
  11. import wandb
  12. from wandb.errors import UnsupportedError
  13. class _Requires:
  14. """Internal feature class."""
  15. _features: tuple[str, ...]
  16. def __init__(self, features: str | Iterable[str]) -> None:
  17. self._features = (
  18. tuple([features]) if isinstance(features, str) else tuple(features)
  19. )
  20. def require_require(self) -> None:
  21. pass
  22. def require_service(self) -> None:
  23. # Legacy no-op kept solely for backward compatibility:
  24. # some integrations (e.g. PyTorch Lightning) still call
  25. # `wandb.require('service')`, which routes here.
  26. wandb.termwarn(
  27. "`wandb.require('service')` is a no-op as it is now the default behavior."
  28. )
  29. def require_core(self) -> None:
  30. # Legacy no-op kept solely for backward compatibility:
  31. # many public codebases still call `wandb.require('core')`.
  32. wandb.termwarn(
  33. "`wandb.require('core')` is a no-op as it is now the default behavior."
  34. )
  35. def apply(self) -> None:
  36. """Call require_* method for supported features."""
  37. last_message: str = ""
  38. for feature_item in self._features:
  39. full_feature = feature_item.split("@", 2)[0]
  40. feature = full_feature.split(":", 2)[0]
  41. func_str = "require_{}".format(feature.replace("-", "_"))
  42. func = getattr(self, func_str, None)
  43. if not func:
  44. last_message = f"require() unsupported requirement: {feature}"
  45. wandb.termwarn(last_message)
  46. continue
  47. func()
  48. if last_message:
  49. raise UnsupportedError(last_message)
  50. def require(
  51. requirement: str | Iterable[str] | None = None,
  52. experiment: str | Iterable[str] | None = None,
  53. ) -> None:
  54. """Indicate which experimental features are used by the script.
  55. This should be called before any other `wandb` functions, ideally right
  56. after importing `wandb`.
  57. Args:
  58. requirement: The name of a feature to require or an iterable of
  59. feature names.
  60. experiment: An alias for `requirement`.
  61. Raises:
  62. wandb.errors.UnsupportedError: If a feature name is unknown.
  63. """
  64. features = requirement or experiment
  65. if not features:
  66. return
  67. f = _Requires(features=features)
  68. f.apply()