ob-metaflow-stubs 6.0.3.186__py2.py3-none-any.whl → 6.0.3.187__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (216) hide show
  1. metaflow-stubs/__init__.pyi +714 -702
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +45 -45
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +1 -1
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +1 -1
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +4 -4
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  88. metaflow-stubs/multicore_utils.pyi +1 -1
  89. metaflow-stubs/ob_internal.pyi +1 -1
  90. metaflow-stubs/parameters.pyi +2 -2
  91. metaflow-stubs/plugins/__init__.pyi +14 -14
  92. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  93. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  94. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  96. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  99. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  100. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  101. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  103. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  104. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  105. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  106. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  108. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  109. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  110. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  111. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  112. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  113. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  114. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  115. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  116. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  117. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  118. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  119. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  121. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  122. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  123. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  124. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  125. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  126. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  127. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  128. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  129. metaflow-stubs/plugins/cards/__init__.pyi +5 -5
  130. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  131. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  133. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  135. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  136. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  138. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  140. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  142. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  143. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  144. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  145. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  146. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  148. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  149. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  151. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  152. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  153. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  154. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  155. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  156. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  157. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  159. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  160. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  161. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  163. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  164. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  165. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  169. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  170. metaflow-stubs/plugins/ollama/__init__.pyi +1 -1
  171. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  172. metaflow-stubs/plugins/perimeters.pyi +1 -1
  173. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  174. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  175. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  177. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  179. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  181. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  182. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  184. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  185. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  186. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  187. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  188. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  189. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  190. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  192. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  193. metaflow-stubs/profilers/__init__.pyi +1 -1
  194. metaflow-stubs/pylint_wrapper.pyi +1 -1
  195. metaflow-stubs/runner/__init__.pyi +1 -1
  196. metaflow-stubs/runner/deployer.pyi +4 -4
  197. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  198. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  199. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  200. metaflow-stubs/runner/nbrun.pyi +1 -1
  201. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  202. metaflow-stubs/runner/utils.pyi +3 -3
  203. metaflow-stubs/system/__init__.pyi +1 -1
  204. metaflow-stubs/system/system_logger.pyi +2 -2
  205. metaflow-stubs/system/system_monitor.pyi +1 -1
  206. metaflow-stubs/tagging_util.pyi +1 -1
  207. metaflow-stubs/tuple_util.pyi +1 -1
  208. metaflow-stubs/user_configs/__init__.pyi +1 -1
  209. metaflow-stubs/user_configs/config_decorators.pyi +4 -4
  210. metaflow-stubs/user_configs/config_options.pyi +2 -2
  211. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  212. {ob_metaflow_stubs-6.0.3.186.dist-info → ob_metaflow_stubs-6.0.3.187.dist-info}/METADATA +1 -1
  213. ob_metaflow_stubs-6.0.3.187.dist-info/RECORD +216 -0
  214. ob_metaflow_stubs-6.0.3.186.dist-info/RECORD +0 -216
  215. {ob_metaflow_stubs-6.0.3.186.dist-info → ob_metaflow_stubs-6.0.3.187.dist-info}/WHEEL +0 -0
  216. {ob_metaflow_stubs-6.0.3.186.dist-info → ob_metaflow_stubs-6.0.3.187.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.18.1+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-07-01T08:40:33.108391 #
4
+ # Generated on 2025-07-01T15:21:03.639611 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,18 +35,18 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
+ from . import events as events
38
39
  from . import cards as cards
39
- from . import metaflow_git as metaflow_git
40
40
  from . import tuple_util as tuple_util
41
- from . import events as events
41
+ from . import metaflow_git as metaflow_git
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
48
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
52
52
  from .client.core import get_namespace as get_namespace
@@ -156,185 +156,65 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
156
156
  ...
157
157
 
158
158
  @typing.overload
159
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
160
- """
161
- Specifies the PyPI packages for the step.
162
-
163
- Information in this decorator will augment any
164
- attributes set in the `@pyi_base` flow-level decorator. Hence,
165
- you can use `@pypi_base` to set packages required by all
166
- steps and use `@pypi` to specify step-specific overrides.
167
-
168
-
169
- Parameters
170
- ----------
171
- packages : Dict[str, str], default: {}
172
- Packages to use for this step. The key is the name of the package
173
- and the value is the version to use.
174
- python : str, optional, default: None
175
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
176
- that the version used will correspond to the version of the Python interpreter used to start the run.
177
- """
178
- ...
179
-
180
- @typing.overload
181
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
182
- ...
183
-
184
- @typing.overload
185
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
186
- ...
187
-
188
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
159
+ def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
189
160
  """
190
- Specifies the PyPI packages for the step.
161
+ Creates a human-readable report, a Metaflow Card, after this step completes.
191
162
 
192
- Information in this decorator will augment any
193
- attributes set in the `@pyi_base` flow-level decorator. Hence,
194
- you can use `@pypi_base` to set packages required by all
195
- steps and use `@pypi` to specify step-specific overrides.
163
+ Note that you may add multiple `@card` decorators in a step with different parameters.
196
164
 
197
165
 
198
166
  Parameters
199
167
  ----------
200
- packages : Dict[str, str], default: {}
201
- Packages to use for this step. The key is the name of the package
202
- and the value is the version to use.
203
- python : str, optional, default: None
204
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
205
- that the version used will correspond to the version of the Python interpreter used to start the run.
168
+ type : str, default 'default'
169
+ Card type.
170
+ id : str, optional, default None
171
+ If multiple cards are present, use this id to identify this card.
172
+ options : Dict[str, Any], default {}
173
+ Options passed to the card. The contents depend on the card type.
174
+ timeout : int, default 45
175
+ Interrupt reporting if it takes more than this many seconds.
206
176
  """
207
177
  ...
208
178
 
209
179
  @typing.overload
210
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
211
- """
212
- Internal decorator to support Fast bakery
213
- """
180
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
214
181
  ...
215
182
 
216
183
  @typing.overload
217
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
218
- ...
219
-
220
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
221
- """
222
- Internal decorator to support Fast bakery
223
- """
184
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
224
185
  ...
225
186
 
226
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
187
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
227
188
  """
228
- Decorator that helps cache, version and store models/datasets from huggingface hub.
229
-
230
- > Examples
231
-
232
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
233
- ```python
234
- @huggingface_hub
235
- @step
236
- def pull_model_from_huggingface(self):
237
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
238
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
239
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
240
- # value of the function is a reference to the model in the backend storage.
241
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
242
-
243
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
244
- self.llama_model = current.huggingface_hub.snapshot_download(
245
- repo_id=self.model_id,
246
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
247
- )
248
- self.next(self.train)
249
- ```
250
-
251
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
252
- ```python
253
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
254
- @step
255
- def pull_model_from_huggingface(self):
256
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
257
- ```
258
-
259
- ```python
260
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
261
- @step
262
- def finetune_model(self):
263
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
264
- # path_to_model will be /my-directory
265
- ```
266
-
267
- ```python
268
- # Takes all the arguments passed to `snapshot_download`
269
- # except for `local_dir`
270
- @huggingface_hub(load=[
271
- {
272
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
273
- },
274
- {
275
- "repo_id": "myorg/mistral-lora",
276
- "repo_type": "model",
277
- },
278
- ])
279
- @step
280
- def finetune_model(self):
281
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
282
- # path_to_model will be /my-directory
283
- ```
284
-
285
-
286
- Parameters
287
- ----------
288
- temp_dir_root : str, optional
289
- The root directory that will hold the temporary directory where objects will be downloaded.
290
-
291
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
292
- The list of repos (models/datasets) to load.
293
-
294
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
295
-
296
- - If repo (model/dataset) is not found in the datastore:
297
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
298
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
299
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
189
+ Creates a human-readable report, a Metaflow Card, after this step completes.
300
190
 
301
- - If repo is found in the datastore:
302
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
303
- """
304
- ...
305
-
306
- @typing.overload
307
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
308
- """
309
- Specifies secrets to be retrieved and injected as environment variables prior to
310
- the execution of a step.
191
+ Note that you may add multiple `@card` decorators in a step with different parameters.
311
192
 
312
193
 
313
194
  Parameters
314
195
  ----------
315
- sources : List[Union[str, Dict[str, Any]]], default: []
316
- List of secret specs, defining how the secrets are to be retrieved
196
+ type : str, default 'default'
197
+ Card type.
198
+ id : str, optional, default None
199
+ If multiple cards are present, use this id to identify this card.
200
+ options : Dict[str, Any], default {}
201
+ Options passed to the card. The contents depend on the card type.
202
+ timeout : int, default 45
203
+ Interrupt reporting if it takes more than this many seconds.
317
204
  """
318
205
  ...
319
206
 
320
- @typing.overload
321
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
322
- ...
323
-
324
- @typing.overload
325
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
326
- ...
327
-
328
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
207
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
329
208
  """
330
- Specifies secrets to be retrieved and injected as environment variables prior to
331
- the execution of a step.
209
+ Specifies that this step should execute on DGX cloud.
332
210
 
333
211
 
334
212
  Parameters
335
213
  ----------
336
- sources : List[Union[str, Dict[str, Any]]], default: []
337
- List of secret specs, defining how the secrets are to be retrieved
214
+ gpu : int
215
+ Number of GPUs to use.
216
+ gpu_type : str
217
+ Type of Nvidia GPU to use.
338
218
  """
339
219
  ...
340
220
 
@@ -389,91 +269,92 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
389
269
  """
390
270
  ...
391
271
 
392
- @typing.overload
393
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
394
- """
395
- Specifies a timeout for your step.
396
-
397
- This decorator is useful if this step may hang indefinitely.
398
-
399
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
400
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
401
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
402
-
403
- Note that all the values specified in parameters are added together so if you specify
404
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
405
-
406
-
407
- Parameters
408
- ----------
409
- seconds : int, default 0
410
- Number of seconds to wait prior to timing out.
411
- minutes : int, default 0
412
- Number of minutes to wait prior to timing out.
413
- hours : int, default 0
414
- Number of hours to wait prior to timing out.
415
- """
416
- ...
417
-
418
- @typing.overload
419
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
420
- ...
421
-
422
- @typing.overload
423
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
424
- ...
425
-
426
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
427
- """
428
- Specifies a timeout for your step.
429
-
430
- This decorator is useful if this step may hang indefinitely.
431
-
432
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
433
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
434
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
435
-
436
- Note that all the values specified in parameters are added together so if you specify
437
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
438
-
439
-
440
- Parameters
441
- ----------
442
- seconds : int, default 0
443
- Number of seconds to wait prior to timing out.
444
- minutes : int, default 0
445
- Number of minutes to wait prior to timing out.
446
- hours : int, default 0
447
- Number of hours to wait prior to timing out.
448
- """
449
- ...
450
-
451
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
272
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
452
273
  """
453
- Specifies that this step should execute on DGX cloud.
274
+ Specifies that this step should execute on Kubernetes.
454
275
 
455
276
 
456
277
  Parameters
457
278
  ----------
458
- gpu : int
459
- Number of GPUs to use.
460
- gpu_type : str
461
- Type of Nvidia GPU to use.
462
- """
463
- ...
464
-
465
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
466
- """
467
- Specifies that this step is used to deploy an instance of the app.
468
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
469
-
279
+ cpu : int, default 1
280
+ Number of CPUs required for this step. If `@resources` is
281
+ also present, the maximum value from all decorators is used.
282
+ memory : int, default 4096
283
+ Memory size (in MB) required for this step. If
284
+ `@resources` is also present, the maximum value from all decorators is
285
+ used.
286
+ disk : int, default 10240
287
+ Disk size (in MB) required for this step. If
288
+ `@resources` is also present, the maximum value from all decorators is
289
+ used.
290
+ image : str, optional, default None
291
+ Docker image to use when launching on Kubernetes. If not specified, and
292
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
293
+ not, a default Docker image mapping to the current version of Python is used.
294
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
295
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
296
+ image_pull_secrets: List[str], default []
297
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
298
+ Kubernetes image pull secrets to use when pulling container images
299
+ in Kubernetes.
300
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
301
+ Kubernetes service account to use when launching pod in Kubernetes.
302
+ secrets : List[str], optional, default None
303
+ Kubernetes secrets to use when launching pod in Kubernetes. These
304
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
305
+ in Metaflow configuration.
306
+ node_selector: Union[Dict[str,str], str], optional, default None
307
+ Kubernetes node selector(s) to apply to the pod running the task.
308
+ Can be passed in as a comma separated string of values e.g.
309
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
310
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
311
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
312
+ Kubernetes namespace to use when launching pod in Kubernetes.
313
+ gpu : int, optional, default None
314
+ Number of GPUs required for this step. A value of zero implies that
315
+ the scheduled node should not have GPUs.
316
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
317
+ The vendor of the GPUs to be used for this step.
318
+ tolerations : List[str], default []
319
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
320
+ Kubernetes tolerations to use when launching pod in Kubernetes.
321
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
322
+ Kubernetes labels to use when launching pod in Kubernetes.
323
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
324
+ Kubernetes annotations to use when launching pod in Kubernetes.
325
+ use_tmpfs : bool, default False
326
+ This enables an explicit tmpfs mount for this step.
327
+ tmpfs_tempdir : bool, default True
328
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
329
+ tmpfs_size : int, optional, default: None
330
+ The value for the size (in MiB) of the tmpfs mount for this step.
331
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
332
+ memory allocated for this step.
333
+ tmpfs_path : str, optional, default /metaflow_temp
334
+ Path to tmpfs mount for this step.
335
+ persistent_volume_claims : Dict[str, str], optional, default None
336
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
337
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
338
+ shared_memory: int, optional
339
+ Shared memory size (in MiB) required for this step
340
+ port: int, optional
341
+ Port number to specify in the Kubernetes job object
342
+ compute_pool : str, optional, default None
343
+ Compute pool to be used for for this step.
344
+ If not specified, any accessible compute pool within the perimeter is used.
345
+ hostname_resolution_timeout: int, default 10 * 60
346
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
347
+ Only applicable when @parallel is used.
348
+ qos: str, default: Burstable
349
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
470
350
 
471
- Parameters
472
- ----------
473
- app_port : int
474
- Number of GPUs to use.
475
- app_name : str
476
- Name of the app to deploy.
351
+ security_context: Dict[str, Any], optional, default None
352
+ Container security context. Applies to the task container. Allows the following keys:
353
+ - privileged: bool, optional, default None
354
+ - allow_privilege_escalation: bool, optional, default None
355
+ - run_as_user: int, optional, default None
356
+ - run_as_group: int, optional, default None
357
+ - run_as_non_root: bool, optional, default None
477
358
  """
478
359
  ...
479
360
 
@@ -493,6 +374,187 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
493
374
  """
494
375
  ...
495
376
 
377
+ @typing.overload
378
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
379
+ """
380
+ Specifies the resources needed when executing this step.
381
+
382
+ Use `@resources` to specify the resource requirements
383
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
384
+
385
+ You can choose the compute layer on the command line by executing e.g.
386
+ ```
387
+ python myflow.py run --with batch
388
+ ```
389
+ or
390
+ ```
391
+ python myflow.py run --with kubernetes
392
+ ```
393
+ which executes the flow on the desired system using the
394
+ requirements specified in `@resources`.
395
+
396
+
397
+ Parameters
398
+ ----------
399
+ cpu : int, default 1
400
+ Number of CPUs required for this step.
401
+ gpu : int, optional, default None
402
+ Number of GPUs required for this step.
403
+ disk : int, optional, default None
404
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
405
+ memory : int, default 4096
406
+ Memory size (in MB) required for this step.
407
+ shared_memory : int, optional, default None
408
+ The value for the size (in MiB) of the /dev/shm volume for this step.
409
+ This parameter maps to the `--shm-size` option in Docker.
410
+ """
411
+ ...
412
+
413
+ @typing.overload
414
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
415
+ ...
416
+
417
+ @typing.overload
418
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
419
+ ...
420
+
421
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
422
+ """
423
+ Specifies the resources needed when executing this step.
424
+
425
+ Use `@resources` to specify the resource requirements
426
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
427
+
428
+ You can choose the compute layer on the command line by executing e.g.
429
+ ```
430
+ python myflow.py run --with batch
431
+ ```
432
+ or
433
+ ```
434
+ python myflow.py run --with kubernetes
435
+ ```
436
+ which executes the flow on the desired system using the
437
+ requirements specified in `@resources`.
438
+
439
+
440
+ Parameters
441
+ ----------
442
+ cpu : int, default 1
443
+ Number of CPUs required for this step.
444
+ gpu : int, optional, default None
445
+ Number of GPUs required for this step.
446
+ disk : int, optional, default None
447
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
448
+ memory : int, default 4096
449
+ Memory size (in MB) required for this step.
450
+ shared_memory : int, optional, default None
451
+ The value for the size (in MiB) of the /dev/shm volume for this step.
452
+ This parameter maps to the `--shm-size` option in Docker.
453
+ """
454
+ ...
455
+
456
+ @typing.overload
457
+ def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
458
+ """
459
+ Specifies environment variables to be set prior to the execution of a step.
460
+
461
+
462
+ Parameters
463
+ ----------
464
+ vars : Dict[str, str], default {}
465
+ Dictionary of environment variables to set.
466
+ """
467
+ ...
468
+
469
+ @typing.overload
470
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
471
+ ...
472
+
473
+ @typing.overload
474
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
475
+ ...
476
+
477
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
478
+ """
479
+ Specifies environment variables to be set prior to the execution of a step.
480
+
481
+
482
+ Parameters
483
+ ----------
484
+ vars : Dict[str, str], default {}
485
+ Dictionary of environment variables to set.
486
+ """
487
+ ...
488
+
489
+ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card_refresh_interval: int, max_retries: int, retry_alert_frequency: int, engine_args: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
490
+ """
491
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
492
+
493
+ User code call
494
+ --------------
495
+ @vllm(
496
+ model="...",
497
+ ...
498
+ )
499
+
500
+ Valid backend options
501
+ ---------------------
502
+ - 'local': Run as a separate process on the local task machine.
503
+
504
+ Valid model options
505
+ -------------------
506
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
507
+
508
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
509
+ If you need multiple models, you must create multiple @vllm decorators.
510
+
511
+
512
+ Parameters
513
+ ----------
514
+ model: str
515
+ HuggingFace model identifier to be served by vLLM.
516
+ backend: str
517
+ Determines where and how to run the vLLM process.
518
+ openai_api_server: bool
519
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
520
+ Default is False (uses native engine).
521
+ Set to True for backward compatibility with existing code.
522
+ debug: bool
523
+ Whether to turn on verbose debugging logs.
524
+ card_refresh_interval: int
525
+ Interval in seconds for refreshing the vLLM status card.
526
+ Only used when openai_api_server=True.
527
+ max_retries: int
528
+ Maximum number of retries checking for vLLM server startup.
529
+ Only used when openai_api_server=True.
530
+ retry_alert_frequency: int
531
+ Frequency of alert logs for vLLM server startup retries.
532
+ Only used when openai_api_server=True.
533
+ engine_args : dict
534
+ Additional keyword arguments to pass to the vLLM engine.
535
+ For example, `tensor_parallel_size=2`.
536
+ """
537
+ ...
538
+
539
+ @typing.overload
540
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
541
+ """
542
+ Decorator prototype for all step decorators. This function gets specialized
543
+ and imported for all decorators types by _import_plugin_decorators().
544
+ """
545
+ ...
546
+
547
+ @typing.overload
548
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
549
+ ...
550
+
551
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
552
+ """
553
+ Decorator prototype for all step decorators. This function gets specialized
554
+ and imported for all decorators types by _import_plugin_decorators().
555
+ """
556
+ ...
557
+
496
558
  def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
497
559
  """
498
560
  This decorator is used to run Ollama APIs as Metaflow task sidecars.
@@ -537,84 +599,118 @@ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy:
537
599
  ...
538
600
 
539
601
  @typing.overload
540
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
602
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
541
603
  """
542
- Enables checkpointing for a step.
604
+ Specifies secrets to be retrieved and injected as environment variables prior to
605
+ the execution of a step.
543
606
 
544
- > Examples
545
607
 
546
- - Saving Checkpoints
608
+ Parameters
609
+ ----------
610
+ sources : List[Union[str, Dict[str, Any]]], default: []
611
+ List of secret specs, defining how the secrets are to be retrieved
612
+ """
613
+ ...
614
+
615
+ @typing.overload
616
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
617
+ ...
618
+
619
+ @typing.overload
620
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
621
+ ...
622
+
623
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
624
+ """
625
+ Specifies secrets to be retrieved and injected as environment variables prior to
626
+ the execution of a step.
547
627
 
548
- ```python
549
- @checkpoint
550
- @step
551
- def train(self):
552
- model = create_model(self.parameters, checkpoint_path = None)
553
- for i in range(self.epochs):
554
- # some training logic
555
- loss = model.train(self.dataset)
556
- if i % 10 == 0:
557
- model.save(
558
- current.checkpoint.directory,
559
- )
560
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
561
- # and returns a reference dictionary to the checkpoint saved in the datastore
562
- self.latest_checkpoint = current.checkpoint.save(
563
- name="epoch_checkpoint",
564
- metadata={
565
- "epoch": i,
566
- "loss": loss,
567
- }
568
- )
569
- ```
570
628
 
571
- - Using Loaded Checkpoints
629
+ Parameters
630
+ ----------
631
+ sources : List[Union[str, Dict[str, Any]]], default: []
632
+ List of secret specs, defining how the secrets are to be retrieved
633
+ """
634
+ ...
635
+
636
+ @typing.overload
637
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
638
+ """
639
+ Specifies a timeout for your step.
572
640
 
573
- ```python
574
- @retry(times=3)
575
- @checkpoint
576
- @step
577
- def train(self):
578
- # Assume that the task has restarted and the previous attempt of the task
579
- # saved a checkpoint
580
- checkpoint_path = None
581
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
582
- print("Loaded checkpoint from the previous attempt")
583
- checkpoint_path = current.checkpoint.directory
641
+ This decorator is useful if this step may hang indefinitely.
584
642
 
585
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
586
- for i in range(self.epochs):
587
- ...
588
- ```
643
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
644
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
645
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
646
+
647
+ Note that all the values specified in parameters are added together so if you specify
648
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
649
+
650
+
651
+ Parameters
652
+ ----------
653
+ seconds : int, default 0
654
+ Number of seconds to wait prior to timing out.
655
+ minutes : int, default 0
656
+ Number of minutes to wait prior to timing out.
657
+ hours : int, default 0
658
+ Number of hours to wait prior to timing out.
659
+ """
660
+ ...
661
+
662
+ @typing.overload
663
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
664
+ ...
665
+
666
+ @typing.overload
667
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
668
+ ...
669
+
670
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
671
+ """
672
+ Specifies a timeout for your step.
673
+
674
+ This decorator is useful if this step may hang indefinitely.
675
+
676
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
677
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
678
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
679
+
680
+ Note that all the values specified in parameters are added together so if you specify
681
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
589
682
 
590
683
 
591
684
  Parameters
592
685
  ----------
593
- load_policy : str, default: "fresh"
594
- The policy for loading the checkpoint. The following policies are supported:
595
- - "eager": Loads the the latest available checkpoint within the namespace.
596
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
597
- will be loaded at the start of the task.
598
- - "none": Do not load any checkpoint
599
- - "fresh": Loads the lastest checkpoint created within the running Task.
600
- This mode helps loading checkpoints across various retry attempts of the same task.
601
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
602
- created within the task will be loaded when the task is retries execution on failure.
603
-
604
- temp_dir_root : str, default: None
605
- The root directory under which `current.checkpoint.directory` will be created.
686
+ seconds : int, default 0
687
+ Number of seconds to wait prior to timing out.
688
+ minutes : int, default 0
689
+ Number of minutes to wait prior to timing out.
690
+ hours : int, default 0
691
+ Number of hours to wait prior to timing out.
606
692
  """
607
693
  ...
608
694
 
609
695
  @typing.overload
610
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
696
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
697
+ """
698
+ Internal decorator to support Fast bakery
699
+ """
611
700
  ...
612
701
 
613
702
  @typing.overload
614
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
703
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
615
704
  ...
616
705
 
617
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
706
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
707
+ """
708
+ Internal decorator to support Fast bakery
709
+ """
710
+ ...
711
+
712
+ @typing.overload
713
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
618
714
  """
619
715
  Enables checkpointing for a step.
620
716
 
@@ -672,288 +768,91 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
672
768
  - "eager": Loads the the latest available checkpoint within the namespace.
673
769
  With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
674
770
  will be loaded at the start of the task.
675
- - "none": Do not load any checkpoint
676
- - "fresh": Loads the lastest checkpoint created within the running Task.
677
- This mode helps loading checkpoints across various retry attempts of the same task.
678
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
679
- created within the task will be loaded when the task is retries execution on failure.
680
-
681
- temp_dir_root : str, default: None
682
- The root directory under which `current.checkpoint.directory` will be created.
683
- """
684
- ...
685
-
686
- def vllm(*, model: str, backend: str, debug: bool, kwargs: typing.Any) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
687
- """
688
- This decorator is used to run vllm APIs as Metaflow task sidecars.
689
-
690
- User code call
691
- --------------
692
- @vllm(
693
- model="...",
694
- ...
695
- )
696
-
697
- Valid backend options
698
- ---------------------
699
- - 'local': Run as a separate process on the local task machine.
700
-
701
- Valid model options
702
- -------------------
703
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
704
-
705
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
706
- If you need multiple models, you must create multiple @vllm decorators.
707
-
708
-
709
- Parameters
710
- ----------
711
- model: str
712
- HuggingFace model identifier to be served by vLLM.
713
- backend: str
714
- Determines where and how to run the vLLM process.
715
- debug: bool
716
- Whether to turn on verbose debugging logs.
717
- kwargs : Any
718
- Any other keyword arguments are passed directly to the vLLM engine.
719
- This allows for flexible configuration of vLLM server settings.
720
- For example, `tensor_parallel_size=2`.
721
- """
722
- ...
723
-
724
- @typing.overload
725
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
726
- """
727
- Specifies the resources needed when executing this step.
728
-
729
- Use `@resources` to specify the resource requirements
730
- independently of the specific compute layer (`@batch`, `@kubernetes`).
731
-
732
- You can choose the compute layer on the command line by executing e.g.
733
- ```
734
- python myflow.py run --with batch
735
- ```
736
- or
737
- ```
738
- python myflow.py run --with kubernetes
739
- ```
740
- which executes the flow on the desired system using the
741
- requirements specified in `@resources`.
742
-
743
-
744
- Parameters
745
- ----------
746
- cpu : int, default 1
747
- Number of CPUs required for this step.
748
- gpu : int, optional, default None
749
- Number of GPUs required for this step.
750
- disk : int, optional, default None
751
- Disk size (in MB) required for this step. Only applies on Kubernetes.
752
- memory : int, default 4096
753
- Memory size (in MB) required for this step.
754
- shared_memory : int, optional, default None
755
- The value for the size (in MiB) of the /dev/shm volume for this step.
756
- This parameter maps to the `--shm-size` option in Docker.
757
- """
758
- ...
759
-
760
- @typing.overload
761
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
762
- ...
763
-
764
- @typing.overload
765
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
766
- ...
767
-
768
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
769
- """
770
- Specifies the resources needed when executing this step.
771
-
772
- Use `@resources` to specify the resource requirements
773
- independently of the specific compute layer (`@batch`, `@kubernetes`).
774
-
775
- You can choose the compute layer on the command line by executing e.g.
776
- ```
777
- python myflow.py run --with batch
778
- ```
779
- or
780
- ```
781
- python myflow.py run --with kubernetes
782
- ```
783
- which executes the flow on the desired system using the
784
- requirements specified in `@resources`.
785
-
786
-
787
- Parameters
788
- ----------
789
- cpu : int, default 1
790
- Number of CPUs required for this step.
791
- gpu : int, optional, default None
792
- Number of GPUs required for this step.
793
- disk : int, optional, default None
794
- Disk size (in MB) required for this step. Only applies on Kubernetes.
795
- memory : int, default 4096
796
- Memory size (in MB) required for this step.
797
- shared_memory : int, optional, default None
798
- The value for the size (in MiB) of the /dev/shm volume for this step.
799
- This parameter maps to the `--shm-size` option in Docker.
800
- """
801
- ...
802
-
803
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
804
- """
805
- Specifies that this step should execute on Kubernetes.
806
-
807
-
808
- Parameters
809
- ----------
810
- cpu : int, default 1
811
- Number of CPUs required for this step. If `@resources` is
812
- also present, the maximum value from all decorators is used.
813
- memory : int, default 4096
814
- Memory size (in MB) required for this step. If
815
- `@resources` is also present, the maximum value from all decorators is
816
- used.
817
- disk : int, default 10240
818
- Disk size (in MB) required for this step. If
819
- `@resources` is also present, the maximum value from all decorators is
820
- used.
821
- image : str, optional, default None
822
- Docker image to use when launching on Kubernetes. If not specified, and
823
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
824
- not, a default Docker image mapping to the current version of Python is used.
825
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
826
- If given, the imagePullPolicy to be applied to the Docker image of the step.
827
- image_pull_secrets: List[str], default []
828
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
829
- Kubernetes image pull secrets to use when pulling container images
830
- in Kubernetes.
831
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
832
- Kubernetes service account to use when launching pod in Kubernetes.
833
- secrets : List[str], optional, default None
834
- Kubernetes secrets to use when launching pod in Kubernetes. These
835
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
836
- in Metaflow configuration.
837
- node_selector: Union[Dict[str,str], str], optional, default None
838
- Kubernetes node selector(s) to apply to the pod running the task.
839
- Can be passed in as a comma separated string of values e.g.
840
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
841
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
842
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
843
- Kubernetes namespace to use when launching pod in Kubernetes.
844
- gpu : int, optional, default None
845
- Number of GPUs required for this step. A value of zero implies that
846
- the scheduled node should not have GPUs.
847
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
848
- The vendor of the GPUs to be used for this step.
849
- tolerations : List[str], default []
850
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
851
- Kubernetes tolerations to use when launching pod in Kubernetes.
852
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
853
- Kubernetes labels to use when launching pod in Kubernetes.
854
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
855
- Kubernetes annotations to use when launching pod in Kubernetes.
856
- use_tmpfs : bool, default False
857
- This enables an explicit tmpfs mount for this step.
858
- tmpfs_tempdir : bool, default True
859
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
860
- tmpfs_size : int, optional, default: None
861
- The value for the size (in MiB) of the tmpfs mount for this step.
862
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
863
- memory allocated for this step.
864
- tmpfs_path : str, optional, default /metaflow_temp
865
- Path to tmpfs mount for this step.
866
- persistent_volume_claims : Dict[str, str], optional, default None
867
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
868
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
869
- shared_memory: int, optional
870
- Shared memory size (in MiB) required for this step
871
- port: int, optional
872
- Port number to specify in the Kubernetes job object
873
- compute_pool : str, optional, default None
874
- Compute pool to be used for for this step.
875
- If not specified, any accessible compute pool within the perimeter is used.
876
- hostname_resolution_timeout: int, default 10 * 60
877
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
878
- Only applicable when @parallel is used.
879
- qos: str, default: Burstable
880
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
881
-
882
- security_context: Dict[str, Any], optional, default None
883
- Container security context. Applies to the task container. Allows the following keys:
884
- - privileged: bool, optional, default None
885
- - allow_privilege_escalation: bool, optional, default None
886
- - run_as_user: int, optional, default None
887
- - run_as_group: int, optional, default None
888
- - run_as_non_root: bool, optional, default None
889
- """
890
- ...
891
-
892
- @typing.overload
893
- def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
894
- """
895
- Creates a human-readable report, a Metaflow Card, after this step completes.
896
-
897
- Note that you may add multiple `@card` decorators in a step with different parameters.
898
-
771
+ - "none": Do not load any checkpoint
772
+ - "fresh": Loads the lastest checkpoint created within the running Task.
773
+ This mode helps loading checkpoints across various retry attempts of the same task.
774
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
775
+ created within the task will be loaded when the task is retries execution on failure.
899
776
 
900
- Parameters
901
- ----------
902
- type : str, default 'default'
903
- Card type.
904
- id : str, optional, default None
905
- If multiple cards are present, use this id to identify this card.
906
- options : Dict[str, Any], default {}
907
- Options passed to the card. The contents depend on the card type.
908
- timeout : int, default 45
909
- Interrupt reporting if it takes more than this many seconds.
777
+ temp_dir_root : str, default: None
778
+ The root directory under which `current.checkpoint.directory` will be created.
910
779
  """
911
780
  ...
912
781
 
913
782
  @typing.overload
914
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
783
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
915
784
  ...
916
785
 
917
786
  @typing.overload
918
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
787
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
919
788
  ...
920
789
 
921
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
790
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
922
791
  """
923
- Creates a human-readable report, a Metaflow Card, after this step completes.
792
+ Enables checkpointing for a step.
924
793
 
925
- Note that you may add multiple `@card` decorators in a step with different parameters.
794
+ > Examples
795
+
796
+ - Saving Checkpoints
797
+
798
+ ```python
799
+ @checkpoint
800
+ @step
801
+ def train(self):
802
+ model = create_model(self.parameters, checkpoint_path = None)
803
+ for i in range(self.epochs):
804
+ # some training logic
805
+ loss = model.train(self.dataset)
806
+ if i % 10 == 0:
807
+ model.save(
808
+ current.checkpoint.directory,
809
+ )
810
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
811
+ # and returns a reference dictionary to the checkpoint saved in the datastore
812
+ self.latest_checkpoint = current.checkpoint.save(
813
+ name="epoch_checkpoint",
814
+ metadata={
815
+ "epoch": i,
816
+ "loss": loss,
817
+ }
818
+ )
819
+ ```
820
+
821
+ - Using Loaded Checkpoints
822
+
823
+ ```python
824
+ @retry(times=3)
825
+ @checkpoint
826
+ @step
827
+ def train(self):
828
+ # Assume that the task has restarted and the previous attempt of the task
829
+ # saved a checkpoint
830
+ checkpoint_path = None
831
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
832
+ print("Loaded checkpoint from the previous attempt")
833
+ checkpoint_path = current.checkpoint.directory
834
+
835
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
836
+ for i in range(self.epochs):
837
+ ...
838
+ ```
926
839
 
927
840
 
928
841
  Parameters
929
842
  ----------
930
- type : str, default 'default'
931
- Card type.
932
- id : str, optional, default None
933
- If multiple cards are present, use this id to identify this card.
934
- options : Dict[str, Any], default {}
935
- Options passed to the card. The contents depend on the card type.
936
- timeout : int, default 45
937
- Interrupt reporting if it takes more than this many seconds.
938
- """
939
- ...
940
-
941
- @typing.overload
942
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
943
- """
944
- Decorator prototype for all step decorators. This function gets specialized
945
- and imported for all decorators types by _import_plugin_decorators().
946
- """
947
- ...
948
-
949
- @typing.overload
950
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
951
- ...
952
-
953
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
954
- """
955
- Decorator prototype for all step decorators. This function gets specialized
956
- and imported for all decorators types by _import_plugin_decorators().
843
+ load_policy : str, default: "fresh"
844
+ The policy for loading the checkpoint. The following policies are supported:
845
+ - "eager": Loads the the latest available checkpoint within the namespace.
846
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
847
+ will be loaded at the start of the task.
848
+ - "none": Do not load any checkpoint
849
+ - "fresh": Loads the lastest checkpoint created within the running Task.
850
+ This mode helps loading checkpoints across various retry attempts of the same task.
851
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
852
+ created within the task will be loaded when the task is retries execution on failure.
853
+
854
+ temp_dir_root : str, default: None
855
+ The root directory under which `current.checkpoint.directory` will be created.
957
856
  """
958
857
  ...
959
858
 
@@ -1012,6 +911,21 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1012
911
  """
1013
912
  ...
1014
913
 
914
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
915
+ """
916
+ Specifies that this step is used to deploy an instance of the app.
917
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
918
+
919
+
920
+ Parameters
921
+ ----------
922
+ app_port : int
923
+ Number of GPUs to use.
924
+ app_name : str
925
+ Name of the app to deploy.
926
+ """
927
+ ...
928
+
1015
929
  @typing.overload
1016
930
  def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1017
931
  """
@@ -1071,6 +985,57 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1071
985
  """
1072
986
  ...
1073
987
 
988
+ @typing.overload
989
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
990
+ """
991
+ Specifies the PyPI packages for the step.
992
+
993
+ Information in this decorator will augment any
994
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
995
+ you can use `@pypi_base` to set packages required by all
996
+ steps and use `@pypi` to specify step-specific overrides.
997
+
998
+
999
+ Parameters
1000
+ ----------
1001
+ packages : Dict[str, str], default: {}
1002
+ Packages to use for this step. The key is the name of the package
1003
+ and the value is the version to use.
1004
+ python : str, optional, default: None
1005
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1006
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1007
+ """
1008
+ ...
1009
+
1010
+ @typing.overload
1011
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1012
+ ...
1013
+
1014
+ @typing.overload
1015
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1016
+ ...
1017
+
1018
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1019
+ """
1020
+ Specifies the PyPI packages for the step.
1021
+
1022
+ Information in this decorator will augment any
1023
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
1024
+ you can use `@pypi_base` to set packages required by all
1025
+ steps and use `@pypi` to specify step-specific overrides.
1026
+
1027
+
1028
+ Parameters
1029
+ ----------
1030
+ packages : Dict[str, str], default: {}
1031
+ Packages to use for this step. The key is the name of the package
1032
+ and the value is the version to use.
1033
+ python : str, optional, default: None
1034
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1035
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1036
+ """
1037
+ ...
1038
+
1074
1039
  @typing.overload
1075
1040
  def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1076
1041
  """
@@ -1200,36 +1165,83 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1200
1165
  """
1201
1166
  ...
1202
1167
 
1203
- @typing.overload
1204
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1168
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1205
1169
  """
1206
- Specifies environment variables to be set prior to the execution of a step.
1170
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1171
+
1172
+ > Examples
1173
+
1174
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
1175
+ ```python
1176
+ @huggingface_hub
1177
+ @step
1178
+ def pull_model_from_huggingface(self):
1179
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
1180
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
1181
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
1182
+ # value of the function is a reference to the model in the backend storage.
1183
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
1184
+
1185
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
1186
+ self.llama_model = current.huggingface_hub.snapshot_download(
1187
+ repo_id=self.model_id,
1188
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1189
+ )
1190
+ self.next(self.train)
1191
+ ```
1192
+
1193
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
1194
+ ```python
1195
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1196
+ @step
1197
+ def pull_model_from_huggingface(self):
1198
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1199
+ ```
1200
+
1201
+ ```python
1202
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1203
+ @step
1204
+ def finetune_model(self):
1205
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1206
+ # path_to_model will be /my-directory
1207
+ ```
1208
+
1209
+ ```python
1210
+ # Takes all the arguments passed to `snapshot_download`
1211
+ # except for `local_dir`
1212
+ @huggingface_hub(load=[
1213
+ {
1214
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1215
+ },
1216
+ {
1217
+ "repo_id": "myorg/mistral-lora",
1218
+ "repo_type": "model",
1219
+ },
1220
+ ])
1221
+ @step
1222
+ def finetune_model(self):
1223
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1224
+ # path_to_model will be /my-directory
1225
+ ```
1207
1226
 
1208
1227
 
1209
1228
  Parameters
1210
1229
  ----------
1211
- vars : Dict[str, str], default {}
1212
- Dictionary of environment variables to set.
1213
- """
1214
- ...
1215
-
1216
- @typing.overload
1217
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1218
- ...
1219
-
1220
- @typing.overload
1221
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1222
- ...
1223
-
1224
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1225
- """
1226
- Specifies environment variables to be set prior to the execution of a step.
1230
+ temp_dir_root : str, optional
1231
+ The root directory that will hold the temporary directory where objects will be downloaded.
1232
+
1233
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1234
+ The list of repos (models/datasets) to load.
1227
1235
 
1236
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1228
1237
 
1229
- Parameters
1230
- ----------
1231
- vars : Dict[str, str], default {}
1232
- Dictionary of environment variables to set.
1238
+ - If repo (model/dataset) is not found in the datastore:
1239
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1240
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1241
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1242
+
1243
+ - If repo is found in the datastore:
1244
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
1233
1245
  """
1234
1246
  ...
1235
1247
 
@@ -1274,84 +1286,6 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1274
1286
  """
1275
1287
  ...
1276
1288
 
1277
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1278
- """
1279
- Specifies what flows belong to the same project.
1280
-
1281
- A project-specific namespace is created for all flows that
1282
- use the same `@project(name)`.
1283
-
1284
-
1285
- Parameters
1286
- ----------
1287
- name : str
1288
- Project name. Make sure that the name is unique amongst all
1289
- projects that use the same production scheduler. The name may
1290
- contain only lowercase alphanumeric characters and underscores.
1291
-
1292
- branch : Optional[str], default None
1293
- The branch to use. If not specified, the branch is set to
1294
- `user.<username>` unless `production` is set to `True`. This can
1295
- also be set on the command line using `--branch` as a top-level option.
1296
- It is an error to specify `branch` in the decorator and on the command line.
1297
-
1298
- production : bool, default False
1299
- Whether or not the branch is the production branch. This can also be set on the
1300
- command line using `--production` as a top-level option. It is an error to specify
1301
- `production` in the decorator and on the command line.
1302
- The project branch name will be:
1303
- - if `branch` is specified:
1304
- - if `production` is True: `prod.<branch>`
1305
- - if `production` is False: `test.<branch>`
1306
- - if `branch` is not specified:
1307
- - if `production` is True: `prod`
1308
- - if `production` is False: `user.<username>`
1309
- """
1310
- ...
1311
-
1312
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1313
- """
1314
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1315
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1316
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1317
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1318
- starts only after all sensors finish.
1319
-
1320
-
1321
- Parameters
1322
- ----------
1323
- timeout : int
1324
- Time, in seconds before the task times out and fails. (Default: 3600)
1325
- poke_interval : int
1326
- Time in seconds that the job should wait in between each try. (Default: 60)
1327
- mode : str
1328
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1329
- exponential_backoff : bool
1330
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1331
- pool : str
1332
- the slot pool this task should run in,
1333
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1334
- soft_fail : bool
1335
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1336
- name : str
1337
- Name of the sensor on Airflow
1338
- description : str
1339
- Description of sensor in the Airflow UI
1340
- bucket_key : Union[str, List[str]]
1341
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1342
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1343
- bucket_name : str
1344
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1345
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1346
- wildcard_match : bool
1347
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1348
- aws_conn_id : str
1349
- a reference to the s3 connection on Airflow. (Default: None)
1350
- verify : bool
1351
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1352
- """
1353
- ...
1354
-
1355
1289
  @typing.overload
1356
1290
  def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1357
1291
  """
@@ -1445,49 +1379,6 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1445
1379
  """
1446
1380
  ...
1447
1381
 
1448
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1449
- """
1450
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1451
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1452
-
1453
-
1454
- Parameters
1455
- ----------
1456
- timeout : int
1457
- Time, in seconds before the task times out and fails. (Default: 3600)
1458
- poke_interval : int
1459
- Time in seconds that the job should wait in between each try. (Default: 60)
1460
- mode : str
1461
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1462
- exponential_backoff : bool
1463
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1464
- pool : str
1465
- the slot pool this task should run in,
1466
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1467
- soft_fail : bool
1468
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1469
- name : str
1470
- Name of the sensor on Airflow
1471
- description : str
1472
- Description of sensor in the Airflow UI
1473
- external_dag_id : str
1474
- The dag_id that contains the task you want to wait for.
1475
- external_task_ids : List[str]
1476
- The list of task_ids that you want to wait for.
1477
- If None (default value) the sensor waits for the DAG. (Default: None)
1478
- allowed_states : List[str]
1479
- Iterable of allowed states, (Default: ['success'])
1480
- failed_states : List[str]
1481
- Iterable of failed or dis-allowed states. (Default: None)
1482
- execution_delta : datetime.timedelta
1483
- time difference with the previous execution to look at,
1484
- the default is the same logical date as the current task or DAG. (Default: None)
1485
- check_existence: bool
1486
- Set to True to check if the external task exists or check if
1487
- the DAG to wait for exists. (Default: True)
1488
- """
1489
- ...
1490
-
1491
1382
  @typing.overload
1492
1383
  def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1493
1384
  """
@@ -1640,6 +1531,49 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1640
1531
  """
1641
1532
  ...
1642
1533
 
1534
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1535
+ """
1536
+ The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1537
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1538
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1539
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1540
+ starts only after all sensors finish.
1541
+
1542
+
1543
+ Parameters
1544
+ ----------
1545
+ timeout : int
1546
+ Time, in seconds before the task times out and fails. (Default: 3600)
1547
+ poke_interval : int
1548
+ Time in seconds that the job should wait in between each try. (Default: 60)
1549
+ mode : str
1550
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1551
+ exponential_backoff : bool
1552
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1553
+ pool : str
1554
+ the slot pool this task should run in,
1555
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1556
+ soft_fail : bool
1557
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1558
+ name : str
1559
+ Name of the sensor on Airflow
1560
+ description : str
1561
+ Description of sensor in the Airflow UI
1562
+ bucket_key : Union[str, List[str]]
1563
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1564
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1565
+ bucket_name : str
1566
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1567
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1568
+ wildcard_match : bool
1569
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1570
+ aws_conn_id : str
1571
+ a reference to the s3 connection on Airflow. (Default: None)
1572
+ verify : bool
1573
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1574
+ """
1575
+ ...
1576
+
1643
1577
  @typing.overload
1644
1578
  def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1645
1579
  """
@@ -1691,6 +1625,49 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1691
1625
  """
1692
1626
  ...
1693
1627
 
1628
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1629
+ """
1630
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1631
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1632
+
1633
+
1634
+ Parameters
1635
+ ----------
1636
+ timeout : int
1637
+ Time, in seconds before the task times out and fails. (Default: 3600)
1638
+ poke_interval : int
1639
+ Time in seconds that the job should wait in between each try. (Default: 60)
1640
+ mode : str
1641
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1642
+ exponential_backoff : bool
1643
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1644
+ pool : str
1645
+ the slot pool this task should run in,
1646
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1647
+ soft_fail : bool
1648
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1649
+ name : str
1650
+ Name of the sensor on Airflow
1651
+ description : str
1652
+ Description of sensor in the Airflow UI
1653
+ external_dag_id : str
1654
+ The dag_id that contains the task you want to wait for.
1655
+ external_task_ids : List[str]
1656
+ The list of task_ids that you want to wait for.
1657
+ If None (default value) the sensor waits for the DAG. (Default: None)
1658
+ allowed_states : List[str]
1659
+ Iterable of allowed states, (Default: ['success'])
1660
+ failed_states : List[str]
1661
+ Iterable of failed or dis-allowed states. (Default: None)
1662
+ execution_delta : datetime.timedelta
1663
+ time difference with the previous execution to look at,
1664
+ the default is the same logical date as the current task or DAG. (Default: None)
1665
+ check_existence: bool
1666
+ Set to True to check if the external task exists or check if
1667
+ the DAG to wait for exists. (Default: True)
1668
+ """
1669
+ ...
1670
+
1694
1671
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1695
1672
  """
1696
1673
  Allows setting external datastores to save data for the
@@ -1805,5 +1782,40 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1805
1782
  """
1806
1783
  ...
1807
1784
 
1785
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1786
+ """
1787
+ Specifies what flows belong to the same project.
1788
+
1789
+ A project-specific namespace is created for all flows that
1790
+ use the same `@project(name)`.
1791
+
1792
+
1793
+ Parameters
1794
+ ----------
1795
+ name : str
1796
+ Project name. Make sure that the name is unique amongst all
1797
+ projects that use the same production scheduler. The name may
1798
+ contain only lowercase alphanumeric characters and underscores.
1799
+
1800
+ branch : Optional[str], default None
1801
+ The branch to use. If not specified, the branch is set to
1802
+ `user.<username>` unless `production` is set to `True`. This can
1803
+ also be set on the command line using `--branch` as a top-level option.
1804
+ It is an error to specify `branch` in the decorator and on the command line.
1805
+
1806
+ production : bool, default False
1807
+ Whether or not the branch is the production branch. This can also be set on the
1808
+ command line using `--production` as a top-level option. It is an error to specify
1809
+ `production` in the decorator and on the command line.
1810
+ The project branch name will be:
1811
+ - if `branch` is specified:
1812
+ - if `production` is True: `prod.<branch>`
1813
+ - if `production` is False: `test.<branch>`
1814
+ - if `branch` is not specified:
1815
+ - if `production` is True: `prod`
1816
+ - if `production` is False: `user.<username>`
1817
+ """
1818
+ ...
1819
+
1808
1820
  pkg_name: str
1809
1821