ob-metaflow-stubs 6.0.3.148__py2.py3-none-any.whl → 6.0.3.151__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 (198) hide show
  1. metaflow-stubs/__init__.pyi +668 -668
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +6 -6
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +3 -3
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +5 -5
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +3 -3
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +141 -141
  21. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  22. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +3 -3
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +5 -5
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +5 -5
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +4 -4
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +5 -5
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +3 -3
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +5 -5
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  61. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  62. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  63. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  74. metaflow-stubs/multicore_utils.pyi +2 -2
  75. metaflow-stubs/parameters.pyi +5 -5
  76. metaflow-stubs/plugins/__init__.pyi +12 -12
  77. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  78. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  79. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  80. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  81. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  82. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  83. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  84. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  85. metaflow-stubs/plugins/argo/argo_client.pyi +3 -3
  86. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  87. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  88. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +4 -4
  89. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  90. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +6 -4
  91. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  92. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  93. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  94. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  96. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  97. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  98. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  100. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  101. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  102. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  103. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  104. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  105. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  107. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  108. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  109. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  110. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  111. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  112. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  113. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  114. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  115. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  116. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  117. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  118. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  119. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  121. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  123. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  125. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  126. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  127. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  128. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  129. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  130. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  132. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  133. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  134. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  135. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  136. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  137. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  138. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  139. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  140. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  141. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  142. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  143. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  144. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  145. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  146. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  147. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  148. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  149. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  150. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  151. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  152. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  153. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  154. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  155. metaflow-stubs/plugins/ollama/__init__.pyi +3 -3
  156. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  157. metaflow-stubs/plugins/perimeters.pyi +2 -2
  158. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  159. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  160. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  162. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  164. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  165. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  166. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  169. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  170. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  171. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  173. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  174. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  175. metaflow-stubs/profilers/__init__.pyi +2 -2
  176. metaflow-stubs/pylint_wrapper.pyi +2 -2
  177. metaflow-stubs/runner/__init__.pyi +2 -2
  178. metaflow-stubs/runner/deployer.pyi +8 -6
  179. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  180. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  181. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  182. metaflow-stubs/runner/nbrun.pyi +2 -2
  183. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  184. metaflow-stubs/runner/utils.pyi +3 -3
  185. metaflow-stubs/system/__init__.pyi +2 -2
  186. metaflow-stubs/system/system_logger.pyi +2 -2
  187. metaflow-stubs/system/system_monitor.pyi +2 -2
  188. metaflow-stubs/tagging_util.pyi +2 -2
  189. metaflow-stubs/tuple_util.pyi +2 -2
  190. metaflow-stubs/user_configs/__init__.pyi +2 -2
  191. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  192. metaflow-stubs/user_configs/config_options.pyi +4 -4
  193. metaflow-stubs/user_configs/config_parameters.pyi +7 -7
  194. {ob_metaflow_stubs-6.0.3.148.dist-info → ob_metaflow_stubs-6.0.3.151.dist-info}/METADATA +1 -1
  195. ob_metaflow_stubs-6.0.3.151.dist-info/RECORD +198 -0
  196. ob_metaflow_stubs-6.0.3.148.dist-info/RECORD +0 -198
  197. {ob_metaflow_stubs-6.0.3.148.dist-info → ob_metaflow_stubs-6.0.3.151.dist-info}/WHEEL +0 -0
  198. {ob_metaflow_stubs-6.0.3.148.dist-info → ob_metaflow_stubs-6.0.3.151.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.3.1+obcheckpoint(0.1.9);ob(v1) #
4
- # Generated on 2025-03-03T22:55:42.139654 #
3
+ # MF version: 2.15.6.1+obcheckpoint(0.1.9);ob(v1) #
4
+ # Generated on 2025-03-17T20:51:51.819256 #
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,8 +35,8 @@ 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 cards as cards
39
38
  from . import tuple_util as tuple_util
39
+ from . import cards as cards
40
40
  from . import events as events
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
@@ -44,8 +44,8 @@ from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package imp
44
44
  from . import includefile as includefile
45
45
  from .includefile import IncludeFile as IncludeFile
46
46
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
47
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
47
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
49
  from . import client as client
50
50
  from .client.core import namespace as namespace
51
51
  from .client.core import get_namespace as get_namespace
@@ -149,39 +149,6 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
149
149
  """
150
150
  ...
151
151
 
152
- def nim(*, models: "list[NIM]", backend: str, queue_timeout: int) -> 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]]]:
153
- """
154
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
155
-
156
- User code call
157
- -----------
158
- @nim(
159
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
160
- backend='managed'
161
- )
162
-
163
- Valid backend options
164
- ---------------------
165
- - 'managed': Outerbounds selects a compute provider based on the model.
166
-
167
- Valid model options
168
- ----------------
169
- - 'meta/llama3-8b-instruct': 8B parameter model
170
- - 'meta/llama3-70b-instruct': 70B parameter model
171
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
172
-
173
-
174
- Parameters
175
- ----------
176
- models: list[NIM]
177
- List of NIM containers running models in sidecars.
178
- backend: str
179
- Compute provider to run the NIM container.
180
- queue_timeout : int
181
- Time to keep the job in NVCF's queue.
182
- """
183
- ...
184
-
185
152
  def ollama(*, models: "list[Ollama]", backend: 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]]]:
186
153
  """
187
154
  This decorator is used to run Ollama APIs as Metaflow task sidecars.
@@ -216,97 +183,35 @@ def ollama(*, models: "list[Ollama]", backend: str) -> typing.Callable[[typing.U
216
183
  ...
217
184
 
218
185
  @typing.overload
219
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
220
- """
221
- Decorator prototype for all step decorators. This function gets specialized
222
- and imported for all decorators types by _import_plugin_decorators().
223
- """
224
- ...
225
-
226
- @typing.overload
227
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
228
- ...
229
-
230
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
231
- """
232
- Decorator prototype for all step decorators. This function gets specialized
233
- and imported for all decorators types by _import_plugin_decorators().
234
- """
235
- ...
236
-
237
- 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]]]:
238
- """
239
- Decorator that helps cache, version and store models/datasets from huggingface hub.
240
-
241
-
242
- Parameters
243
- ----------
244
- temp_dir_root : str, optional
245
- The root directory that will hold the temporary directory where objects will be downloaded.
246
-
247
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
248
- The list of repos (models/datasets) to load.
249
-
250
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
251
-
252
- - If repo (model/dataset) is not found in the datastore:
253
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
254
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
255
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
256
-
257
- - If repo is found in the datastore:
258
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
259
- """
260
- ...
261
-
262
- @typing.overload
263
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> 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]]]:
186
+ 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]]]:
264
187
  """
265
- Specifies that the step will success under all circumstances.
266
-
267
- The decorator will create an optional artifact, specified by `var`, which
268
- contains the exception raised. You can use it to detect the presence
269
- of errors, indicating that all happy-path artifacts produced by the step
270
- are missing.
188
+ Specifies environment variables to be set prior to the execution of a step.
271
189
 
272
190
 
273
191
  Parameters
274
192
  ----------
275
- var : str, optional, default None
276
- Name of the artifact in which to store the caught exception.
277
- If not specified, the exception is not stored.
278
- print_exception : bool, default True
279
- Determines whether or not the exception is printed to
280
- stdout when caught.
193
+ vars : Dict[str, str], default {}
194
+ Dictionary of environment variables to set.
281
195
  """
282
196
  ...
283
197
 
284
198
  @typing.overload
285
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
199
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
286
200
  ...
287
201
 
288
202
  @typing.overload
289
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
203
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
290
204
  ...
291
205
 
292
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
206
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
293
207
  """
294
- Specifies that the step will success under all circumstances.
295
-
296
- The decorator will create an optional artifact, specified by `var`, which
297
- contains the exception raised. You can use it to detect the presence
298
- of errors, indicating that all happy-path artifacts produced by the step
299
- are missing.
208
+ Specifies environment variables to be set prior to the execution of a step.
300
209
 
301
210
 
302
211
  Parameters
303
212
  ----------
304
- var : str, optional, default None
305
- Name of the artifact in which to store the caught exception.
306
- If not specified, the exception is not stored.
307
- print_exception : bool, default True
308
- Determines whether or not the exception is printed to
309
- stdout when caught.
213
+ vars : Dict[str, str], default {}
214
+ Dictionary of environment variables to set.
310
215
  """
311
216
  ...
312
217
 
@@ -365,272 +270,95 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
365
270
  """
366
271
  ...
367
272
 
368
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', 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') -> 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]]]:
369
- """
370
- Specifies that this step should execute on Kubernetes.
371
-
372
-
373
- Parameters
374
- ----------
375
- cpu : int, default 1
376
- Number of CPUs required for this step. If `@resources` is
377
- also present, the maximum value from all decorators is used.
378
- memory : int, default 4096
379
- Memory size (in MB) required for this step. If
380
- `@resources` is also present, the maximum value from all decorators is
381
- used.
382
- disk : int, default 10240
383
- Disk size (in MB) required for this step. If
384
- `@resources` is also present, the maximum value from all decorators is
385
- used.
386
- image : str, optional, default None
387
- Docker image to use when launching on Kubernetes. If not specified, and
388
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
389
- not, a default Docker image mapping to the current version of Python is used.
390
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
391
- If given, the imagePullPolicy to be applied to the Docker image of the step.
392
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
393
- Kubernetes service account to use when launching pod in Kubernetes.
394
- secrets : List[str], optional, default None
395
- Kubernetes secrets to use when launching pod in Kubernetes. These
396
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
397
- in Metaflow configuration.
398
- node_selector: Union[Dict[str,str], str], optional, default None
399
- Kubernetes node selector(s) to apply to the pod running the task.
400
- Can be passed in as a comma separated string of values e.g.
401
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
402
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
403
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
404
- Kubernetes namespace to use when launching pod in Kubernetes.
405
- gpu : int, optional, default None
406
- Number of GPUs required for this step. A value of zero implies that
407
- the scheduled node should not have GPUs.
408
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
409
- The vendor of the GPUs to be used for this step.
410
- tolerations : List[str], default []
411
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
412
- Kubernetes tolerations to use when launching pod in Kubernetes.
413
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
414
- Kubernetes labels to use when launching pod in Kubernetes.
415
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
416
- Kubernetes annotations to use when launching pod in Kubernetes.
417
- use_tmpfs : bool, default False
418
- This enables an explicit tmpfs mount for this step.
419
- tmpfs_tempdir : bool, default True
420
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
421
- tmpfs_size : int, optional, default: None
422
- The value for the size (in MiB) of the tmpfs mount for this step.
423
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
424
- memory allocated for this step.
425
- tmpfs_path : str, optional, default /metaflow_temp
426
- Path to tmpfs mount for this step.
427
- persistent_volume_claims : Dict[str, str], optional, default None
428
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
429
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
430
- shared_memory: int, optional
431
- Shared memory size (in MiB) required for this step
432
- port: int, optional
433
- Port number to specify in the Kubernetes job object
434
- compute_pool : str, optional, default None
435
- Compute pool to be used for for this step.
436
- If not specified, any accessible compute pool within the perimeter is used.
437
- hostname_resolution_timeout: int, default 10 * 60
438
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
439
- Only applicable when @parallel is used.
440
- qos: str, default: Burstable
441
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
442
- """
443
- ...
444
-
445
273
  @typing.overload
446
- 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]]]:
274
+ 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]]]:
447
275
  """
448
- Specifies the Conda environment for the step.
449
-
450
- Information in this decorator will augment any
451
- attributes set in the `@conda_base` flow-level decorator. Hence,
452
- you can use `@conda_base` to set packages required by all
453
- steps and use `@conda` to specify step-specific overrides.
276
+ Specifies secrets to be retrieved and injected as environment variables prior to
277
+ the execution of a step.
454
278
 
455
279
 
456
280
  Parameters
457
281
  ----------
458
- packages : Dict[str, str], default {}
459
- Packages to use for this step. The key is the name of the package
460
- and the value is the version to use.
461
- libraries : Dict[str, str], default {}
462
- Supported for backward compatibility. When used with packages, packages will take precedence.
463
- python : str, optional, default None
464
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
465
- that the version used will correspond to the version of the Python interpreter used to start the run.
466
- disabled : bool, default False
467
- If set to True, disables @conda.
282
+ sources : List[Union[str, Dict[str, Any]]], default: []
283
+ List of secret specs, defining how the secrets are to be retrieved
468
284
  """
469
285
  ...
470
286
 
471
287
  @typing.overload
472
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
288
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
473
289
  ...
474
290
 
475
291
  @typing.overload
476
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
292
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
477
293
  ...
478
294
 
479
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
295
+ 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]]] = []):
480
296
  """
481
- Specifies the Conda environment for the step.
482
-
483
- Information in this decorator will augment any
484
- attributes set in the `@conda_base` flow-level decorator. Hence,
485
- you can use `@conda_base` to set packages required by all
486
- steps and use `@conda` to specify step-specific overrides.
297
+ Specifies secrets to be retrieved and injected as environment variables prior to
298
+ the execution of a step.
487
299
 
488
300
 
489
301
  Parameters
490
302
  ----------
491
- packages : Dict[str, str], default {}
492
- Packages to use for this step. The key is the name of the package
493
- and the value is the version to use.
494
- libraries : Dict[str, str], default {}
495
- Supported for backward compatibility. When used with packages, packages will take precedence.
496
- python : str, optional, default None
497
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
498
- that the version used will correspond to the version of the Python interpreter used to start the run.
499
- disabled : bool, default False
500
- If set to True, disables @conda.
303
+ sources : List[Union[str, Dict[str, Any]]], default: []
304
+ List of secret specs, defining how the secrets are to be retrieved
501
305
  """
502
306
  ...
503
307
 
504
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> 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
+ @typing.overload
309
+ 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]]]:
505
310
  """
506
- Specifies that this step should execute on DGX cloud.
311
+ Enables checkpointing for a step.
312
+
507
313
 
508
314
 
509
315
  Parameters
510
316
  ----------
511
- gpu : int
512
- Number of GPUs to use.
513
- gpu_type : str
514
- Type of Nvidia GPU to use.
515
- queue_timeout : int
516
- Time to keep the job in NVCF's queue.
517
- """
518
- ...
519
-
520
- @typing.overload
521
- 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]]]:
522
- """
523
- Creates a human-readable report, a Metaflow Card, after this step completes.
524
-
525
- Note that you may add multiple `@card` decorators in a step with different parameters.
526
-
527
-
528
- Parameters
529
- ----------
530
- type : str, default 'default'
531
- Card type.
532
- id : str, optional, default None
533
- If multiple cards are present, use this id to identify this card.
534
- options : Dict[str, Any], default {}
535
- Options passed to the card. The contents depend on the card type.
536
- timeout : int, default 45
537
- Interrupt reporting if it takes more than this many seconds.
538
- """
539
- ...
540
-
541
- @typing.overload
542
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
543
- ...
544
-
545
- @typing.overload
546
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
- ...
548
-
549
- 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):
550
- """
551
- Creates a human-readable report, a Metaflow Card, after this step completes.
552
-
553
- Note that you may add multiple `@card` decorators in a step with different parameters.
554
-
317
+ load_policy : str, default: "fresh"
318
+ The policy for loading the checkpoint. The following policies are supported:
319
+ - "eager": Loads the the latest available checkpoint within the namespace.
320
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
321
+ will be loaded at the start of the task.
322
+ - "none": Do not load any checkpoint
323
+ - "fresh": Loads the lastest checkpoint created within the running Task.
324
+ This mode helps loading checkpoints across various retry attempts of the same task.
325
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
326
+ created within the task will be loaded when the task is retries execution on failure.
555
327
 
556
- Parameters
557
- ----------
558
- type : str, default 'default'
559
- Card type.
560
- id : str, optional, default None
561
- If multiple cards are present, use this id to identify this card.
562
- options : Dict[str, Any], default {}
563
- Options passed to the card. The contents depend on the card type.
564
- timeout : int, default 45
565
- Interrupt reporting if it takes more than this many seconds.
328
+ temp_dir_root : str, default: None
329
+ The root directory under which `current.checkpoint.directory` will be created.
566
330
  """
567
331
  ...
568
332
 
569
333
  @typing.overload
570
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
571
- """
572
- Internal decorator to support Fast bakery
573
- """
334
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
574
335
  ...
575
336
 
576
337
  @typing.overload
577
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
578
- ...
579
-
580
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
581
- """
582
- Internal decorator to support Fast bakery
583
- """
338
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
584
339
  ...
585
340
 
586
- @typing.overload
587
- 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]]]:
341
+ 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):
588
342
  """
589
- Specifies the PyPI packages for the step.
343
+ Enables checkpointing for a step.
590
344
 
591
- Information in this decorator will augment any
592
- attributes set in the `@pyi_base` flow-level decorator. Hence,
593
- you can use `@pypi_base` to set packages required by all
594
- steps and use `@pypi` to specify step-specific overrides.
595
345
 
596
346
 
597
347
  Parameters
598
348
  ----------
599
- packages : Dict[str, str], default: {}
600
- Packages to use for this step. The key is the name of the package
601
- and the value is the version to use.
602
- python : str, optional, default: None
603
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
604
- that the version used will correspond to the version of the Python interpreter used to start the run.
605
- """
606
- ...
607
-
608
- @typing.overload
609
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
610
- ...
611
-
612
- @typing.overload
613
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
614
- ...
615
-
616
- 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):
617
- """
618
- Specifies the PyPI packages for the step.
619
-
620
- Information in this decorator will augment any
621
- attributes set in the `@pyi_base` flow-level decorator. Hence,
622
- you can use `@pypi_base` to set packages required by all
623
- steps and use `@pypi` to specify step-specific overrides.
624
-
349
+ load_policy : str, default: "fresh"
350
+ The policy for loading the checkpoint. The following policies are supported:
351
+ - "eager": Loads the the latest available checkpoint within the namespace.
352
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
353
+ will be loaded at the start of the task.
354
+ - "none": Do not load any checkpoint
355
+ - "fresh": Loads the lastest checkpoint created within the running Task.
356
+ This mode helps loading checkpoints across various retry attempts of the same task.
357
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
358
+ created within the task will be loaded when the task is retries execution on failure.
625
359
 
626
- Parameters
627
- ----------
628
- packages : Dict[str, str], default: {}
629
- Packages to use for this step. The key is the name of the package
630
- and the value is the version to use.
631
- python : str, optional, default: None
632
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
633
- that the version used will correspond to the version of the Python interpreter used to start the run.
360
+ temp_dir_root : str, default: None
361
+ The root directory under which `current.checkpoint.directory` will be created.
634
362
  """
635
363
  ...
636
364
 
@@ -691,187 +419,234 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
691
419
  """
692
420
  ...
693
421
 
694
- @typing.overload
695
- 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]]]:
422
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', 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') -> 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]]]:
696
423
  """
697
- Specifies secrets to be retrieved and injected as environment variables prior to
698
- the execution of a step.
424
+ Specifies that this step should execute on Kubernetes.
699
425
 
700
426
 
701
427
  Parameters
702
428
  ----------
703
- sources : List[Union[str, Dict[str, Any]]], default: []
704
- List of secret specs, defining how the secrets are to be retrieved
429
+ cpu : int, default 1
430
+ Number of CPUs required for this step. If `@resources` is
431
+ also present, the maximum value from all decorators is used.
432
+ memory : int, default 4096
433
+ Memory size (in MB) required for this step. If
434
+ `@resources` is also present, the maximum value from all decorators is
435
+ used.
436
+ disk : int, default 10240
437
+ Disk size (in MB) required for this step. If
438
+ `@resources` is also present, the maximum value from all decorators is
439
+ used.
440
+ image : str, optional, default None
441
+ Docker image to use when launching on Kubernetes. If not specified, and
442
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
443
+ not, a default Docker image mapping to the current version of Python is used.
444
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
445
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
446
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
447
+ Kubernetes service account to use when launching pod in Kubernetes.
448
+ secrets : List[str], optional, default None
449
+ Kubernetes secrets to use when launching pod in Kubernetes. These
450
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
451
+ in Metaflow configuration.
452
+ node_selector: Union[Dict[str,str], str], optional, default None
453
+ Kubernetes node selector(s) to apply to the pod running the task.
454
+ Can be passed in as a comma separated string of values e.g.
455
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
456
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
457
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
458
+ Kubernetes namespace to use when launching pod in Kubernetes.
459
+ gpu : int, optional, default None
460
+ Number of GPUs required for this step. A value of zero implies that
461
+ the scheduled node should not have GPUs.
462
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
463
+ The vendor of the GPUs to be used for this step.
464
+ tolerations : List[str], default []
465
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
466
+ Kubernetes tolerations to use when launching pod in Kubernetes.
467
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
468
+ Kubernetes labels to use when launching pod in Kubernetes.
469
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
470
+ Kubernetes annotations to use when launching pod in Kubernetes.
471
+ use_tmpfs : bool, default False
472
+ This enables an explicit tmpfs mount for this step.
473
+ tmpfs_tempdir : bool, default True
474
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
475
+ tmpfs_size : int, optional, default: None
476
+ The value for the size (in MiB) of the tmpfs mount for this step.
477
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
478
+ memory allocated for this step.
479
+ tmpfs_path : str, optional, default /metaflow_temp
480
+ Path to tmpfs mount for this step.
481
+ persistent_volume_claims : Dict[str, str], optional, default None
482
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
483
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
484
+ shared_memory: int, optional
485
+ Shared memory size (in MiB) required for this step
486
+ port: int, optional
487
+ Port number to specify in the Kubernetes job object
488
+ compute_pool : str, optional, default None
489
+ Compute pool to be used for for this step.
490
+ If not specified, any accessible compute pool within the perimeter is used.
491
+ hostname_resolution_timeout: int, default 10 * 60
492
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
493
+ Only applicable when @parallel is used.
494
+ qos: str, default: Burstable
495
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
705
496
  """
706
497
  ...
707
498
 
708
499
  @typing.overload
709
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
500
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
501
+ """
502
+ Decorator prototype for all step decorators. This function gets specialized
503
+ and imported for all decorators types by _import_plugin_decorators().
504
+ """
710
505
  ...
711
506
 
712
507
  @typing.overload
713
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
508
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
714
509
  ...
715
510
 
716
- 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]]] = []):
511
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
717
512
  """
718
- Specifies secrets to be retrieved and injected as environment variables prior to
719
- the execution of a step.
720
-
721
-
722
- Parameters
723
- ----------
724
- sources : List[Union[str, Dict[str, Any]]], default: []
725
- List of secret specs, defining how the secrets are to be retrieved
513
+ Decorator prototype for all step decorators. This function gets specialized
514
+ and imported for all decorators types by _import_plugin_decorators().
726
515
  """
727
516
  ...
728
517
 
729
518
  @typing.overload
730
- 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]]]:
519
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
731
520
  """
732
- Specifies a timeout for your step.
733
-
734
- This decorator is useful if this step may hang indefinitely.
735
-
736
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
737
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
738
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
739
-
740
- Note that all the values specified in parameters are added together so if you specify
741
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
742
-
743
-
744
- Parameters
745
- ----------
746
- seconds : int, default 0
747
- Number of seconds to wait prior to timing out.
748
- minutes : int, default 0
749
- Number of minutes to wait prior to timing out.
750
- hours : int, default 0
751
- Number of hours to wait prior to timing out.
521
+ Internal decorator to support Fast bakery
752
522
  """
753
523
  ...
754
524
 
755
525
  @typing.overload
756
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
526
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
757
527
  ...
758
528
 
759
- @typing.overload
760
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
761
- ...
762
-
763
- 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):
764
- """
765
- Specifies a timeout for your step.
766
-
767
- This decorator is useful if this step may hang indefinitely.
768
-
769
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
770
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
771
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
772
-
773
- Note that all the values specified in parameters are added together so if you specify
774
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
775
-
776
-
777
- Parameters
778
- ----------
779
- seconds : int, default 0
780
- Number of seconds to wait prior to timing out.
781
- minutes : int, default 0
782
- Number of minutes to wait prior to timing out.
783
- hours : int, default 0
784
- Number of hours to wait prior to timing out.
785
- """
529
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
530
+ """
531
+ Internal decorator to support Fast bakery
532
+ """
786
533
  ...
787
534
 
788
535
  @typing.overload
789
- 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]]]:
536
+ 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]]]:
790
537
  """
791
- Enables checkpointing for a step.
538
+ Specifies the Conda environment for the step.
792
539
 
540
+ Information in this decorator will augment any
541
+ attributes set in the `@conda_base` flow-level decorator. Hence,
542
+ you can use `@conda_base` to set packages required by all
543
+ steps and use `@conda` to specify step-specific overrides.
793
544
 
794
545
 
795
546
  Parameters
796
547
  ----------
797
- load_policy : str, default: "fresh"
798
- The policy for loading the checkpoint. The following policies are supported:
799
- - "eager": Loads the the latest available checkpoint within the namespace.
800
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
801
- will be loaded at the start of the task.
802
- - "none": Do not load any checkpoint
803
- - "fresh": Loads the lastest checkpoint created within the running Task.
804
- This mode helps loading checkpoints across various retry attempts of the same task.
805
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
806
- created within the task will be loaded when the task is retries execution on failure.
807
-
808
- temp_dir_root : str, default: None
809
- The root directory under which `current.checkpoint.directory` will be created.
548
+ packages : Dict[str, str], default {}
549
+ Packages to use for this step. The key is the name of the package
550
+ and the value is the version to use.
551
+ libraries : Dict[str, str], default {}
552
+ Supported for backward compatibility. When used with packages, packages will take precedence.
553
+ python : str, optional, default None
554
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
555
+ that the version used will correspond to the version of the Python interpreter used to start the run.
556
+ disabled : bool, default False
557
+ If set to True, disables @conda.
810
558
  """
811
559
  ...
812
560
 
813
561
  @typing.overload
814
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
562
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
815
563
  ...
816
564
 
817
565
  @typing.overload
818
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
566
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
819
567
  ...
820
568
 
821
- 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):
569
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
822
570
  """
823
- Enables checkpointing for a step.
571
+ Specifies the Conda environment for the step.
824
572
 
573
+ Information in this decorator will augment any
574
+ attributes set in the `@conda_base` flow-level decorator. Hence,
575
+ you can use `@conda_base` to set packages required by all
576
+ steps and use `@conda` to specify step-specific overrides.
825
577
 
826
578
 
827
579
  Parameters
828
580
  ----------
829
- load_policy : str, default: "fresh"
830
- The policy for loading the checkpoint. The following policies are supported:
831
- - "eager": Loads the the latest available checkpoint within the namespace.
832
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
833
- will be loaded at the start of the task.
834
- - "none": Do not load any checkpoint
835
- - "fresh": Loads the lastest checkpoint created within the running Task.
836
- This mode helps loading checkpoints across various retry attempts of the same task.
837
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
838
- created within the task will be loaded when the task is retries execution on failure.
839
-
840
- temp_dir_root : str, default: None
841
- The root directory under which `current.checkpoint.directory` will be created.
581
+ packages : Dict[str, str], default {}
582
+ Packages to use for this step. The key is the name of the package
583
+ and the value is the version to use.
584
+ libraries : Dict[str, str], default {}
585
+ Supported for backward compatibility. When used with packages, packages will take precedence.
586
+ python : str, optional, default None
587
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
588
+ that the version used will correspond to the version of the Python interpreter used to start the run.
589
+ disabled : bool, default False
590
+ If set to True, disables @conda.
842
591
  """
843
592
  ...
844
593
 
845
594
  @typing.overload
846
- 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]]]:
595
+ 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]]]:
847
596
  """
848
- Specifies environment variables to be set prior to the execution of a step.
597
+ Specifies a timeout for your step.
598
+
599
+ This decorator is useful if this step may hang indefinitely.
600
+
601
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
602
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
603
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
604
+
605
+ Note that all the values specified in parameters are added together so if you specify
606
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
849
607
 
850
608
 
851
609
  Parameters
852
610
  ----------
853
- vars : Dict[str, str], default {}
854
- Dictionary of environment variables to set.
611
+ seconds : int, default 0
612
+ Number of seconds to wait prior to timing out.
613
+ minutes : int, default 0
614
+ Number of minutes to wait prior to timing out.
615
+ hours : int, default 0
616
+ Number of hours to wait prior to timing out.
855
617
  """
856
618
  ...
857
619
 
858
620
  @typing.overload
859
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
621
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
860
622
  ...
861
623
 
862
624
  @typing.overload
863
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
625
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
864
626
  ...
865
627
 
866
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
628
+ 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):
867
629
  """
868
- Specifies environment variables to be set prior to the execution of a step.
630
+ Specifies a timeout for your step.
631
+
632
+ This decorator is useful if this step may hang indefinitely.
633
+
634
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
635
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
636
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
637
+
638
+ Note that all the values specified in parameters are added together so if you specify
639
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
869
640
 
870
641
 
871
642
  Parameters
872
643
  ----------
873
- vars : Dict[str, str], default {}
874
- Dictionary of environment variables to set.
644
+ seconds : int, default 0
645
+ Number of seconds to wait prior to timing out.
646
+ minutes : int, default 0
647
+ Number of minutes to wait prior to timing out.
648
+ hours : int, default 0
649
+ Number of hours to wait prior to timing out.
875
650
  """
876
651
  ...
877
652
 
@@ -954,190 +729,279 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
954
729
  """
955
730
  ...
956
731
 
957
- @typing.overload
958
- 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]]:
732
+ 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]]]:
959
733
  """
960
- Specifies the flow(s) that this flow depends on.
734
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
961
735
 
962
- ```
963
- @trigger_on_finish(flow='FooFlow')
964
- ```
965
- or
966
- ```
967
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
968
- ```
969
- This decorator respects the @project decorator and triggers the flow
970
- when upstream runs within the same namespace complete successfully
971
736
 
972
- Additionally, you can specify project aware upstream flow dependencies
973
- by specifying the fully qualified project_flow_name.
974
- ```
975
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
976
- ```
977
- or
978
- ```
979
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
980
- ```
737
+ Parameters
738
+ ----------
739
+ temp_dir_root : str, optional
740
+ The root directory that will hold the temporary directory where objects will be downloaded.
981
741
 
982
- You can also specify just the project or project branch (other values will be
983
- inferred from the current project or project branch):
984
- ```
985
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
986
- ```
742
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
743
+ The list of repos (models/datasets) to load.
987
744
 
988
- Note that `branch` is typically one of:
989
- - `prod`
990
- - `user.bob`
991
- - `test.my_experiment`
992
- - `prod.staging`
745
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
746
+
747
+ - If repo (model/dataset) is not found in the datastore:
748
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
749
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
750
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
751
+
752
+ - If repo is found in the datastore:
753
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
754
+ """
755
+ ...
756
+
757
+ @typing.overload
758
+ 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]]]:
759
+ """
760
+ Creates a human-readable report, a Metaflow Card, after this step completes.
761
+
762
+ Note that you may add multiple `@card` decorators in a step with different parameters.
993
763
 
994
764
 
995
765
  Parameters
996
766
  ----------
997
- flow : Union[str, Dict[str, str]], optional, default None
998
- Upstream flow dependency for this flow.
999
- flows : List[Union[str, Dict[str, str]]], default []
1000
- Upstream flow dependencies for this flow.
767
+ type : str, default 'default'
768
+ Card type.
769
+ id : str, optional, default None
770
+ If multiple cards are present, use this id to identify this card.
1001
771
  options : Dict[str, Any], default {}
1002
- Backend-specific configuration for tuning eventing behavior.
772
+ Options passed to the card. The contents depend on the card type.
773
+ timeout : int, default 45
774
+ Interrupt reporting if it takes more than this many seconds.
1003
775
  """
1004
776
  ...
1005
777
 
1006
778
  @typing.overload
1007
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
779
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1008
780
  ...
1009
781
 
1010
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, 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] = {}):
782
+ @typing.overload
783
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
784
+ ...
785
+
786
+ 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):
1011
787
  """
1012
- Specifies the flow(s) that this flow depends on.
1013
-
1014
- ```
1015
- @trigger_on_finish(flow='FooFlow')
1016
- ```
1017
- or
1018
- ```
1019
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1020
- ```
1021
- This decorator respects the @project decorator and triggers the flow
1022
- when upstream runs within the same namespace complete successfully
1023
-
1024
- Additionally, you can specify project aware upstream flow dependencies
1025
- by specifying the fully qualified project_flow_name.
1026
- ```
1027
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1028
- ```
1029
- or
1030
- ```
1031
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1032
- ```
1033
-
1034
- You can also specify just the project or project branch (other values will be
1035
- inferred from the current project or project branch):
1036
- ```
1037
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1038
- ```
788
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1039
789
 
1040
- Note that `branch` is typically one of:
1041
- - `prod`
1042
- - `user.bob`
1043
- - `test.my_experiment`
1044
- - `prod.staging`
790
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1045
791
 
1046
792
 
1047
793
  Parameters
1048
794
  ----------
1049
- flow : Union[str, Dict[str, str]], optional, default None
1050
- Upstream flow dependency for this flow.
1051
- flows : List[Union[str, Dict[str, str]]], default []
1052
- Upstream flow dependencies for this flow.
795
+ type : str, default 'default'
796
+ Card type.
797
+ id : str, optional, default None
798
+ If multiple cards are present, use this id to identify this card.
1053
799
  options : Dict[str, Any], default {}
1054
- Backend-specific configuration for tuning eventing behavior.
800
+ Options passed to the card. The contents depend on the card type.
801
+ timeout : int, default 45
802
+ Interrupt reporting if it takes more than this many seconds.
1055
803
  """
1056
804
  ...
1057
805
 
1058
- 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]]:
806
+ @typing.overload
807
+ 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]]]:
1059
808
  """
1060
- 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)
1061
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1062
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1063
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1064
- starts only after all sensors finish.
809
+ Specifies the PyPI packages for the step.
810
+
811
+ Information in this decorator will augment any
812
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
813
+ you can use `@pypi_base` to set packages required by all
814
+ steps and use `@pypi` to specify step-specific overrides.
1065
815
 
1066
816
 
1067
817
  Parameters
1068
818
  ----------
1069
- timeout : int
1070
- Time, in seconds before the task times out and fails. (Default: 3600)
1071
- poke_interval : int
1072
- Time in seconds that the job should wait in between each try. (Default: 60)
1073
- mode : str
1074
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1075
- exponential_backoff : bool
1076
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1077
- pool : str
1078
- the slot pool this task should run in,
1079
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1080
- soft_fail : bool
1081
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1082
- name : str
1083
- Name of the sensor on Airflow
1084
- description : str
1085
- Description of sensor in the Airflow UI
1086
- bucket_key : Union[str, List[str]]
1087
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1088
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1089
- bucket_name : str
1090
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1091
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1092
- wildcard_match : bool
1093
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1094
- aws_conn_id : str
1095
- a reference to the s3 connection on Airflow. (Default: None)
1096
- verify : bool
1097
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
819
+ packages : Dict[str, str], default: {}
820
+ Packages to use for this step. The key is the name of the package
821
+ and the value is the version to use.
822
+ python : str, optional, default: None
823
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
824
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1098
825
  """
1099
826
  ...
1100
827
 
1101
- 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]]:
828
+ @typing.overload
829
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
830
+ ...
831
+
832
+ @typing.overload
833
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
834
+ ...
835
+
836
+ 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):
1102
837
  """
1103
- 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.
1104
- 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.
838
+ Specifies the PyPI packages for the step.
839
+
840
+ Information in this decorator will augment any
841
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
842
+ you can use `@pypi_base` to set packages required by all
843
+ steps and use `@pypi` to specify step-specific overrides.
1105
844
 
1106
845
 
1107
846
  Parameters
1108
847
  ----------
1109
- timeout : int
1110
- Time, in seconds before the task times out and fails. (Default: 3600)
1111
- poke_interval : int
1112
- Time in seconds that the job should wait in between each try. (Default: 60)
1113
- mode : str
1114
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1115
- exponential_backoff : bool
1116
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1117
- pool : str
1118
- the slot pool this task should run in,
1119
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1120
- soft_fail : bool
1121
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1122
- name : str
1123
- Name of the sensor on Airflow
1124
- description : str
1125
- Description of sensor in the Airflow UI
1126
- external_dag_id : str
1127
- The dag_id that contains the task you want to wait for.
1128
- external_task_ids : List[str]
1129
- The list of task_ids that you want to wait for.
1130
- If None (default value) the sensor waits for the DAG. (Default: None)
1131
- allowed_states : List[str]
1132
- Iterable of allowed states, (Default: ['success'])
1133
- failed_states : List[str]
1134
- Iterable of failed or dis-allowed states. (Default: None)
1135
- execution_delta : datetime.timedelta
1136
- time difference with the previous execution to look at,
1137
- the default is the same logical date as the current task or DAG. (Default: None)
1138
- check_existence: bool
1139
- Set to True to check if the external task exists or check if
1140
- the DAG to wait for exists. (Default: True)
848
+ packages : Dict[str, str], default: {}
849
+ Packages to use for this step. The key is the name of the package
850
+ and the value is the version to use.
851
+ python : str, optional, default: None
852
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
853
+ that the version used will correspond to the version of the Python interpreter used to start the run.
854
+ """
855
+ ...
856
+
857
+ def nim(*, models: "list[NIM]", backend: str, queue_timeout: int) -> 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]]]:
858
+ """
859
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
860
+
861
+ User code call
862
+ -----------
863
+ @nim(
864
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
865
+ backend='managed'
866
+ )
867
+
868
+ Valid backend options
869
+ ---------------------
870
+ - 'managed': Outerbounds selects a compute provider based on the model.
871
+
872
+ Valid model options
873
+ ----------------
874
+ - 'meta/llama3-8b-instruct': 8B parameter model
875
+ - 'meta/llama3-70b-instruct': 70B parameter model
876
+ - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
877
+
878
+
879
+ Parameters
880
+ ----------
881
+ models: list[NIM]
882
+ List of NIM containers running models in sidecars.
883
+ backend: str
884
+ Compute provider to run the NIM container.
885
+ queue_timeout : int
886
+ Time to keep the job in NVCF's queue.
887
+ """
888
+ ...
889
+
890
+ @typing.overload
891
+ def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> 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]]]:
892
+ """
893
+ Specifies that the step will success under all circumstances.
894
+
895
+ The decorator will create an optional artifact, specified by `var`, which
896
+ contains the exception raised. You can use it to detect the presence
897
+ of errors, indicating that all happy-path artifacts produced by the step
898
+ are missing.
899
+
900
+
901
+ Parameters
902
+ ----------
903
+ var : str, optional, default None
904
+ Name of the artifact in which to store the caught exception.
905
+ If not specified, the exception is not stored.
906
+ print_exception : bool, default True
907
+ Determines whether or not the exception is printed to
908
+ stdout when caught.
909
+ """
910
+ ...
911
+
912
+ @typing.overload
913
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
914
+ ...
915
+
916
+ @typing.overload
917
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
918
+ ...
919
+
920
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
921
+ """
922
+ Specifies that the step will success under all circumstances.
923
+
924
+ The decorator will create an optional artifact, specified by `var`, which
925
+ contains the exception raised. You can use it to detect the presence
926
+ of errors, indicating that all happy-path artifacts produced by the step
927
+ are missing.
928
+
929
+
930
+ Parameters
931
+ ----------
932
+ var : str, optional, default None
933
+ Name of the artifact in which to store the caught exception.
934
+ If not specified, the exception is not stored.
935
+ print_exception : bool, default True
936
+ Determines whether or not the exception is printed to
937
+ stdout when caught.
938
+ """
939
+ ...
940
+
941
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> 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]]]:
942
+ """
943
+ Specifies that this step should execute on DGX cloud.
944
+
945
+
946
+ Parameters
947
+ ----------
948
+ gpu : int
949
+ Number of GPUs to use.
950
+ gpu_type : str
951
+ Type of Nvidia GPU to use.
952
+ queue_timeout : int
953
+ Time to keep the job in NVCF's queue.
954
+ """
955
+ ...
956
+
957
+ @typing.overload
958
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
959
+ """
960
+ Specifies the times when the flow should be run when running on a
961
+ production scheduler.
962
+
963
+
964
+ Parameters
965
+ ----------
966
+ hourly : bool, default False
967
+ Run the workflow hourly.
968
+ daily : bool, default True
969
+ Run the workflow daily.
970
+ weekly : bool, default False
971
+ Run the workflow weekly.
972
+ cron : str, optional, default None
973
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
974
+ specified by this expression.
975
+ timezone : str, optional, default None
976
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
977
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
978
+ """
979
+ ...
980
+
981
+ @typing.overload
982
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
983
+ ...
984
+
985
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
986
+ """
987
+ Specifies the times when the flow should be run when running on a
988
+ production scheduler.
989
+
990
+
991
+ Parameters
992
+ ----------
993
+ hourly : bool, default False
994
+ Run the workflow hourly.
995
+ daily : bool, default True
996
+ Run the workflow daily.
997
+ weekly : bool, default False
998
+ Run the workflow weekly.
999
+ cron : str, optional, default None
1000
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1001
+ specified by this expression.
1002
+ timezone : str, optional, default None
1003
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1004
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1141
1005
  """
1142
1006
  ...
1143
1007
 
@@ -1176,44 +1040,46 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1176
1040
  """
1177
1041
  ...
1178
1042
 
1179
- @typing.overload
1180
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1181
- """
1182
- Specifies the PyPI packages for all steps of the flow.
1183
-
1184
- Use `@pypi_base` to set common packages required by all
1185
- steps and use `@pypi` to specify step-specific overrides.
1186
-
1187
- Parameters
1188
- ----------
1189
- packages : Dict[str, str], default: {}
1190
- Packages to use for this flow. The key is the name of the package
1191
- and the value is the version to use.
1192
- python : str, optional, default: None
1193
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1194
- that the version used will correspond to the version of the Python interpreter used to start the run.
1195
- """
1196
- ...
1197
-
1198
- @typing.overload
1199
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1200
- ...
1201
-
1202
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1043
+ 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]]:
1203
1044
  """
1204
- Specifies the PyPI packages for all steps of the flow.
1045
+ 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.
1046
+ 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.
1205
1047
 
1206
- Use `@pypi_base` to set common packages required by all
1207
- steps and use `@pypi` to specify step-specific overrides.
1208
1048
 
1209
1049
  Parameters
1210
1050
  ----------
1211
- packages : Dict[str, str], default: {}
1212
- Packages to use for this flow. The key is the name of the package
1213
- and the value is the version to use.
1214
- python : str, optional, default: None
1215
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1216
- that the version used will correspond to the version of the Python interpreter used to start the run.
1051
+ timeout : int
1052
+ Time, in seconds before the task times out and fails. (Default: 3600)
1053
+ poke_interval : int
1054
+ Time in seconds that the job should wait in between each try. (Default: 60)
1055
+ mode : str
1056
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1057
+ exponential_backoff : bool
1058
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1059
+ pool : str
1060
+ the slot pool this task should run in,
1061
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1062
+ soft_fail : bool
1063
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1064
+ name : str
1065
+ Name of the sensor on Airflow
1066
+ description : str
1067
+ Description of sensor in the Airflow UI
1068
+ external_dag_id : str
1069
+ The dag_id that contains the task you want to wait for.
1070
+ external_task_ids : List[str]
1071
+ The list of task_ids that you want to wait for.
1072
+ If None (default value) the sensor waits for the DAG. (Default: None)
1073
+ allowed_states : List[str]
1074
+ Iterable of allowed states, (Default: ['success'])
1075
+ failed_states : List[str]
1076
+ Iterable of failed or dis-allowed states. (Default: None)
1077
+ execution_delta : datetime.timedelta
1078
+ time difference with the previous execution to look at,
1079
+ the default is the same logical date as the current task or DAG. (Default: None)
1080
+ check_existence: bool
1081
+ Set to True to check if the external task exists or check if
1082
+ the DAG to wait for exists. (Default: True)
1217
1083
  """
1218
1084
  ...
1219
1085
 
@@ -1310,57 +1176,6 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1310
1176
  """
1311
1177
  ...
1312
1178
 
1313
- @typing.overload
1314
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1315
- """
1316
- Specifies the times when the flow should be run when running on a
1317
- production scheduler.
1318
-
1319
-
1320
- Parameters
1321
- ----------
1322
- hourly : bool, default False
1323
- Run the workflow hourly.
1324
- daily : bool, default True
1325
- Run the workflow daily.
1326
- weekly : bool, default False
1327
- Run the workflow weekly.
1328
- cron : str, optional, default None
1329
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1330
- specified by this expression.
1331
- timezone : str, optional, default None
1332
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1333
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1334
- """
1335
- ...
1336
-
1337
- @typing.overload
1338
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1339
- ...
1340
-
1341
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1342
- """
1343
- Specifies the times when the flow should be run when running on a
1344
- production scheduler.
1345
-
1346
-
1347
- Parameters
1348
- ----------
1349
- hourly : bool, default False
1350
- Run the workflow hourly.
1351
- daily : bool, default True
1352
- Run the workflow daily.
1353
- weekly : bool, default False
1354
- Run the workflow weekly.
1355
- cron : str, optional, default None
1356
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1357
- specified by this expression.
1358
- timezone : str, optional, default None
1359
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1360
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1361
- """
1362
- ...
1363
-
1364
1179
  @typing.overload
1365
1180
  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]]:
1366
1181
  """
@@ -1412,5 +1227,190 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1412
1227
  """
1413
1228
  ...
1414
1229
 
1230
+ @typing.overload
1231
+ 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]]:
1232
+ """
1233
+ Specifies the flow(s) that this flow depends on.
1234
+
1235
+ ```
1236
+ @trigger_on_finish(flow='FooFlow')
1237
+ ```
1238
+ or
1239
+ ```
1240
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1241
+ ```
1242
+ This decorator respects the @project decorator and triggers the flow
1243
+ when upstream runs within the same namespace complete successfully
1244
+
1245
+ Additionally, you can specify project aware upstream flow dependencies
1246
+ by specifying the fully qualified project_flow_name.
1247
+ ```
1248
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1249
+ ```
1250
+ or
1251
+ ```
1252
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1253
+ ```
1254
+
1255
+ You can also specify just the project or project branch (other values will be
1256
+ inferred from the current project or project branch):
1257
+ ```
1258
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1259
+ ```
1260
+
1261
+ Note that `branch` is typically one of:
1262
+ - `prod`
1263
+ - `user.bob`
1264
+ - `test.my_experiment`
1265
+ - `prod.staging`
1266
+
1267
+
1268
+ Parameters
1269
+ ----------
1270
+ flow : Union[str, Dict[str, str]], optional, default None
1271
+ Upstream flow dependency for this flow.
1272
+ flows : List[Union[str, Dict[str, str]]], default []
1273
+ Upstream flow dependencies for this flow.
1274
+ options : Dict[str, Any], default {}
1275
+ Backend-specific configuration for tuning eventing behavior.
1276
+ """
1277
+ ...
1278
+
1279
+ @typing.overload
1280
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1281
+ ...
1282
+
1283
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, 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] = {}):
1284
+ """
1285
+ Specifies the flow(s) that this flow depends on.
1286
+
1287
+ ```
1288
+ @trigger_on_finish(flow='FooFlow')
1289
+ ```
1290
+ or
1291
+ ```
1292
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1293
+ ```
1294
+ This decorator respects the @project decorator and triggers the flow
1295
+ when upstream runs within the same namespace complete successfully
1296
+
1297
+ Additionally, you can specify project aware upstream flow dependencies
1298
+ by specifying the fully qualified project_flow_name.
1299
+ ```
1300
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1301
+ ```
1302
+ or
1303
+ ```
1304
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1305
+ ```
1306
+
1307
+ You can also specify just the project or project branch (other values will be
1308
+ inferred from the current project or project branch):
1309
+ ```
1310
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1311
+ ```
1312
+
1313
+ Note that `branch` is typically one of:
1314
+ - `prod`
1315
+ - `user.bob`
1316
+ - `test.my_experiment`
1317
+ - `prod.staging`
1318
+
1319
+
1320
+ Parameters
1321
+ ----------
1322
+ flow : Union[str, Dict[str, str]], optional, default None
1323
+ Upstream flow dependency for this flow.
1324
+ flows : List[Union[str, Dict[str, str]]], default []
1325
+ Upstream flow dependencies for this flow.
1326
+ options : Dict[str, Any], default {}
1327
+ Backend-specific configuration for tuning eventing behavior.
1328
+ """
1329
+ ...
1330
+
1331
+ 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]]:
1332
+ """
1333
+ 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)
1334
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1335
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1336
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1337
+ starts only after all sensors finish.
1338
+
1339
+
1340
+ Parameters
1341
+ ----------
1342
+ timeout : int
1343
+ Time, in seconds before the task times out and fails. (Default: 3600)
1344
+ poke_interval : int
1345
+ Time in seconds that the job should wait in between each try. (Default: 60)
1346
+ mode : str
1347
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1348
+ exponential_backoff : bool
1349
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1350
+ pool : str
1351
+ the slot pool this task should run in,
1352
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1353
+ soft_fail : bool
1354
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1355
+ name : str
1356
+ Name of the sensor on Airflow
1357
+ description : str
1358
+ Description of sensor in the Airflow UI
1359
+ bucket_key : Union[str, List[str]]
1360
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1361
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1362
+ bucket_name : str
1363
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1364
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1365
+ wildcard_match : bool
1366
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1367
+ aws_conn_id : str
1368
+ a reference to the s3 connection on Airflow. (Default: None)
1369
+ verify : bool
1370
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1371
+ """
1372
+ ...
1373
+
1374
+ @typing.overload
1375
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1376
+ """
1377
+ Specifies the PyPI packages for all steps of the flow.
1378
+
1379
+ Use `@pypi_base` to set common packages required by all
1380
+ steps and use `@pypi` to specify step-specific overrides.
1381
+
1382
+ Parameters
1383
+ ----------
1384
+ packages : Dict[str, str], default: {}
1385
+ Packages to use for this flow. The key is the name of the package
1386
+ and the value is the version to use.
1387
+ python : str, optional, default: None
1388
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1389
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1390
+ """
1391
+ ...
1392
+
1393
+ @typing.overload
1394
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1395
+ ...
1396
+
1397
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1398
+ """
1399
+ Specifies the PyPI packages for all steps of the flow.
1400
+
1401
+ Use `@pypi_base` to set common packages required by all
1402
+ steps and use `@pypi` to specify step-specific overrides.
1403
+
1404
+ Parameters
1405
+ ----------
1406
+ packages : Dict[str, str], default: {}
1407
+ Packages to use for this flow. The key is the name of the package
1408
+ and the value is the version to use.
1409
+ python : str, optional, default: None
1410
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1411
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1412
+ """
1413
+ ...
1414
+
1415
1415
  pkg_name: str
1416
1416