ob-metaflow-stubs 6.0.3.143rc0__py2.py3-none-any.whl → 6.0.3.144__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (199) hide show
  1. metaflow-stubs/__init__.pyi +727 -725
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -7
  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 +95 -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 +3 -3
  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 +26 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +3 -5
  20. metaflow-stubs/metaflow_current.pyi +85 -85
  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 +2 -2
  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 +4 -4
  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 +4 -4
  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 +4 -6
  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 +5 -8
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +5 -6
  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 +3 -5
  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 +8 -4
  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 +5 -12
  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 +3 -3
  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 +4 -4
  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 +3 -3
  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 +3 -3
  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 +6 -0
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +23 -0
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +6 -0
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +29 -0
  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 +3 -3
  76. metaflow-stubs/plugins/__init__.pyi +11 -10
  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 +2 -2
  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 +3 -3
  89. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  90. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +4 -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 +3 -3
  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 +3 -3
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +4 -4
  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 +3 -3
  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 +2 -2
  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 +3 -3
  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 +4 -4
  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 +3 -3
  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 +4 -4
  151. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  152. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -5
  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 +53 -0
  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 +6 -6
  162. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  164. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  165. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  166. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  168. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +4 -4
  169. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  170. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  171. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  172. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  173. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  174. metaflow-stubs/profilers/__init__.pyi +2 -2
  175. metaflow-stubs/pylint_wrapper.pyi +2 -2
  176. metaflow-stubs/runner/__init__.pyi +2 -2
  177. metaflow-stubs/runner/deployer.pyi +6 -6
  178. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  179. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  180. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  181. metaflow-stubs/runner/nbrun.pyi +2 -2
  182. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  183. metaflow-stubs/runner/utils.pyi +3 -3
  184. metaflow-stubs/system/__init__.pyi +2 -2
  185. metaflow-stubs/system/system_logger.pyi +3 -3
  186. metaflow-stubs/system/system_monitor.pyi +2 -2
  187. metaflow-stubs/tagging_util.pyi +2 -2
  188. metaflow-stubs/tuple_util.pyi +2 -2
  189. metaflow-stubs/user_configs/__init__.pyi +2 -2
  190. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  191. metaflow-stubs/user_configs/config_options.pyi +3 -3
  192. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  193. {ob_metaflow_stubs-6.0.3.143rc0.dist-info → ob_metaflow_stubs-6.0.3.144.dist-info}/METADATA +1 -1
  194. ob_metaflow_stubs-6.0.3.144.dist-info/RECORD +197 -0
  195. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +0 -79
  196. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +0 -46
  197. ob_metaflow_stubs-6.0.3.143rc0.dist-info/RECORD +0 -194
  198. {ob_metaflow_stubs-6.0.3.143rc0.dist-info → ob_metaflow_stubs-6.0.3.144.dist-info}/WHEEL +0 -0
  199. {ob_metaflow_stubs-6.0.3.143rc0.dist-info → ob_metaflow_stubs-6.0.3.144.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.14.0.2rc0+obcheckpoint(0.1.8.1rc0);ob(v1) #
4
- # Generated on 2025-02-17T19:27:36.178389 #
3
+ # MF version: 2.14.2.1+obcheckpoint(0.1.8);ob(v1) #
4
+ # Generated on 2025-02-21T20:43:13.534812 #
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
 
@@ -63,7 +63,6 @@ from .runner.deployer import DeployedFlow as DeployedFlow
63
63
  from .runner.nbdeploy import NBDeployer as NBDeployer
64
64
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.checkpoints.final_api import Checkpoint as Checkpoint
65
65
  from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastructures import load_model as load_model
66
- from .mf_extensions.obcheckpoint.plugins.machine_learning_utilities.datastore.context import using_artifact_store as using_artifact_store
67
66
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_aws_client as get_aws_client
68
67
  from .mf_extensions.outerbounds.plugins.snowflake.snowflake import Snowflake as Snowflake
69
68
  from . import cli_components as cli_components
@@ -147,160 +146,6 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
147
146
  """
148
147
  ...
149
148
 
150
- 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]]]:
151
- """
152
- Decorator that helps cache, version and store models/datasets from huggingface hub.
153
-
154
-
155
- Parameters
156
- ----------
157
- temp_dir_root : str, optional
158
- The root directory that will hold the temporary directory where objects will be downloaded.
159
-
160
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
161
- The list of repos (models/datasets) to load.
162
-
163
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
164
-
165
- - If repo (model/dataset) is not found in the datastore:
166
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
167
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
168
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
169
-
170
- - If repo is found in the datastore:
171
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
172
- """
173
- ...
174
-
175
- @typing.overload
176
- 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]]]:
177
- """
178
- Specifies a timeout for your step.
179
-
180
- This decorator is useful if this step may hang indefinitely.
181
-
182
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
183
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
184
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
185
-
186
- Note that all the values specified in parameters are added together so if you specify
187
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
188
-
189
-
190
- Parameters
191
- ----------
192
- seconds : int, default 0
193
- Number of seconds to wait prior to timing out.
194
- minutes : int, default 0
195
- Number of minutes to wait prior to timing out.
196
- hours : int, default 0
197
- Number of hours to wait prior to timing out.
198
- """
199
- ...
200
-
201
- @typing.overload
202
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
203
- ...
204
-
205
- @typing.overload
206
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
207
- ...
208
-
209
- 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):
210
- """
211
- Specifies a timeout for your step.
212
-
213
- This decorator is useful if this step may hang indefinitely.
214
-
215
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
216
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
217
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
218
-
219
- Note that all the values specified in parameters are added together so if you specify
220
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
221
-
222
-
223
- Parameters
224
- ----------
225
- seconds : int, default 0
226
- Number of seconds to wait prior to timing out.
227
- minutes : int, default 0
228
- Number of minutes to wait prior to timing out.
229
- hours : int, default 0
230
- Number of hours to wait prior to timing out.
231
- """
232
- ...
233
-
234
- @typing.overload
235
- 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]]]:
236
- """
237
- Specifies that the step will success under all circumstances.
238
-
239
- The decorator will create an optional artifact, specified by `var`, which
240
- contains the exception raised. You can use it to detect the presence
241
- of errors, indicating that all happy-path artifacts produced by the step
242
- are missing.
243
-
244
-
245
- Parameters
246
- ----------
247
- var : str, optional, default None
248
- Name of the artifact in which to store the caught exception.
249
- If not specified, the exception is not stored.
250
- print_exception : bool, default True
251
- Determines whether or not the exception is printed to
252
- stdout when caught.
253
- """
254
- ...
255
-
256
- @typing.overload
257
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
258
- ...
259
-
260
- @typing.overload
261
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
262
- ...
263
-
264
- 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):
265
- """
266
- Specifies that the step will success under all circumstances.
267
-
268
- The decorator will create an optional artifact, specified by `var`, which
269
- contains the exception raised. You can use it to detect the presence
270
- of errors, indicating that all happy-path artifacts produced by the step
271
- are missing.
272
-
273
-
274
- Parameters
275
- ----------
276
- var : str, optional, default None
277
- Name of the artifact in which to store the caught exception.
278
- If not specified, the exception is not stored.
279
- print_exception : bool, default True
280
- Determines whether or not the exception is printed to
281
- stdout when caught.
282
- """
283
- ...
284
-
285
- @typing.overload
286
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
287
- """
288
- Parameters
289
- ----------
290
- """
291
- ...
292
-
293
- @typing.overload
294
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
295
- ...
296
-
297
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
298
- """
299
- Parameters
300
- ----------
301
- """
302
- ...
303
-
304
149
  @typing.overload
305
150
  def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
306
151
  """
@@ -358,236 +203,44 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
358
203
  """
359
204
  ...
360
205
 
361
- 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]]]:
362
- """
363
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
364
-
365
- User code call
366
- -----------
367
- @nim(
368
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
369
- backend='managed'
370
- )
371
-
372
- Valid backend options
373
- ---------------------
374
- - 'managed': Outerbounds selects a compute provider based on the model.
375
-
376
- Valid model options
377
- ----------------
378
- - 'meta/llama3-8b-instruct': 8B parameter model
379
- - 'meta/llama3-70b-instruct': 70B parameter model
380
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
381
-
382
-
383
- Parameters
384
- ----------
385
- models: list[NIM]
386
- List of NIM containers running models in sidecars.
387
- backend: str
388
- Compute provider to run the NIM container.
389
- queue_timeout : int
390
- Time to keep the job in NVCF's queue.
391
- """
392
- ...
393
-
394
206
  @typing.overload
395
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> 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]]]:
207
+ 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]]]:
396
208
  """
397
- Specifies the number of times the task corresponding
398
- to a step needs to be retried.
399
-
400
- This decorator is useful for handling transient errors, such as networking issues.
401
- If your task contains operations that can't be retried safely, e.g. database updates,
402
- it is advisable to annotate it with `@retry(times=0)`.
403
-
404
- This can be used in conjunction with the `@catch` decorator. The `@catch`
405
- decorator will execute a no-op task after all retries have been exhausted,
406
- ensuring that the flow execution can continue.
209
+ Specifies secrets to be retrieved and injected as environment variables prior to
210
+ the execution of a step.
407
211
 
408
212
 
409
213
  Parameters
410
214
  ----------
411
- times : int, default 3
412
- Number of times to retry this task.
413
- minutes_between_retries : int, default 2
414
- Number of minutes between retries.
215
+ sources : List[Union[str, Dict[str, Any]]], default: []
216
+ List of secret specs, defining how the secrets are to be retrieved
415
217
  """
416
218
  ...
417
219
 
418
220
  @typing.overload
419
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
221
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
420
222
  ...
421
223
 
422
224
  @typing.overload
423
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
225
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
424
226
  ...
425
227
 
426
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
228
+ 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]]] = []):
427
229
  """
428
- Specifies the number of times the task corresponding
429
- to a step needs to be retried.
430
-
431
- This decorator is useful for handling transient errors, such as networking issues.
432
- If your task contains operations that can't be retried safely, e.g. database updates,
433
- it is advisable to annotate it with `@retry(times=0)`.
434
-
435
- This can be used in conjunction with the `@catch` decorator. The `@catch`
436
- decorator will execute a no-op task after all retries have been exhausted,
437
- ensuring that the flow execution can continue.
230
+ Specifies secrets to be retrieved and injected as environment variables prior to
231
+ the execution of a step.
438
232
 
439
233
 
440
234
  Parameters
441
235
  ----------
442
- times : int, default 3
443
- Number of times to retry this task.
444
- minutes_between_retries : int, default 2
445
- Number of minutes between retries.
236
+ sources : List[Union[str, Dict[str, Any]]], default: []
237
+ List of secret specs, defining how the secrets are to be retrieved
446
238
  """
447
239
  ...
448
240
 
449
- @typing.overload
450
- 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]]]:
241
+ 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]]]:
451
242
  """
452
- Enables checkpointing for a step.
453
-
454
-
455
-
456
- Parameters
457
- ----------
458
- load_policy : str, default: "fresh"
459
- The policy for loading the checkpoint. The following policies are supported:
460
- - "eager": Loads the the latest available checkpoint within the namespace.
461
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
462
- will be loaded at the start of the task.
463
- - "none": Do not load any checkpoint
464
- - "fresh": Loads the lastest checkpoint created within the running Task.
465
- This mode helps loading checkpoints across various retry attempts of the same task.
466
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
467
- created within the task will be loaded when the task is retries execution on failure.
468
-
469
- temp_dir_root : str, default: None
470
- The root directory under which `current.checkpoint.directory` will be created.
471
- """
472
- ...
473
-
474
- @typing.overload
475
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
476
- ...
477
-
478
- @typing.overload
479
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
480
- ...
481
-
482
- 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):
483
- """
484
- Enables checkpointing for a step.
485
-
486
-
487
-
488
- Parameters
489
- ----------
490
- load_policy : str, default: "fresh"
491
- The policy for loading the checkpoint. The following policies are supported:
492
- - "eager": Loads the the latest available checkpoint within the namespace.
493
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
494
- will be loaded at the start of the task.
495
- - "none": Do not load any checkpoint
496
- - "fresh": Loads the lastest checkpoint created within the running Task.
497
- This mode helps loading checkpoints across various retry attempts of the same task.
498
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
499
- created within the task will be loaded when the task is retries execution on failure.
500
-
501
- temp_dir_root : str, default: None
502
- The root directory under which `current.checkpoint.directory` will be created.
503
- """
504
- ...
505
-
506
- @typing.overload
507
- 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]]]:
508
- """
509
- Creates a human-readable report, a Metaflow Card, after this step completes.
510
-
511
- Note that you may add multiple `@card` decorators in a step with different parameters.
512
-
513
-
514
- Parameters
515
- ----------
516
- type : str, default 'default'
517
- Card type.
518
- id : str, optional, default None
519
- If multiple cards are present, use this id to identify this card.
520
- options : Dict[str, Any], default {}
521
- Options passed to the card. The contents depend on the card type.
522
- timeout : int, default 45
523
- Interrupt reporting if it takes more than this many seconds.
524
- """
525
- ...
526
-
527
- @typing.overload
528
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
529
- ...
530
-
531
- @typing.overload
532
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
533
- ...
534
-
535
- 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):
536
- """
537
- Creates a human-readable report, a Metaflow Card, after this step completes.
538
-
539
- Note that you may add multiple `@card` decorators in a step with different parameters.
540
-
541
-
542
- Parameters
543
- ----------
544
- type : str, default 'default'
545
- Card type.
546
- id : str, optional, default None
547
- If multiple cards are present, use this id to identify this card.
548
- options : Dict[str, Any], default {}
549
- Options passed to the card. The contents depend on the card type.
550
- timeout : int, default 45
551
- Interrupt reporting if it takes more than this many seconds.
552
- """
553
- ...
554
-
555
- @typing.overload
556
- 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]]]:
557
- """
558
- Specifies environment variables to be set prior to the execution of a step.
559
-
560
-
561
- Parameters
562
- ----------
563
- vars : Dict[str, str], default {}
564
- Dictionary of environment variables to set.
565
- """
566
- ...
567
-
568
- @typing.overload
569
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
570
- ...
571
-
572
- @typing.overload
573
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
574
- ...
575
-
576
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
577
- """
578
- Specifies environment variables to be set prior to the execution of a step.
579
-
580
-
581
- Parameters
582
- ----------
583
- vars : Dict[str, str], default {}
584
- Dictionary of environment variables to set.
585
- """
586
- ...
587
-
588
- 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', extended_resources: typing.Optional[typing.Dict[str, 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]]]:
589
- """
590
- Specifies that this step should execute on Kubernetes.
243
+ Specifies that this step should execute on Kubernetes.
591
244
 
592
245
 
593
246
  Parameters
@@ -659,9 +312,39 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
659
312
  Only applicable when @parallel is used.
660
313
  qos: str, default: Burstable
661
314
  Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
662
- extended_resources: Dict[str, str], optional, default None
663
- Extended resources to be requested for the pod.
664
- https://kubernetes.io/docs/tasks/administer-cluster/extended-resource-node/
315
+ """
316
+ ...
317
+
318
+ @typing.overload
319
+ 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]]]:
320
+ """
321
+ Specifies environment variables to be set prior to the execution of a step.
322
+
323
+
324
+ Parameters
325
+ ----------
326
+ vars : Dict[str, str], default {}
327
+ Dictionary of environment variables to set.
328
+ """
329
+ ...
330
+
331
+ @typing.overload
332
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
333
+ ...
334
+
335
+ @typing.overload
336
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
337
+ ...
338
+
339
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
340
+ """
341
+ Specifies environment variables to be set prior to the execution of a step.
342
+
343
+
344
+ Parameters
345
+ ----------
346
+ vars : Dict[str, str], default {}
347
+ Dictionary of environment variables to set.
665
348
  """
666
349
  ...
667
350
 
@@ -725,130 +408,344 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
725
408
  ...
726
409
 
727
410
  @typing.overload
728
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
411
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> 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]]]:
729
412
  """
730
- Specifies the resources needed when executing this step.
413
+ Specifies the number of times the task corresponding
414
+ to a step needs to be retried.
731
415
 
732
- Use `@resources` to specify the resource requirements
733
- independently of the specific compute layer (`@batch`, `@kubernetes`).
416
+ This decorator is useful for handling transient errors, such as networking issues.
417
+ If your task contains operations that can't be retried safely, e.g. database updates,
418
+ it is advisable to annotate it with `@retry(times=0)`.
734
419
 
735
- You can choose the compute layer on the command line by executing e.g.
736
- ```
737
- python myflow.py run --with batch
738
- ```
739
- or
740
- ```
741
- python myflow.py run --with kubernetes
742
- ```
743
- which executes the flow on the desired system using the
744
- requirements specified in `@resources`.
420
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
421
+ decorator will execute a no-op task after all retries have been exhausted,
422
+ ensuring that the flow execution can continue.
745
423
 
746
424
 
747
425
  Parameters
748
426
  ----------
749
- cpu : int, default 1
750
- Number of CPUs required for this step.
751
- gpu : int, optional, default None
752
- Number of GPUs required for this step.
753
- disk : int, optional, default None
754
- Disk size (in MB) required for this step. Only applies on Kubernetes.
755
- memory : int, default 4096
756
- Memory size (in MB) required for this step.
757
- shared_memory : int, optional, default None
758
- The value for the size (in MiB) of the /dev/shm volume for this step.
759
- This parameter maps to the `--shm-size` option in Docker.
427
+ times : int, default 3
428
+ Number of times to retry this task.
429
+ minutes_between_retries : int, default 2
430
+ Number of minutes between retries.
760
431
  """
761
432
  ...
762
433
 
763
434
  @typing.overload
764
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
435
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
765
436
  ...
766
437
 
767
438
  @typing.overload
768
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
439
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
769
440
  ...
770
441
 
771
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
442
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
772
443
  """
773
- Specifies the resources needed when executing this step.
444
+ Specifies the number of times the task corresponding
445
+ to a step needs to be retried.
774
446
 
775
- Use `@resources` to specify the resource requirements
776
- independently of the specific compute layer (`@batch`, `@kubernetes`).
447
+ This decorator is useful for handling transient errors, such as networking issues.
448
+ If your task contains operations that can't be retried safely, e.g. database updates,
449
+ it is advisable to annotate it with `@retry(times=0)`.
777
450
 
778
- You can choose the compute layer on the command line by executing e.g.
779
- ```
780
- python myflow.py run --with batch
781
- ```
782
- or
783
- ```
784
- python myflow.py run --with kubernetes
785
- ```
786
- which executes the flow on the desired system using the
787
- requirements specified in `@resources`.
451
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
452
+ decorator will execute a no-op task after all retries have been exhausted,
453
+ ensuring that the flow execution can continue.
788
454
 
789
455
 
790
456
  Parameters
791
457
  ----------
792
- cpu : int, default 1
793
- Number of CPUs required for this step.
794
- gpu : int, optional, default None
795
- Number of GPUs required for this step.
796
- disk : int, optional, default None
797
- Disk size (in MB) required for this step. Only applies on Kubernetes.
798
- memory : int, default 4096
799
- Memory size (in MB) required for this step.
800
- shared_memory : int, optional, default None
801
- The value for the size (in MiB) of the /dev/shm volume for this step.
802
- This parameter maps to the `--shm-size` option in Docker.
458
+ times : int, default 3
459
+ Number of times to retry this task.
460
+ minutes_between_retries : int, default 2
461
+ Number of minutes between retries.
803
462
  """
804
463
  ...
805
464
 
806
- @typing.overload
807
- 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]]]:
465
+ 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]]]:
808
466
  """
809
- Specifies secrets to be retrieved and injected as environment variables prior to
810
- the execution of a step.
467
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
811
468
 
812
469
 
813
470
  Parameters
814
471
  ----------
815
- sources : List[Union[str, Dict[str, Any]]], default: []
816
- List of secret specs, defining how the secrets are to be retrieved
472
+ temp_dir_root : str, optional
473
+ The root directory that will hold the temporary directory where objects will be downloaded.
474
+
475
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
476
+ The list of repos (models/datasets) to load.
477
+
478
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
479
+
480
+ - If repo (model/dataset) is not found in the datastore:
481
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
482
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
483
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
484
+
485
+ - If repo is found in the datastore:
486
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
817
487
  """
818
488
  ...
819
489
 
820
490
  @typing.overload
821
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
822
- ...
491
+ 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]]]:
492
+ """
493
+ Creates a human-readable report, a Metaflow Card, after this step completes.
494
+
495
+ Note that you may add multiple `@card` decorators in a step with different parameters.
496
+
497
+
498
+ Parameters
499
+ ----------
500
+ type : str, default 'default'
501
+ Card type.
502
+ id : str, optional, default None
503
+ If multiple cards are present, use this id to identify this card.
504
+ options : Dict[str, Any], default {}
505
+ Options passed to the card. The contents depend on the card type.
506
+ timeout : int, default 45
507
+ Interrupt reporting if it takes more than this many seconds.
508
+ """
509
+ ...
823
510
 
824
511
  @typing.overload
825
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
512
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
826
513
  ...
827
514
 
828
- 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]]] = []):
515
+ @typing.overload
516
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
517
+ ...
518
+
519
+ 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):
829
520
  """
830
- Specifies secrets to be retrieved and injected as environment variables prior to
831
- the execution of a step.
521
+ Creates a human-readable report, a Metaflow Card, after this step completes.
522
+
523
+ Note that you may add multiple `@card` decorators in a step with different parameters.
832
524
 
833
525
 
834
526
  Parameters
835
527
  ----------
836
- sources : List[Union[str, Dict[str, Any]]], default: []
837
- List of secret specs, defining how the secrets are to be retrieved
528
+ type : str, default 'default'
529
+ Card type.
530
+ id : str, optional, default None
531
+ If multiple cards are present, use this id to identify this card.
532
+ options : Dict[str, Any], default {}
533
+ Options passed to the card. The contents depend on the card type.
534
+ timeout : int, default 45
535
+ Interrupt reporting if it takes more than this many seconds.
838
536
  """
839
537
  ...
840
538
 
841
- 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]]]:
539
+ 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]]]:
842
540
  """
843
- Specifies that this step should execute on DGX cloud.
541
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
542
+
543
+ User code call
544
+ -----------
545
+ @ollama(
546
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
547
+ backend='local'
548
+ )
549
+
550
+ Valid backend options
551
+ ---------------------
552
+ - 'local': Run as a separate process on the local task machine.
553
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
554
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
555
+
556
+ Valid model options
557
+ ----------------
558
+ - 'llama3.2'
559
+ - 'llama3.3'
560
+ - any model here https://ollama.com/search
844
561
 
845
562
 
846
563
  Parameters
847
564
  ----------
848
- gpu : int
849
- Number of GPUs to use.
850
- gpu_type : str
851
- Type of Nvidia GPU to use.
565
+ models: list[Ollama]
566
+ List of Ollama containers running models in sidecars.
567
+ backend: str
568
+ Determines where and how to run the Ollama process.
569
+ """
570
+ ...
571
+
572
+ @typing.overload
573
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
574
+ """
575
+ Parameters
576
+ ----------
577
+ """
578
+ ...
579
+
580
+ @typing.overload
581
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
582
+ ...
583
+
584
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
585
+ """
586
+ Parameters
587
+ ----------
588
+ """
589
+ ...
590
+
591
+ @typing.overload
592
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
593
+ """
594
+ Specifies the resources needed when executing this step.
595
+
596
+ Use `@resources` to specify the resource requirements
597
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
598
+
599
+ You can choose the compute layer on the command line by executing e.g.
600
+ ```
601
+ python myflow.py run --with batch
602
+ ```
603
+ or
604
+ ```
605
+ python myflow.py run --with kubernetes
606
+ ```
607
+ which executes the flow on the desired system using the
608
+ requirements specified in `@resources`.
609
+
610
+
611
+ Parameters
612
+ ----------
613
+ cpu : int, default 1
614
+ Number of CPUs required for this step.
615
+ gpu : int, optional, default None
616
+ Number of GPUs required for this step.
617
+ disk : int, optional, default None
618
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
619
+ memory : int, default 4096
620
+ Memory size (in MB) required for this step.
621
+ shared_memory : int, optional, default None
622
+ The value for the size (in MiB) of the /dev/shm volume for this step.
623
+ This parameter maps to the `--shm-size` option in Docker.
624
+ """
625
+ ...
626
+
627
+ @typing.overload
628
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
629
+ ...
630
+
631
+ @typing.overload
632
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
633
+ ...
634
+
635
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
636
+ """
637
+ Specifies the resources needed when executing this step.
638
+
639
+ Use `@resources` to specify the resource requirements
640
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
641
+
642
+ You can choose the compute layer on the command line by executing e.g.
643
+ ```
644
+ python myflow.py run --with batch
645
+ ```
646
+ or
647
+ ```
648
+ python myflow.py run --with kubernetes
649
+ ```
650
+ which executes the flow on the desired system using the
651
+ requirements specified in `@resources`.
652
+
653
+
654
+ Parameters
655
+ ----------
656
+ cpu : int, default 1
657
+ Number of CPUs required for this step.
658
+ gpu : int, optional, default None
659
+ Number of GPUs required for this step.
660
+ disk : int, optional, default None
661
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
662
+ memory : int, default 4096
663
+ Memory size (in MB) required for this step.
664
+ shared_memory : int, optional, default None
665
+ The value for the size (in MiB) of the /dev/shm volume for this step.
666
+ This parameter maps to the `--shm-size` option in Docker.
667
+ """
668
+ ...
669
+
670
+ @typing.overload
671
+ 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]]]:
672
+ """
673
+ Specifies that the step will success under all circumstances.
674
+
675
+ The decorator will create an optional artifact, specified by `var`, which
676
+ contains the exception raised. You can use it to detect the presence
677
+ of errors, indicating that all happy-path artifacts produced by the step
678
+ are missing.
679
+
680
+
681
+ Parameters
682
+ ----------
683
+ var : str, optional, default None
684
+ Name of the artifact in which to store the caught exception.
685
+ If not specified, the exception is not stored.
686
+ print_exception : bool, default True
687
+ Determines whether or not the exception is printed to
688
+ stdout when caught.
689
+ """
690
+ ...
691
+
692
+ @typing.overload
693
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
694
+ ...
695
+
696
+ @typing.overload
697
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
698
+ ...
699
+
700
+ 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):
701
+ """
702
+ Specifies that the step will success under all circumstances.
703
+
704
+ The decorator will create an optional artifact, specified by `var`, which
705
+ contains the exception raised. You can use it to detect the presence
706
+ of errors, indicating that all happy-path artifacts produced by the step
707
+ are missing.
708
+
709
+
710
+ Parameters
711
+ ----------
712
+ var : str, optional, default None
713
+ Name of the artifact in which to store the caught exception.
714
+ If not specified, the exception is not stored.
715
+ print_exception : bool, default True
716
+ Determines whether or not the exception is printed to
717
+ stdout when caught.
718
+ """
719
+ ...
720
+
721
+ 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]]]:
722
+ """
723
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
724
+
725
+ User code call
726
+ -----------
727
+ @nim(
728
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
729
+ backend='managed'
730
+ )
731
+
732
+ Valid backend options
733
+ ---------------------
734
+ - 'managed': Outerbounds selects a compute provider based on the model.
735
+
736
+ Valid model options
737
+ ----------------
738
+ - 'meta/llama3-8b-instruct': 8B parameter model
739
+ - 'meta/llama3-70b-instruct': 70B parameter model
740
+ - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
741
+
742
+
743
+ Parameters
744
+ ----------
745
+ models: list[NIM]
746
+ List of NIM containers running models in sidecars.
747
+ backend: str
748
+ Compute provider to run the NIM container.
852
749
  queue_timeout : int
853
750
  Time to keep the job in NVCF's queue.
854
751
  """
@@ -877,52 +774,235 @@ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] =
877
774
  ...
878
775
 
879
776
  @typing.overload
880
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
777
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
778
+ ...
779
+
780
+ @typing.overload
781
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
782
+ ...
783
+
784
+ 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):
785
+ """
786
+ Specifies the PyPI packages for the step.
787
+
788
+ Information in this decorator will augment any
789
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
790
+ you can use `@pypi_base` to set packages required by all
791
+ steps and use `@pypi` to specify step-specific overrides.
792
+
793
+
794
+ Parameters
795
+ ----------
796
+ packages : Dict[str, str], default: {}
797
+ Packages to use for this step. The key is the name of the package
798
+ and the value is the version to use.
799
+ python : str, optional, default: None
800
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
801
+ that the version used will correspond to the version of the Python interpreter used to start the run.
802
+ """
803
+ ...
804
+
805
+ @typing.overload
806
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
807
+ """
808
+ Internal decorator to support Fast bakery
809
+ Parameters
810
+ ----------
811
+ """
812
+ ...
813
+
814
+ @typing.overload
815
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
816
+ ...
817
+
818
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
819
+ """
820
+ Internal decorator to support Fast bakery
821
+ Parameters
822
+ ----------
823
+ """
824
+ ...
825
+
826
+ 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]]]:
827
+ """
828
+ Specifies that this step should execute on DGX cloud.
829
+
830
+
831
+ Parameters
832
+ ----------
833
+ gpu : int
834
+ Number of GPUs to use.
835
+ gpu_type : str
836
+ Type of Nvidia GPU to use.
837
+ queue_timeout : int
838
+ Time to keep the job in NVCF's queue.
839
+ """
840
+ ...
841
+
842
+ @typing.overload
843
+ 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]]]:
844
+ """
845
+ Specifies a timeout for your step.
846
+
847
+ This decorator is useful if this step may hang indefinitely.
848
+
849
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
850
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
851
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
852
+
853
+ Note that all the values specified in parameters are added together so if you specify
854
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
855
+
856
+
857
+ Parameters
858
+ ----------
859
+ seconds : int, default 0
860
+ Number of seconds to wait prior to timing out.
861
+ minutes : int, default 0
862
+ Number of minutes to wait prior to timing out.
863
+ hours : int, default 0
864
+ Number of hours to wait prior to timing out.
865
+ """
866
+ ...
867
+
868
+ @typing.overload
869
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
870
+ ...
871
+
872
+ @typing.overload
873
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
874
+ ...
875
+
876
+ 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):
877
+ """
878
+ Specifies a timeout for your step.
879
+
880
+ This decorator is useful if this step may hang indefinitely.
881
+
882
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
883
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
884
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
885
+
886
+ Note that all the values specified in parameters are added together so if you specify
887
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
888
+
889
+
890
+ Parameters
891
+ ----------
892
+ seconds : int, default 0
893
+ Number of seconds to wait prior to timing out.
894
+ minutes : int, default 0
895
+ Number of minutes to wait prior to timing out.
896
+ hours : int, default 0
897
+ Number of hours to wait prior to timing out.
898
+ """
899
+ ...
900
+
901
+ @typing.overload
902
+ 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]]]:
903
+ """
904
+ Enables checkpointing for a step.
905
+
906
+
907
+
908
+ Parameters
909
+ ----------
910
+ load_policy : str, default: "fresh"
911
+ The policy for loading the checkpoint. The following policies are supported:
912
+ - "eager": Loads the the latest available checkpoint within the namespace.
913
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
914
+ will be loaded at the start of the task.
915
+ - "none": Do not load any checkpoint
916
+ - "fresh": Loads the lastest checkpoint created within the running Task.
917
+ This mode helps loading checkpoints across various retry attempts of the same task.
918
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
919
+ created within the task will be loaded when the task is retries execution on failure.
920
+
921
+ temp_dir_root : str, default: None
922
+ The root directory under which `current.checkpoint.directory` will be created.
923
+ """
924
+ ...
925
+
926
+ @typing.overload
927
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
881
928
  ...
882
929
 
883
930
  @typing.overload
884
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
931
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
885
932
  ...
886
933
 
887
- 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):
934
+ 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):
888
935
  """
889
- Specifies the PyPI packages for the step.
936
+ Enables checkpointing for a step.
890
937
 
891
- Information in this decorator will augment any
892
- attributes set in the `@pyi_base` flow-level decorator. Hence,
893
- you can use `@pypi_base` to set packages required by all
894
- steps and use `@pypi` to specify step-specific overrides.
895
938
 
896
939
 
897
940
  Parameters
898
941
  ----------
899
- packages : Dict[str, str], default: {}
900
- Packages to use for this step. The key is the name of the package
901
- and the value is the version to use.
902
- python : str, optional, default: None
903
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
904
- that the version used will correspond to the version of the Python interpreter used to start the run.
942
+ load_policy : str, default: "fresh"
943
+ The policy for loading the checkpoint. The following policies are supported:
944
+ - "eager": Loads the the latest available checkpoint within the namespace.
945
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
946
+ will be loaded at the start of the task.
947
+ - "none": Do not load any checkpoint
948
+ - "fresh": Loads the lastest checkpoint created within the running Task.
949
+ This mode helps loading checkpoints across various retry attempts of the same task.
950
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
951
+ created within the task will be loaded when the task is retries execution on failure.
952
+
953
+ temp_dir_root : str, default: None
954
+ The root directory under which `current.checkpoint.directory` will be created.
905
955
  """
906
956
  ...
907
957
 
908
958
  @typing.overload
909
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
959
+ 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]]:
910
960
  """
911
- Internal decorator to support Fast bakery
961
+ Specifies the times when the flow should be run when running on a
962
+ production scheduler.
963
+
964
+
912
965
  Parameters
913
966
  ----------
967
+ hourly : bool, default False
968
+ Run the workflow hourly.
969
+ daily : bool, default True
970
+ Run the workflow daily.
971
+ weekly : bool, default False
972
+ Run the workflow weekly.
973
+ cron : str, optional, default None
974
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
975
+ specified by this expression.
976
+ timezone : str, optional, default None
977
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
978
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
914
979
  """
915
980
  ...
916
981
 
917
982
  @typing.overload
918
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
983
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
919
984
  ...
920
985
 
921
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
986
+ 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):
922
987
  """
923
- Internal decorator to support Fast bakery
988
+ Specifies the times when the flow should be run when running on a
989
+ production scheduler.
990
+
991
+
924
992
  Parameters
925
993
  ----------
994
+ hourly : bool, default False
995
+ Run the workflow hourly.
996
+ daily : bool, default True
997
+ Run the workflow daily.
998
+ weekly : bool, default False
999
+ Run the workflow weekly.
1000
+ cron : str, optional, default None
1001
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1002
+ specified by this expression.
1003
+ timezone : str, optional, default None
1004
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1005
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
926
1006
  """
927
1007
  ...
928
1008
 
@@ -977,157 +1057,104 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
977
1057
  """
978
1058
  ...
979
1059
 
980
- 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]]:
1060
+ @typing.overload
1061
+ 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]]:
981
1062
  """
982
- 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.
983
- 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.
1063
+ Specifies the flow(s) that this flow depends on.
984
1064
 
1065
+ ```
1066
+ @trigger_on_finish(flow='FooFlow')
1067
+ ```
1068
+ or
1069
+ ```
1070
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1071
+ ```
1072
+ This decorator respects the @project decorator and triggers the flow
1073
+ when upstream runs within the same namespace complete successfully
985
1074
 
986
- Parameters
987
- ----------
988
- timeout : int
989
- Time, in seconds before the task times out and fails. (Default: 3600)
990
- poke_interval : int
991
- Time in seconds that the job should wait in between each try. (Default: 60)
992
- mode : str
993
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
994
- exponential_backoff : bool
995
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
996
- pool : str
997
- the slot pool this task should run in,
998
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
999
- soft_fail : bool
1000
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1001
- name : str
1002
- Name of the sensor on Airflow
1003
- description : str
1004
- Description of sensor in the Airflow UI
1005
- external_dag_id : str
1006
- The dag_id that contains the task you want to wait for.
1007
- external_task_ids : List[str]
1008
- The list of task_ids that you want to wait for.
1009
- If None (default value) the sensor waits for the DAG. (Default: None)
1010
- allowed_states : List[str]
1011
- Iterable of allowed states, (Default: ['success'])
1012
- failed_states : List[str]
1013
- Iterable of failed or dis-allowed states. (Default: None)
1014
- execution_delta : datetime.timedelta
1015
- time difference with the previous execution to look at,
1016
- the default is the same logical date as the current task or DAG. (Default: None)
1017
- check_existence: bool
1018
- Set to True to check if the external task exists or check if
1019
- the DAG to wait for exists. (Default: True)
1020
- """
1021
- ...
1022
-
1023
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1024
- """
1025
- Allows setting external datastores to save data for the
1026
- `@checkpoint`/`@model`/`@huggingface_hub` decorators. This is useful
1027
- when the compute medium lives in a different geographical location
1028
- than metaflow's datastore (e.g. S3, GCS, Azure Blob Storage).
1075
+ Additionally, you can specify project aware upstream flow dependencies
1076
+ by specifying the fully qualified project_flow_name.
1077
+ ```
1078
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1079
+ ```
1080
+ or
1081
+ ```
1082
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1083
+ ```
1029
1084
 
1030
- Parameters:
1031
- ----------
1085
+ You can also specify just the project or project branch (other values will be
1086
+ inferred from the current project or project branch):
1087
+ ```
1088
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1089
+ ```
1032
1090
 
1033
- type: str
1034
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1035
-
1036
- config: dict or Callable
1037
- Dictionary of configuration options for the datastore. The following keys are required:
1038
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1039
- - example: 's3://bucket-name/path/to/root'
1040
- - example: 'gs://bucket-name/path/to/root'
1041
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1042
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1043
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1044
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1045
- Parameters
1046
- ----------
1047
- """
1048
- ...
1049
-
1050
- @typing.overload
1051
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1052
- """
1053
- Specifies the PyPI packages for all steps of the flow.
1091
+ Note that `branch` is typically one of:
1092
+ - `prod`
1093
+ - `user.bob`
1094
+ - `test.my_experiment`
1095
+ - `prod.staging`
1054
1096
 
1055
- Use `@pypi_base` to set common packages required by all
1056
- steps and use `@pypi` to specify step-specific overrides.
1057
1097
 
1058
1098
  Parameters
1059
1099
  ----------
1060
- packages : Dict[str, str], default: {}
1061
- Packages to use for this flow. The key is the name of the package
1062
- and the value is the version to use.
1063
- python : str, optional, default: None
1064
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1065
- that the version used will correspond to the version of the Python interpreter used to start the run.
1100
+ flow : Union[str, Dict[str, str]], optional, default None
1101
+ Upstream flow dependency for this flow.
1102
+ flows : List[Union[str, Dict[str, str]]], default []
1103
+ Upstream flow dependencies for this flow.
1104
+ options : Dict[str, Any], default {}
1105
+ Backend-specific configuration for tuning eventing behavior.
1066
1106
  """
1067
1107
  ...
1068
1108
 
1069
1109
  @typing.overload
1070
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1110
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1071
1111
  ...
1072
1112
 
1073
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1113
+ 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] = {}):
1074
1114
  """
1075
- Specifies the PyPI packages for all steps of the flow.
1115
+ Specifies the flow(s) that this flow depends on.
1076
1116
 
1077
- Use `@pypi_base` to set common packages required by all
1078
- steps and use `@pypi` to specify step-specific overrides.
1117
+ ```
1118
+ @trigger_on_finish(flow='FooFlow')
1119
+ ```
1120
+ or
1121
+ ```
1122
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1123
+ ```
1124
+ This decorator respects the @project decorator and triggers the flow
1125
+ when upstream runs within the same namespace complete successfully
1079
1126
 
1080
- Parameters
1081
- ----------
1082
- packages : Dict[str, str], default: {}
1083
- Packages to use for this flow. The key is the name of the package
1084
- and the value is the version to use.
1085
- python : str, optional, default: None
1086
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1087
- that the version used will correspond to the version of the Python interpreter used to start the run.
1088
- """
1089
- ...
1090
-
1091
- 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]]:
1092
- """
1093
- 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)
1094
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1095
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1096
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1097
- starts only after all sensors finish.
1127
+ Additionally, you can specify project aware upstream flow dependencies
1128
+ by specifying the fully qualified project_flow_name.
1129
+ ```
1130
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1131
+ ```
1132
+ or
1133
+ ```
1134
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1135
+ ```
1098
1136
 
1137
+ You can also specify just the project or project branch (other values will be
1138
+ inferred from the current project or project branch):
1139
+ ```
1140
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1141
+ ```
1099
1142
 
1100
- Parameters
1101
- ----------
1102
- timeout : int
1103
- Time, in seconds before the task times out and fails. (Default: 3600)
1104
- poke_interval : int
1105
- Time in seconds that the job should wait in between each try. (Default: 60)
1106
- mode : str
1107
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1108
- exponential_backoff : bool
1109
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1110
- pool : str
1111
- the slot pool this task should run in,
1112
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1113
- soft_fail : bool
1114
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1115
- name : str
1116
- Name of the sensor on Airflow
1117
- description : str
1118
- Description of sensor in the Airflow UI
1119
- bucket_key : Union[str, List[str]]
1120
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1121
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1122
- bucket_name : str
1123
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1124
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1125
- wildcard_match : bool
1126
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1127
- aws_conn_id : str
1128
- a reference to the s3 connection on Airflow. (Default: None)
1129
- verify : bool
1130
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1143
+ Note that `branch` is typically one of:
1144
+ - `prod`
1145
+ - `user.bob`
1146
+ - `test.my_experiment`
1147
+ - `prod.staging`
1148
+
1149
+
1150
+ Parameters
1151
+ ----------
1152
+ flow : Union[str, Dict[str, str]], optional, default None
1153
+ Upstream flow dependency for this flow.
1154
+ flows : List[Union[str, Dict[str, str]]], default []
1155
+ Upstream flow dependencies for this flow.
1156
+ options : Dict[str, Any], default {}
1157
+ Backend-specific configuration for tuning eventing behavior.
1131
1158
  """
1132
1159
  ...
1133
1160
 
@@ -1224,6 +1251,47 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1224
1251
  """
1225
1252
  ...
1226
1253
 
1254
+ @typing.overload
1255
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1256
+ """
1257
+ Specifies the PyPI packages for all steps of the flow.
1258
+
1259
+ Use `@pypi_base` to set common packages required by all
1260
+ steps and use `@pypi` to specify step-specific overrides.
1261
+
1262
+ Parameters
1263
+ ----------
1264
+ packages : Dict[str, str], default: {}
1265
+ Packages to use for this flow. The key is the name of the package
1266
+ and the value is the version to use.
1267
+ python : str, optional, default: None
1268
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1269
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1270
+ """
1271
+ ...
1272
+
1273
+ @typing.overload
1274
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1275
+ ...
1276
+
1277
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1278
+ """
1279
+ Specifies the PyPI packages for all steps of the flow.
1280
+
1281
+ Use `@pypi_base` to set common packages required by all
1282
+ steps and use `@pypi` to specify step-specific overrides.
1283
+
1284
+ Parameters
1285
+ ----------
1286
+ packages : Dict[str, str], default: {}
1287
+ Packages to use for this flow. The key is the name of the package
1288
+ and the value is the version to use.
1289
+ python : str, optional, default: None
1290
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1291
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1292
+ """
1293
+ ...
1294
+
1227
1295
  def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1228
1296
  """
1229
1297
  Specifies what flows belong to the same project.
@@ -1259,155 +1327,89 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1259
1327
  """
1260
1328
  ...
1261
1329
 
1262
- @typing.overload
1263
- 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]]:
1264
- """
1265
- Specifies the times when the flow should be run when running on a
1266
- production scheduler.
1267
-
1268
-
1269
- Parameters
1270
- ----------
1271
- hourly : bool, default False
1272
- Run the workflow hourly.
1273
- daily : bool, default True
1274
- Run the workflow daily.
1275
- weekly : bool, default False
1276
- Run the workflow weekly.
1277
- cron : str, optional, default None
1278
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1279
- specified by this expression.
1280
- timezone : str, optional, default None
1281
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1282
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1283
- """
1284
- ...
1285
-
1286
- @typing.overload
1287
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1288
- ...
1289
-
1290
- 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):
1291
- """
1292
- Specifies the times when the flow should be run when running on a
1293
- production scheduler.
1294
-
1295
-
1296
- Parameters
1297
- ----------
1298
- hourly : bool, default False
1299
- Run the workflow hourly.
1300
- daily : bool, default True
1301
- Run the workflow daily.
1302
- weekly : bool, default False
1303
- Run the workflow weekly.
1304
- cron : str, optional, default None
1305
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1306
- specified by this expression.
1307
- timezone : str, optional, default None
1308
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1309
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1310
- """
1311
- ...
1312
-
1313
- @typing.overload
1314
- 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]]:
1330
+ 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]]:
1315
1331
  """
1316
- Specifies the flow(s) that this flow depends on.
1317
-
1318
- ```
1319
- @trigger_on_finish(flow='FooFlow')
1320
- ```
1321
- or
1322
- ```
1323
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1324
- ```
1325
- This decorator respects the @project decorator and triggers the flow
1326
- when upstream runs within the same namespace complete successfully
1327
-
1328
- Additionally, you can specify project aware upstream flow dependencies
1329
- by specifying the fully qualified project_flow_name.
1330
- ```
1331
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1332
- ```
1333
- or
1334
- ```
1335
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1336
- ```
1337
-
1338
- You can also specify just the project or project branch (other values will be
1339
- inferred from the current project or project branch):
1340
- ```
1341
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1342
- ```
1343
-
1344
- Note that `branch` is typically one of:
1345
- - `prod`
1346
- - `user.bob`
1347
- - `test.my_experiment`
1348
- - `prod.staging`
1332
+ 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)
1333
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1334
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1335
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1336
+ starts only after all sensors finish.
1349
1337
 
1350
1338
 
1351
1339
  Parameters
1352
1340
  ----------
1353
- flow : Union[str, Dict[str, str]], optional, default None
1354
- Upstream flow dependency for this flow.
1355
- flows : List[Union[str, Dict[str, str]]], default []
1356
- Upstream flow dependencies for this flow.
1357
- options : Dict[str, Any], default {}
1358
- Backend-specific configuration for tuning eventing behavior.
1341
+ timeout : int
1342
+ Time, in seconds before the task times out and fails. (Default: 3600)
1343
+ poke_interval : int
1344
+ Time in seconds that the job should wait in between each try. (Default: 60)
1345
+ mode : str
1346
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1347
+ exponential_backoff : bool
1348
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1349
+ pool : str
1350
+ the slot pool this task should run in,
1351
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1352
+ soft_fail : bool
1353
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1354
+ name : str
1355
+ Name of the sensor on Airflow
1356
+ description : str
1357
+ Description of sensor in the Airflow UI
1358
+ bucket_key : Union[str, List[str]]
1359
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1360
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1361
+ bucket_name : str
1362
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1363
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1364
+ wildcard_match : bool
1365
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1366
+ aws_conn_id : str
1367
+ a reference to the s3 connection on Airflow. (Default: None)
1368
+ verify : bool
1369
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1359
1370
  """
1360
1371
  ...
1361
1372
 
1362
- @typing.overload
1363
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1364
- ...
1365
-
1366
- 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] = {}):
1373
+ 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]]:
1367
1374
  """
1368
- Specifies the flow(s) that this flow depends on.
1369
-
1370
- ```
1371
- @trigger_on_finish(flow='FooFlow')
1372
- ```
1373
- or
1374
- ```
1375
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1376
- ```
1377
- This decorator respects the @project decorator and triggers the flow
1378
- when upstream runs within the same namespace complete successfully
1379
-
1380
- Additionally, you can specify project aware upstream flow dependencies
1381
- by specifying the fully qualified project_flow_name.
1382
- ```
1383
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1384
- ```
1385
- or
1386
- ```
1387
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1388
- ```
1389
-
1390
- You can also specify just the project or project branch (other values will be
1391
- inferred from the current project or project branch):
1392
- ```
1393
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1394
- ```
1395
-
1396
- Note that `branch` is typically one of:
1397
- - `prod`
1398
- - `user.bob`
1399
- - `test.my_experiment`
1400
- - `prod.staging`
1375
+ 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.
1376
+ 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.
1401
1377
 
1402
1378
 
1403
1379
  Parameters
1404
1380
  ----------
1405
- flow : Union[str, Dict[str, str]], optional, default None
1406
- Upstream flow dependency for this flow.
1407
- flows : List[Union[str, Dict[str, str]]], default []
1408
- Upstream flow dependencies for this flow.
1409
- options : Dict[str, Any], default {}
1410
- Backend-specific configuration for tuning eventing behavior.
1381
+ timeout : int
1382
+ Time, in seconds before the task times out and fails. (Default: 3600)
1383
+ poke_interval : int
1384
+ Time in seconds that the job should wait in between each try. (Default: 60)
1385
+ mode : str
1386
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1387
+ exponential_backoff : bool
1388
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1389
+ pool : str
1390
+ the slot pool this task should run in,
1391
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1392
+ soft_fail : bool
1393
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1394
+ name : str
1395
+ Name of the sensor on Airflow
1396
+ description : str
1397
+ Description of sensor in the Airflow UI
1398
+ external_dag_id : str
1399
+ The dag_id that contains the task you want to wait for.
1400
+ external_task_ids : List[str]
1401
+ The list of task_ids that you want to wait for.
1402
+ If None (default value) the sensor waits for the DAG. (Default: None)
1403
+ allowed_states : List[str]
1404
+ Iterable of allowed states, (Default: ['success'])
1405
+ failed_states : List[str]
1406
+ Iterable of failed or dis-allowed states. (Default: None)
1407
+ execution_delta : datetime.timedelta
1408
+ time difference with the previous execution to look at,
1409
+ the default is the same logical date as the current task or DAG. (Default: None)
1410
+ check_existence: bool
1411
+ Set to True to check if the external task exists or check if
1412
+ the DAG to wait for exists. (Default: True)
1411
1413
  """
1412
1414
  ...
1413
1415