ob-metaflow-stubs 6.0.3.179rc5__py2.py3-none-any.whl → 6.0.3.180__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (220) hide show
  1. metaflow-stubs/__init__.pyi +645 -608
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +2 -2
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +4 -4
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +89 -89
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +1 -1
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +1 -1
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +1 -1
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +1 -1
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +1 -1
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +1 -1
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -2
  88. metaflow-stubs/multicore_utils.pyi +1 -1
  89. metaflow-stubs/parameters.pyi +4 -4
  90. metaflow-stubs/plugins/__init__.pyi +12 -12
  91. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  92. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  93. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  94. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  96. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  98. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  99. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  100. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  101. metaflow-stubs/plugins/argo/argo_workflows.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +1 -1
  103. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  104. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +1 -1
  105. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  106. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  107. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  108. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  109. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  110. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  111. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  112. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  113. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  114. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  115. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  116. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  117. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  118. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  119. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  120. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  121. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  122. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  123. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  124. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  125. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  126. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  127. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  128. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  129. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  130. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  131. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  133. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  135. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  137. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  140. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  141. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  142. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  143. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  144. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  146. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  148. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  149. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  150. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  151. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  152. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  153. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  154. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  155. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  156. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  157. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  158. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  159. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  160. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  161. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  162. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  163. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  164. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  165. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  169. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  170. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  171. metaflow-stubs/plugins/perimeters.pyi +1 -1
  172. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  173. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  174. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  175. metaflow-stubs/plugins/pypi/conda_environment.pyi +2 -2
  176. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  177. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  179. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  180. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  181. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  182. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  183. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  184. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  185. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  186. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  187. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  188. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  189. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  190. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/uv/uv_environment.pyi +1 -1
  192. metaflow-stubs/profilers/__init__.pyi +1 -1
  193. metaflow-stubs/pylint_wrapper.pyi +1 -1
  194. metaflow-stubs/runner/__init__.pyi +1 -1
  195. metaflow-stubs/runner/deployer.pyi +30 -30
  196. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  197. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  198. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  199. metaflow-stubs/runner/nbrun.pyi +1 -1
  200. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  201. metaflow-stubs/runner/utils.pyi +2 -2
  202. metaflow-stubs/system/__init__.pyi +1 -1
  203. metaflow-stubs/system/system_logger.pyi +1 -1
  204. metaflow-stubs/system/system_monitor.pyi +1 -1
  205. metaflow-stubs/tagging_util.pyi +1 -1
  206. metaflow-stubs/tuple_util.pyi +1 -1
  207. metaflow-stubs/user_configs/__init__.pyi +1 -1
  208. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  209. metaflow-stubs/user_configs/config_options.pyi +2 -2
  210. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  211. {ob_metaflow_stubs-6.0.3.179rc5.dist-info → ob_metaflow_stubs-6.0.3.180.dist-info}/METADATA +1 -1
  212. ob_metaflow_stubs-6.0.3.180.dist-info/RECORD +215 -0
  213. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +0 -6
  214. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +0 -51
  215. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +0 -65
  216. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +0 -74
  217. metaflow-stubs/ob_internal.pyi +0 -11
  218. ob_metaflow_stubs-6.0.3.179rc5.dist-info/RECORD +0 -220
  219. {ob_metaflow_stubs-6.0.3.179rc5.dist-info → ob_metaflow_stubs-6.0.3.180.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.179rc5.dist-info → ob_metaflow_stubs-6.0.3.180.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.17.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-13T21:29:38.839071 #
4
+ # Generated on 2025-06-17T23:27:53.375440 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,8 +35,8 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import tuple_util as tuple_util
39
38
  from . import cards as cards
39
+ from . import tuple_util as tuple_util
40
40
  from . import metaflow_git as metaflow_git
41
41
  from . import events as events
42
42
  from . import runner as runner
@@ -78,7 +78,6 @@ from . import system as system
78
78
  from . import pylint_wrapper as pylint_wrapper
79
79
  from . import cli as cli
80
80
  from . import profilers as profilers
81
- from . import ob_internal as ob_internal
82
81
 
83
82
  EXT_PKG: str
84
83
 
@@ -156,143 +155,21 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
156
155
  ...
157
156
 
158
157
  @typing.overload
159
- 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]]]:
160
- """
161
- Specifies environment variables to be set prior to the execution of a step.
162
-
163
-
164
- Parameters
165
- ----------
166
- vars : Dict[str, str], default {}
167
- Dictionary of environment variables to set.
168
- """
169
- ...
170
-
171
- @typing.overload
172
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
173
- ...
174
-
175
- @typing.overload
176
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
177
- ...
178
-
179
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
180
- """
181
- Specifies environment variables to be set prior to the execution of a step.
182
-
183
-
184
- Parameters
185
- ----------
186
- vars : Dict[str, str], default {}
187
- Dictionary of environment variables to set.
188
- """
189
- ...
190
-
191
- @typing.overload
192
- 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]]]:
193
- """
194
- Enables checkpointing for a step.
195
-
196
-
197
-
198
- Parameters
199
- ----------
200
- load_policy : str, default: "fresh"
201
- The policy for loading the checkpoint. The following policies are supported:
202
- - "eager": Loads the the latest available checkpoint within the namespace.
203
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
204
- will be loaded at the start of the task.
205
- - "none": Do not load any checkpoint
206
- - "fresh": Loads the lastest checkpoint created within the running Task.
207
- This mode helps loading checkpoints across various retry attempts of the same task.
208
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
209
- created within the task will be loaded when the task is retries execution on failure.
210
-
211
- temp_dir_root : str, default: None
212
- The root directory under which `current.checkpoint.directory` will be created.
213
- """
214
- ...
215
-
216
- @typing.overload
217
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
218
- ...
219
-
220
- @typing.overload
221
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
222
- ...
223
-
224
- 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):
225
- """
226
- Enables checkpointing for a step.
227
-
228
-
229
-
230
- Parameters
231
- ----------
232
- load_policy : str, default: "fresh"
233
- The policy for loading the checkpoint. The following policies are supported:
234
- - "eager": Loads the the latest available checkpoint within the namespace.
235
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
236
- will be loaded at the start of the task.
237
- - "none": Do not load any checkpoint
238
- - "fresh": Loads the lastest checkpoint created within the running Task.
239
- This mode helps loading checkpoints across various retry attempts of the same task.
240
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
241
- created within the task will be loaded when the task is retries execution on failure.
242
-
243
- temp_dir_root : str, default: None
244
- The root directory under which `current.checkpoint.directory` will be created.
245
- """
246
- ...
247
-
248
- @typing.overload
249
- 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]]]:
158
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
250
159
  """
251
- Specifies that the step will success under all circumstances.
252
-
253
- The decorator will create an optional artifact, specified by `var`, which
254
- contains the exception raised. You can use it to detect the presence
255
- of errors, indicating that all happy-path artifacts produced by the step
256
- are missing.
257
-
258
-
259
- Parameters
260
- ----------
261
- var : str, optional, default None
262
- Name of the artifact in which to store the caught exception.
263
- If not specified, the exception is not stored.
264
- print_exception : bool, default True
265
- Determines whether or not the exception is printed to
266
- stdout when caught.
160
+ Decorator prototype for all step decorators. This function gets specialized
161
+ and imported for all decorators types by _import_plugin_decorators().
267
162
  """
268
163
  ...
269
164
 
270
165
  @typing.overload
271
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
272
- ...
273
-
274
- @typing.overload
275
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
166
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
276
167
  ...
277
168
 
278
- 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):
169
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
279
170
  """
280
- Specifies that the step will success under all circumstances.
281
-
282
- The decorator will create an optional artifact, specified by `var`, which
283
- contains the exception raised. You can use it to detect the presence
284
- of errors, indicating that all happy-path artifacts produced by the step
285
- are missing.
286
-
287
-
288
- Parameters
289
- ----------
290
- var : str, optional, default None
291
- Name of the artifact in which to store the caught exception.
292
- If not specified, the exception is not stored.
293
- print_exception : bool, default True
294
- Determines whether or not the exception is printed to
295
- stdout when caught.
171
+ Decorator prototype for all step decorators. This function gets specialized
172
+ and imported for all decorators types by _import_plugin_decorators().
296
173
  """
297
174
  ...
298
175
 
@@ -385,6 +262,104 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
385
262
  """
386
263
  ...
387
264
 
265
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
266
+ """
267
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
268
+
269
+ User code call
270
+ --------------
271
+ @ollama(
272
+ models=[...],
273
+ ...
274
+ )
275
+
276
+ Valid backend options
277
+ ---------------------
278
+ - 'local': Run as a separate process on the local task machine.
279
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
280
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
281
+
282
+ Valid model options
283
+ -------------------
284
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
285
+
286
+
287
+ Parameters
288
+ ----------
289
+ models: list[str]
290
+ List of Ollama containers running models in sidecars.
291
+ backend: str
292
+ Determines where and how to run the Ollama process.
293
+ force_pull: bool
294
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
295
+ cache_update_policy: str
296
+ Cache update policy: "auto", "force", or "never".
297
+ force_cache_update: bool
298
+ Simple override for "force" cache update policy.
299
+ debug: bool
300
+ Whether to turn on verbose debugging logs.
301
+ circuit_breaker_config: dict
302
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
303
+ timeout_config: dict
304
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
305
+ """
306
+ ...
307
+
308
+ @typing.overload
309
+ 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]]]:
310
+ """
311
+ Specifies the number of times the task corresponding
312
+ to a step needs to be retried.
313
+
314
+ This decorator is useful for handling transient errors, such as networking issues.
315
+ If your task contains operations that can't be retried safely, e.g. database updates,
316
+ it is advisable to annotate it with `@retry(times=0)`.
317
+
318
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
319
+ decorator will execute a no-op task after all retries have been exhausted,
320
+ ensuring that the flow execution can continue.
321
+
322
+
323
+ Parameters
324
+ ----------
325
+ times : int, default 3
326
+ Number of times to retry this task.
327
+ minutes_between_retries : int, default 2
328
+ Number of minutes between retries.
329
+ """
330
+ ...
331
+
332
+ @typing.overload
333
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
334
+ ...
335
+
336
+ @typing.overload
337
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
338
+ ...
339
+
340
+ 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):
341
+ """
342
+ Specifies the number of times the task corresponding
343
+ to a step needs to be retried.
344
+
345
+ This decorator is useful for handling transient errors, such as networking issues.
346
+ If your task contains operations that can't be retried safely, e.g. database updates,
347
+ it is advisable to annotate it with `@retry(times=0)`.
348
+
349
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
350
+ decorator will execute a no-op task after all retries have been exhausted,
351
+ ensuring that the flow execution can continue.
352
+
353
+
354
+ Parameters
355
+ ----------
356
+ times : int, default 3
357
+ Number of times to retry this task.
358
+ minutes_between_retries : int, default 2
359
+ Number of minutes between retries.
360
+ """
361
+ ...
362
+
388
363
  @typing.overload
389
364
  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]]]:
390
365
  """
@@ -464,36 +439,81 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
464
439
  """
465
440
  ...
466
441
 
467
- @typing.overload
468
- 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]]]:
442
+ def vllm(*, model: str, backend: str, debug: bool, kwargs: typing.Any) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
469
443
  """
470
- Creates a human-readable report, a Metaflow Card, after this step completes.
444
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
471
445
 
472
- Note that you may add multiple `@card` decorators in a step with different parameters.
446
+ User code call
447
+ --------------
448
+ @vllm(
449
+ model="...",
450
+ ...
451
+ )
452
+
453
+ Valid backend options
454
+ ---------------------
455
+ - 'local': Run as a separate process on the local task machine.
456
+
457
+ Valid model options
458
+ -------------------
459
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
460
+
461
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
462
+ If you need multiple models, you must create multiple @vllm decorators.
473
463
 
474
464
 
475
465
  Parameters
476
466
  ----------
477
- type : str, default 'default'
478
- Card type.
479
- id : str, optional, default None
480
- If multiple cards are present, use this id to identify this card.
481
- options : Dict[str, Any], default {}
482
- Options passed to the card. The contents depend on the card type.
483
- timeout : int, default 45
484
- Interrupt reporting if it takes more than this many seconds.
467
+ model: str
468
+ HuggingFace model identifier to be served by vLLM.
469
+ backend: str
470
+ Determines where and how to run the vLLM process.
471
+ debug: bool
472
+ Whether to turn on verbose debugging logs.
473
+ kwargs : Any
474
+ Any other keyword arguments are passed directly to the vLLM engine.
475
+ This allows for flexible configuration of vLLM server settings.
476
+ For example, `tensor_parallel_size=2`.
485
477
  """
486
478
  ...
487
479
 
488
480
  @typing.overload
489
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
481
+ 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]]]:
482
+ """
483
+ Specifies secrets to be retrieved and injected as environment variables prior to
484
+ the execution of a step.
485
+
486
+
487
+ Parameters
488
+ ----------
489
+ sources : List[Union[str, Dict[str, Any]]], default: []
490
+ List of secret specs, defining how the secrets are to be retrieved
491
+ """
490
492
  ...
491
493
 
492
494
  @typing.overload
493
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
495
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
494
496
  ...
495
497
 
496
- 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):
498
+ @typing.overload
499
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
500
+ ...
501
+
502
+ 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]]] = []):
503
+ """
504
+ Specifies secrets to be retrieved and injected as environment variables prior to
505
+ the execution of a step.
506
+
507
+
508
+ Parameters
509
+ ----------
510
+ sources : List[Union[str, Dict[str, Any]]], default: []
511
+ List of secret specs, defining how the secrets are to be retrieved
512
+ """
513
+ ...
514
+
515
+ @typing.overload
516
+ 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]]]:
497
517
  """
498
518
  Creates a human-readable report, a Metaflow Card, after this step completes.
499
519
 
@@ -513,156 +533,139 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
513
533
  """
514
534
  ...
515
535
 
516
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
517
- """
518
- Specifies that this step is used to deploy an instance of the app.
519
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
520
-
521
-
522
- Parameters
523
- ----------
524
- app_port : int
525
- Number of GPUs to use.
526
- app_name : str
527
- Name of the app to deploy.
528
- """
536
+ @typing.overload
537
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
529
538
  ...
530
539
 
531
- 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]]]:
540
+ @typing.overload
541
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
542
+ ...
543
+
544
+ 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):
532
545
  """
533
- Decorator that helps cache, version and store models/datasets from huggingface hub.
546
+ Creates a human-readable report, a Metaflow Card, after this step completes.
547
+
548
+ Note that you may add multiple `@card` decorators in a step with different parameters.
534
549
 
535
550
 
536
551
  Parameters
537
552
  ----------
538
- temp_dir_root : str, optional
539
- The root directory that will hold the temporary directory where objects will be downloaded.
540
-
541
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
542
- The list of repos (models/datasets) to load.
543
-
544
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
545
-
546
- - If repo (model/dataset) is not found in the datastore:
547
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
548
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
549
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
550
-
551
- - If repo is found in the datastore:
552
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
553
+ type : str, default 'default'
554
+ Card type.
555
+ id : str, optional, default None
556
+ If multiple cards are present, use this id to identify this card.
557
+ options : Dict[str, Any], default {}
558
+ Options passed to the card. The contents depend on the card type.
559
+ timeout : int, default 45
560
+ Interrupt reporting if it takes more than this many seconds.
553
561
  """
554
562
  ...
555
563
 
556
564
  @typing.overload
557
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
565
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
558
566
  """
559
- Decorator prototype for all step decorators. This function gets specialized
560
- and imported for all decorators types by _import_plugin_decorators().
567
+ Internal decorator to support Fast bakery
561
568
  """
562
569
  ...
563
570
 
564
571
  @typing.overload
565
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
572
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
566
573
  ...
567
574
 
568
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
575
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
569
576
  """
570
- Decorator prototype for all step decorators. This function gets specialized
571
- and imported for all decorators types by _import_plugin_decorators().
577
+ Internal decorator to support Fast bakery
572
578
  """
573
579
  ...
574
580
 
575
581
  @typing.overload
576
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
582
+ 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]]]:
577
583
  """
578
- Specifies the PyPI packages for the step.
584
+ Specifies that the step will success under all circumstances.
579
585
 
580
- Information in this decorator will augment any
581
- attributes set in the `@pyi_base` flow-level decorator. Hence,
582
- you can use `@pypi_base` to set packages required by all
583
- steps and use `@pypi` to specify step-specific overrides.
586
+ The decorator will create an optional artifact, specified by `var`, which
587
+ contains the exception raised. You can use it to detect the presence
588
+ of errors, indicating that all happy-path artifacts produced by the step
589
+ are missing.
584
590
 
585
591
 
586
592
  Parameters
587
593
  ----------
588
- packages : Dict[str, str], default: {}
589
- Packages to use for this step. The key is the name of the package
590
- and the value is the version to use.
591
- python : str, optional, default: None
592
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
593
- that the version used will correspond to the version of the Python interpreter used to start the run.
594
+ var : str, optional, default None
595
+ Name of the artifact in which to store the caught exception.
596
+ If not specified, the exception is not stored.
597
+ print_exception : bool, default True
598
+ Determines whether or not the exception is printed to
599
+ stdout when caught.
594
600
  """
595
601
  ...
596
602
 
597
603
  @typing.overload
598
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
604
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
599
605
  ...
600
606
 
601
607
  @typing.overload
602
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
608
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
603
609
  ...
604
610
 
605
- 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):
611
+ 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):
606
612
  """
607
- Specifies the PyPI packages for the step.
613
+ Specifies that the step will success under all circumstances.
608
614
 
609
- Information in this decorator will augment any
610
- attributes set in the `@pyi_base` flow-level decorator. Hence,
611
- you can use `@pypi_base` to set packages required by all
612
- steps and use `@pypi` to specify step-specific overrides.
615
+ The decorator will create an optional artifact, specified by `var`, which
616
+ contains the exception raised. You can use it to detect the presence
617
+ of errors, indicating that all happy-path artifacts produced by the step
618
+ are missing.
613
619
 
614
620
 
615
621
  Parameters
616
622
  ----------
617
- packages : Dict[str, str], default: {}
618
- Packages to use for this step. The key is the name of the package
619
- and the value is the version to use.
620
- python : str, optional, default: None
621
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
622
- that the version used will correspond to the version of the Python interpreter used to start the run.
623
+ var : str, optional, default None
624
+ Name of the artifact in which to store the caught exception.
625
+ If not specified, the exception is not stored.
626
+ print_exception : bool, default True
627
+ Determines whether or not the exception is printed to
628
+ stdout when caught.
623
629
  """
624
630
  ...
625
631
 
626
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
632
+ 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]]]:
627
633
  """
628
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
634
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
629
635
 
630
- User code call
631
- --------------
632
- @ollama(
633
- models=[...],
634
- ...
635
- )
636
636
 
637
- Valid backend options
638
- ---------------------
639
- - 'local': Run as a separate process on the local task machine.
640
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
641
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
637
+ Parameters
638
+ ----------
639
+ temp_dir_root : str, optional
640
+ The root directory that will hold the temporary directory where objects will be downloaded.
642
641
 
643
- Valid model options
644
- -------------------
645
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
642
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
643
+ The list of repos (models/datasets) to load.
644
+
645
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
646
+
647
+ - If repo (model/dataset) is not found in the datastore:
648
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
649
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
650
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
651
+
652
+ - If repo is found in the datastore:
653
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
654
+ """
655
+ ...
656
+
657
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
658
+ """
659
+ Specifies that this step is used to deploy an instance of the app.
660
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
646
661
 
647
662
 
648
663
  Parameters
649
664
  ----------
650
- models: list[str]
651
- List of Ollama containers running models in sidecars.
652
- backend: str
653
- Determines where and how to run the Ollama process.
654
- force_pull: bool
655
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
656
- cache_update_policy: str
657
- Cache update policy: "auto", "force", or "never".
658
- force_cache_update: bool
659
- Simple override for "force" cache update policy.
660
- debug: bool
661
- Whether to turn on verbose debugging logs.
662
- circuit_breaker_config: dict
663
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
664
- timeout_config: dict
665
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
665
+ app_port : int
666
+ Number of GPUs to use.
667
+ app_name : str
668
+ Name of the app to deploy.
666
669
  """
667
670
  ...
668
671
 
@@ -726,149 +729,67 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
726
729
  ...
727
730
 
728
731
  @typing.overload
729
- 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]]]:
732
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
730
733
  """
731
- Specifies the number of times the task corresponding
732
- to a step needs to be retried.
733
-
734
- This decorator is useful for handling transient errors, such as networking issues.
735
- If your task contains operations that can't be retried safely, e.g. database updates,
736
- it is advisable to annotate it with `@retry(times=0)`.
734
+ Specifies the PyPI packages for the step.
737
735
 
738
- This can be used in conjunction with the `@catch` decorator. The `@catch`
739
- decorator will execute a no-op task after all retries have been exhausted,
740
- ensuring that the flow execution can continue.
736
+ Information in this decorator will augment any
737
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
738
+ you can use `@pypi_base` to set packages required by all
739
+ steps and use `@pypi` to specify step-specific overrides.
741
740
 
742
741
 
743
742
  Parameters
744
743
  ----------
745
- times : int, default 3
746
- Number of times to retry this task.
747
- minutes_between_retries : int, default 2
748
- Number of minutes between retries.
744
+ packages : Dict[str, str], default: {}
745
+ Packages to use for this step. The key is the name of the package
746
+ and the value is the version to use.
747
+ python : str, optional, default: None
748
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
749
+ that the version used will correspond to the version of the Python interpreter used to start the run.
749
750
  """
750
751
  ...
751
752
 
752
753
  @typing.overload
753
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
754
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
754
755
  ...
755
756
 
756
757
  @typing.overload
757
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
758
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
758
759
  ...
759
760
 
760
- 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):
761
+ 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):
761
762
  """
762
- Specifies the number of times the task corresponding
763
- to a step needs to be retried.
764
-
765
- This decorator is useful for handling transient errors, such as networking issues.
766
- If your task contains operations that can't be retried safely, e.g. database updates,
767
- it is advisable to annotate it with `@retry(times=0)`.
763
+ Specifies the PyPI packages for the step.
768
764
 
769
- This can be used in conjunction with the `@catch` decorator. The `@catch`
770
- decorator will execute a no-op task after all retries have been exhausted,
771
- ensuring that the flow execution can continue.
772
-
773
-
774
- Parameters
775
- ----------
776
- times : int, default 3
777
- Number of times to retry this task.
778
- minutes_between_retries : int, default 2
779
- Number of minutes between retries.
780
- """
781
- ...
782
-
783
- @typing.overload
784
- 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]]]:
785
- """
786
- Enables loading / saving of models within a step.
787
-
788
-
789
-
790
- Parameters
791
- ----------
792
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
793
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
794
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
795
- - `current.checkpoint`
796
- - `current.model`
797
- - `current.huggingface_hub`
798
-
799
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
800
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
801
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
802
-
803
- temp_dir_root : str, default: None
804
- The root directory under which `current.model.loaded` will store loaded models
805
- """
806
- ...
807
-
808
- @typing.overload
809
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
810
- ...
811
-
812
- @typing.overload
813
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
814
- ...
815
-
816
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
817
- """
818
- Enables loading / saving of models within a step.
819
-
820
-
821
-
822
- Parameters
823
- ----------
824
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
825
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
826
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
827
- - `current.checkpoint`
828
- - `current.model`
829
- - `current.huggingface_hub`
830
-
831
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
832
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
833
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
834
-
835
- temp_dir_root : str, default: None
836
- The root directory under which `current.model.loaded` will store loaded models
837
- """
838
- ...
839
-
840
- @typing.overload
841
- 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]]]:
842
- """
843
- Specifies secrets to be retrieved and injected as environment variables prior to
844
- the execution of a step.
765
+ Information in this decorator will augment any
766
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
767
+ you can use `@pypi_base` to set packages required by all
768
+ steps and use `@pypi` to specify step-specific overrides.
845
769
 
846
770
 
847
771
  Parameters
848
772
  ----------
849
- sources : List[Union[str, Dict[str, Any]]], default: []
850
- List of secret specs, defining how the secrets are to be retrieved
773
+ packages : Dict[str, str], default: {}
774
+ Packages to use for this step. The key is the name of the package
775
+ and the value is the version to use.
776
+ python : str, optional, default: None
777
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
778
+ that the version used will correspond to the version of the Python interpreter used to start the run.
851
779
  """
852
780
  ...
853
781
 
854
- @typing.overload
855
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
856
- ...
857
-
858
- @typing.overload
859
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
860
- ...
861
-
862
- 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]]] = []):
782
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
863
783
  """
864
- Specifies secrets to be retrieved and injected as environment variables prior to
865
- the execution of a step.
784
+ Specifies that this step should execute on DGX cloud.
866
785
 
867
786
 
868
787
  Parameters
869
788
  ----------
870
- sources : List[Union[str, Dict[str, Any]]], default: []
871
- List of secret specs, defining how the secrets are to be retrieved
789
+ gpu : int
790
+ Number of GPUs to use.
791
+ gpu_type : str
792
+ Type of Nvidia GPU to use.
872
793
  """
873
794
  ...
874
795
 
@@ -931,34 +852,93 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
931
852
  """
932
853
  ...
933
854
 
934
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
855
+ @typing.overload
856
+ 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]]]:
935
857
  """
936
- Specifies that this step should execute on DGX cloud.
858
+ Specifies environment variables to be set prior to the execution of a step.
937
859
 
938
860
 
939
861
  Parameters
940
862
  ----------
941
- gpu : int
942
- Number of GPUs to use.
943
- gpu_type : str
944
- Type of Nvidia GPU to use.
863
+ vars : Dict[str, str], default {}
864
+ Dictionary of environment variables to set.
945
865
  """
946
866
  ...
947
867
 
948
868
  @typing.overload
949
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
869
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
870
+ ...
871
+
872
+ @typing.overload
873
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
874
+ ...
875
+
876
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
950
877
  """
951
- Internal decorator to support Fast bakery
878
+ Specifies environment variables to be set prior to the execution of a step.
879
+
880
+
881
+ Parameters
882
+ ----------
883
+ vars : Dict[str, str], default {}
884
+ Dictionary of environment variables to set.
952
885
  """
953
886
  ...
954
887
 
955
888
  @typing.overload
956
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
889
+ 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]]]:
890
+ """
891
+ Enables loading / saving of models within a step.
892
+
893
+
894
+
895
+ Parameters
896
+ ----------
897
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
898
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
899
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
900
+ - `current.checkpoint`
901
+ - `current.model`
902
+ - `current.huggingface_hub`
903
+
904
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
905
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
906
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
907
+
908
+ temp_dir_root : str, default: None
909
+ The root directory under which `current.model.loaded` will store loaded models
910
+ """
957
911
  ...
958
912
 
959
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
913
+ @typing.overload
914
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
915
+ ...
916
+
917
+ @typing.overload
918
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
919
+ ...
920
+
921
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
960
922
  """
961
- Internal decorator to support Fast bakery
923
+ Enables loading / saving of models within a step.
924
+
925
+
926
+
927
+ Parameters
928
+ ----------
929
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
930
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
931
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
932
+ - `current.checkpoint`
933
+ - `current.model`
934
+ - `current.huggingface_hub`
935
+
936
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
937
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
938
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
939
+
940
+ temp_dir_root : str, default: None
941
+ The root directory under which `current.model.loaded` will store loaded models
962
942
  """
963
943
  ...
964
944
 
@@ -978,241 +958,101 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
978
958
  """
979
959
  ...
980
960
 
981
- 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]]:
961
+ @typing.overload
962
+ 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]]]:
982
963
  """
983
- 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.
984
- 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.
964
+ Enables checkpointing for a step.
965
+
985
966
 
986
967
 
987
968
  Parameters
988
969
  ----------
989
- timeout : int
990
- Time, in seconds before the task times out and fails. (Default: 3600)
991
- poke_interval : int
992
- Time in seconds that the job should wait in between each try. (Default: 60)
993
- mode : str
994
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
995
- exponential_backoff : bool
996
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
997
- pool : str
998
- the slot pool this task should run in,
999
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1000
- soft_fail : bool
1001
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1002
- name : str
1003
- Name of the sensor on Airflow
1004
- description : str
1005
- Description of sensor in the Airflow UI
1006
- external_dag_id : str
1007
- The dag_id that contains the task you want to wait for.
1008
- external_task_ids : List[str]
1009
- The list of task_ids that you want to wait for.
1010
- If None (default value) the sensor waits for the DAG. (Default: None)
1011
- allowed_states : List[str]
1012
- Iterable of allowed states, (Default: ['success'])
1013
- failed_states : List[str]
1014
- Iterable of failed or dis-allowed states. (Default: None)
1015
- execution_delta : datetime.timedelta
1016
- time difference with the previous execution to look at,
1017
- the default is the same logical date as the current task or DAG. (Default: None)
1018
- check_existence: bool
1019
- Set to True to check if the external task exists or check if
1020
- the DAG to wait for exists. (Default: True)
970
+ load_policy : str, default: "fresh"
971
+ The policy for loading the checkpoint. The following policies are supported:
972
+ - "eager": Loads the the latest available checkpoint within the namespace.
973
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
974
+ will be loaded at the start of the task.
975
+ - "none": Do not load any checkpoint
976
+ - "fresh": Loads the lastest checkpoint created within the running Task.
977
+ This mode helps loading checkpoints across various retry attempts of the same task.
978
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
979
+ created within the task will be loaded when the task is retries execution on failure.
980
+
981
+ temp_dir_root : str, default: None
982
+ The root directory under which `current.checkpoint.directory` will be created.
1021
983
  """
1022
984
  ...
1023
985
 
1024
986
  @typing.overload
1025
- 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]]:
1026
- """
1027
- Specifies the flow(s) that this flow depends on.
1028
-
1029
- ```
1030
- @trigger_on_finish(flow='FooFlow')
1031
- ```
1032
- or
1033
- ```
1034
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1035
- ```
1036
- This decorator respects the @project decorator and triggers the flow
1037
- when upstream runs within the same namespace complete successfully
1038
-
1039
- Additionally, you can specify project aware upstream flow dependencies
1040
- by specifying the fully qualified project_flow_name.
1041
- ```
1042
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1043
- ```
1044
- or
1045
- ```
1046
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1047
- ```
1048
-
1049
- You can also specify just the project or project branch (other values will be
1050
- inferred from the current project or project branch):
1051
- ```
1052
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1053
- ```
1054
-
1055
- Note that `branch` is typically one of:
1056
- - `prod`
1057
- - `user.bob`
1058
- - `test.my_experiment`
1059
- - `prod.staging`
1060
-
1061
-
1062
- Parameters
1063
- ----------
1064
- flow : Union[str, Dict[str, str]], optional, default None
1065
- Upstream flow dependency for this flow.
1066
- flows : List[Union[str, Dict[str, str]]], default []
1067
- Upstream flow dependencies for this flow.
1068
- options : Dict[str, Any], default {}
1069
- Backend-specific configuration for tuning eventing behavior.
1070
- """
987
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1071
988
  ...
1072
989
 
1073
990
  @typing.overload
1074
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
991
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1075
992
  ...
1076
993
 
1077
- 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] = {}):
994
+ 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):
1078
995
  """
1079
- Specifies the flow(s) that this flow depends on.
1080
-
1081
- ```
1082
- @trigger_on_finish(flow='FooFlow')
1083
- ```
1084
- or
1085
- ```
1086
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1087
- ```
1088
- This decorator respects the @project decorator and triggers the flow
1089
- when upstream runs within the same namespace complete successfully
1090
-
1091
- Additionally, you can specify project aware upstream flow dependencies
1092
- by specifying the fully qualified project_flow_name.
1093
- ```
1094
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1095
- ```
1096
- or
1097
- ```
1098
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1099
- ```
1100
-
1101
- You can also specify just the project or project branch (other values will be
1102
- inferred from the current project or project branch):
1103
- ```
1104
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1105
- ```
996
+ Enables checkpointing for a step.
1106
997
 
1107
- Note that `branch` is typically one of:
1108
- - `prod`
1109
- - `user.bob`
1110
- - `test.my_experiment`
1111
- - `prod.staging`
1112
998
 
1113
999
 
1114
1000
  Parameters
1115
1001
  ----------
1116
- flow : Union[str, Dict[str, str]], optional, default None
1117
- Upstream flow dependency for this flow.
1118
- flows : List[Union[str, Dict[str, str]]], default []
1119
- Upstream flow dependencies for this flow.
1120
- options : Dict[str, Any], default {}
1121
- Backend-specific configuration for tuning eventing behavior.
1002
+ load_policy : str, default: "fresh"
1003
+ The policy for loading the checkpoint. The following policies are supported:
1004
+ - "eager": Loads the the latest available checkpoint within the namespace.
1005
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1006
+ will be loaded at the start of the task.
1007
+ - "none": Do not load any checkpoint
1008
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1009
+ This mode helps loading checkpoints across various retry attempts of the same task.
1010
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1011
+ created within the task will be loaded when the task is retries execution on failure.
1012
+
1013
+ temp_dir_root : str, default: None
1014
+ The root directory under which `current.checkpoint.directory` will be created.
1122
1015
  """
1123
1016
  ...
1124
1017
 
1125
1018
  @typing.overload
1126
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1019
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1127
1020
  """
1128
- Specifies the Conda environment for all steps of the flow.
1129
-
1130
- Use `@conda_base` to set common libraries required by all
1131
- steps and use `@conda` to specify step-specific additions.
1021
+ Specifies the PyPI packages for all steps of the flow.
1132
1022
 
1023
+ Use `@pypi_base` to set common packages required by all
1024
+ steps and use `@pypi` to specify step-specific overrides.
1133
1025
 
1134
1026
  Parameters
1135
1027
  ----------
1136
- packages : Dict[str, str], default {}
1028
+ packages : Dict[str, str], default: {}
1137
1029
  Packages to use for this flow. The key is the name of the package
1138
1030
  and the value is the version to use.
1139
- libraries : Dict[str, str], default {}
1140
- Supported for backward compatibility. When used with packages, packages will take precedence.
1141
- python : str, optional, default None
1031
+ python : str, optional, default: None
1142
1032
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1143
1033
  that the version used will correspond to the version of the Python interpreter used to start the run.
1144
- disabled : bool, default False
1145
- If set to True, disables Conda.
1146
1034
  """
1147
1035
  ...
1148
1036
 
1149
1037
  @typing.overload
1150
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1038
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1151
1039
  ...
1152
1040
 
1153
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1041
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1154
1042
  """
1155
- Specifies the Conda environment for all steps of the flow.
1156
-
1157
- Use `@conda_base` to set common libraries required by all
1158
- steps and use `@conda` to specify step-specific additions.
1043
+ Specifies the PyPI packages for all steps of the flow.
1159
1044
 
1045
+ Use `@pypi_base` to set common packages required by all
1046
+ steps and use `@pypi` to specify step-specific overrides.
1160
1047
 
1161
1048
  Parameters
1162
1049
  ----------
1163
- packages : Dict[str, str], default {}
1050
+ packages : Dict[str, str], default: {}
1164
1051
  Packages to use for this flow. The key is the name of the package
1165
1052
  and the value is the version to use.
1166
- libraries : Dict[str, str], default {}
1167
- Supported for backward compatibility. When used with packages, packages will take precedence.
1168
- python : str, optional, default None
1053
+ python : str, optional, default: None
1169
1054
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1170
1055
  that the version used will correspond to the version of the Python interpreter used to start the run.
1171
- disabled : bool, default False
1172
- If set to True, disables Conda.
1173
- """
1174
- ...
1175
-
1176
- 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]]:
1177
- """
1178
- 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)
1179
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1180
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1181
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1182
- starts only after all sensors finish.
1183
-
1184
-
1185
- Parameters
1186
- ----------
1187
- timeout : int
1188
- Time, in seconds before the task times out and fails. (Default: 3600)
1189
- poke_interval : int
1190
- Time in seconds that the job should wait in between each try. (Default: 60)
1191
- mode : str
1192
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1193
- exponential_backoff : bool
1194
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1195
- pool : str
1196
- the slot pool this task should run in,
1197
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1198
- soft_fail : bool
1199
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1200
- name : str
1201
- Name of the sensor on Airflow
1202
- description : str
1203
- Description of sensor in the Airflow UI
1204
- bucket_key : Union[str, List[str]]
1205
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1206
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1207
- bucket_name : str
1208
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1209
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1210
- wildcard_match : bool
1211
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1212
- aws_conn_id : str
1213
- a reference to the s3 connection on Airflow. (Default: None)
1214
- verify : bool
1215
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1216
1056
  """
1217
1057
  ...
1218
1058
 
@@ -1252,59 +1092,18 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1252
1092
 
1253
1093
  Parameters
1254
1094
  ----------
1255
- hourly : bool, default False
1256
- Run the workflow hourly.
1257
- daily : bool, default True
1258
- Run the workflow daily.
1259
- weekly : bool, default False
1260
- Run the workflow weekly.
1261
- cron : str, optional, default None
1262
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1263
- specified by this expression.
1264
- timezone : str, optional, default None
1265
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1266
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1267
- """
1268
- ...
1269
-
1270
- @typing.overload
1271
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1272
- """
1273
- Specifies the PyPI packages for all steps of the flow.
1274
-
1275
- Use `@pypi_base` to set common packages required by all
1276
- steps and use `@pypi` to specify step-specific overrides.
1277
-
1278
- Parameters
1279
- ----------
1280
- packages : Dict[str, str], default: {}
1281
- Packages to use for this flow. The key is the name of the package
1282
- and the value is the version to use.
1283
- python : str, optional, default: None
1284
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1285
- that the version used will correspond to the version of the Python interpreter used to start the run.
1286
- """
1287
- ...
1288
-
1289
- @typing.overload
1290
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1291
- ...
1292
-
1293
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1294
- """
1295
- Specifies the PyPI packages for all steps of the flow.
1296
-
1297
- Use `@pypi_base` to set common packages required by all
1298
- steps and use `@pypi` to specify step-specific overrides.
1299
-
1300
- Parameters
1301
- ----------
1302
- packages : Dict[str, str], default: {}
1303
- Packages to use for this flow. The key is the name of the package
1304
- and the value is the version to use.
1305
- python : str, optional, default: None
1306
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1307
- that the version used will correspond to the version of the Python interpreter used to start the run.
1095
+ hourly : bool, default False
1096
+ Run the workflow hourly.
1097
+ daily : bool, default True
1098
+ Run the workflow daily.
1099
+ weekly : bool, default False
1100
+ Run the workflow weekly.
1101
+ cron : str, optional, default None
1102
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1103
+ specified by this expression.
1104
+ timezone : str, optional, default None
1105
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1106
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1308
1107
  """
1309
1108
  ...
1310
1109
 
@@ -1401,6 +1200,150 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1401
1200
  """
1402
1201
  ...
1403
1202
 
1203
+ @typing.overload
1204
+ 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]]:
1205
+ """
1206
+ Specifies the flow(s) that this flow depends on.
1207
+
1208
+ ```
1209
+ @trigger_on_finish(flow='FooFlow')
1210
+ ```
1211
+ or
1212
+ ```
1213
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1214
+ ```
1215
+ This decorator respects the @project decorator and triggers the flow
1216
+ when upstream runs within the same namespace complete successfully
1217
+
1218
+ Additionally, you can specify project aware upstream flow dependencies
1219
+ by specifying the fully qualified project_flow_name.
1220
+ ```
1221
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1222
+ ```
1223
+ or
1224
+ ```
1225
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1226
+ ```
1227
+
1228
+ You can also specify just the project or project branch (other values will be
1229
+ inferred from the current project or project branch):
1230
+ ```
1231
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1232
+ ```
1233
+
1234
+ Note that `branch` is typically one of:
1235
+ - `prod`
1236
+ - `user.bob`
1237
+ - `test.my_experiment`
1238
+ - `prod.staging`
1239
+
1240
+
1241
+ Parameters
1242
+ ----------
1243
+ flow : Union[str, Dict[str, str]], optional, default None
1244
+ Upstream flow dependency for this flow.
1245
+ flows : List[Union[str, Dict[str, str]]], default []
1246
+ Upstream flow dependencies for this flow.
1247
+ options : Dict[str, Any], default {}
1248
+ Backend-specific configuration for tuning eventing behavior.
1249
+ """
1250
+ ...
1251
+
1252
+ @typing.overload
1253
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1254
+ ...
1255
+
1256
+ 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] = {}):
1257
+ """
1258
+ Specifies the flow(s) that this flow depends on.
1259
+
1260
+ ```
1261
+ @trigger_on_finish(flow='FooFlow')
1262
+ ```
1263
+ or
1264
+ ```
1265
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1266
+ ```
1267
+ This decorator respects the @project decorator and triggers the flow
1268
+ when upstream runs within the same namespace complete successfully
1269
+
1270
+ Additionally, you can specify project aware upstream flow dependencies
1271
+ by specifying the fully qualified project_flow_name.
1272
+ ```
1273
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1274
+ ```
1275
+ or
1276
+ ```
1277
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1278
+ ```
1279
+
1280
+ You can also specify just the project or project branch (other values will be
1281
+ inferred from the current project or project branch):
1282
+ ```
1283
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1284
+ ```
1285
+
1286
+ Note that `branch` is typically one of:
1287
+ - `prod`
1288
+ - `user.bob`
1289
+ - `test.my_experiment`
1290
+ - `prod.staging`
1291
+
1292
+
1293
+ Parameters
1294
+ ----------
1295
+ flow : Union[str, Dict[str, str]], optional, default None
1296
+ Upstream flow dependency for this flow.
1297
+ flows : List[Union[str, Dict[str, str]]], default []
1298
+ Upstream flow dependencies for this flow.
1299
+ options : Dict[str, Any], default {}
1300
+ Backend-specific configuration for tuning eventing behavior.
1301
+ """
1302
+ ...
1303
+
1304
+ 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]]:
1305
+ """
1306
+ 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)
1307
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1308
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1309
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1310
+ starts only after all sensors finish.
1311
+
1312
+
1313
+ Parameters
1314
+ ----------
1315
+ timeout : int
1316
+ Time, in seconds before the task times out and fails. (Default: 3600)
1317
+ poke_interval : int
1318
+ Time in seconds that the job should wait in between each try. (Default: 60)
1319
+ mode : str
1320
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1321
+ exponential_backoff : bool
1322
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1323
+ pool : str
1324
+ the slot pool this task should run in,
1325
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1326
+ soft_fail : bool
1327
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1328
+ name : str
1329
+ Name of the sensor on Airflow
1330
+ description : str
1331
+ Description of sensor in the Airflow UI
1332
+ bucket_key : Union[str, List[str]]
1333
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1334
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1335
+ bucket_name : str
1336
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1337
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1338
+ wildcard_match : bool
1339
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1340
+ aws_conn_id : str
1341
+ a reference to the s3 connection on Airflow. (Default: None)
1342
+ verify : bool
1343
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1344
+ """
1345
+ ...
1346
+
1404
1347
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1405
1348
  """
1406
1349
  Allows setting external datastores to save data for the
@@ -1515,6 +1458,100 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1515
1458
  """
1516
1459
  ...
1517
1460
 
1461
+ 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]]:
1462
+ """
1463
+ 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.
1464
+ 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.
1465
+
1466
+
1467
+ Parameters
1468
+ ----------
1469
+ timeout : int
1470
+ Time, in seconds before the task times out and fails. (Default: 3600)
1471
+ poke_interval : int
1472
+ Time in seconds that the job should wait in between each try. (Default: 60)
1473
+ mode : str
1474
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1475
+ exponential_backoff : bool
1476
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1477
+ pool : str
1478
+ the slot pool this task should run in,
1479
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1480
+ soft_fail : bool
1481
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1482
+ name : str
1483
+ Name of the sensor on Airflow
1484
+ description : str
1485
+ Description of sensor in the Airflow UI
1486
+ external_dag_id : str
1487
+ The dag_id that contains the task you want to wait for.
1488
+ external_task_ids : List[str]
1489
+ The list of task_ids that you want to wait for.
1490
+ If None (default value) the sensor waits for the DAG. (Default: None)
1491
+ allowed_states : List[str]
1492
+ Iterable of allowed states, (Default: ['success'])
1493
+ failed_states : List[str]
1494
+ Iterable of failed or dis-allowed states. (Default: None)
1495
+ execution_delta : datetime.timedelta
1496
+ time difference with the previous execution to look at,
1497
+ the default is the same logical date as the current task or DAG. (Default: None)
1498
+ check_existence: bool
1499
+ Set to True to check if the external task exists or check if
1500
+ the DAG to wait for exists. (Default: True)
1501
+ """
1502
+ ...
1503
+
1504
+ @typing.overload
1505
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1506
+ """
1507
+ Specifies the Conda environment for all steps of the flow.
1508
+
1509
+ Use `@conda_base` to set common libraries required by all
1510
+ steps and use `@conda` to specify step-specific additions.
1511
+
1512
+
1513
+ Parameters
1514
+ ----------
1515
+ packages : Dict[str, str], default {}
1516
+ Packages to use for this flow. The key is the name of the package
1517
+ and the value is the version to use.
1518
+ libraries : Dict[str, str], default {}
1519
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1520
+ python : str, optional, default None
1521
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1522
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1523
+ disabled : bool, default False
1524
+ If set to True, disables Conda.
1525
+ """
1526
+ ...
1527
+
1528
+ @typing.overload
1529
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1530
+ ...
1531
+
1532
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1533
+ """
1534
+ Specifies the Conda environment for all steps of the flow.
1535
+
1536
+ Use `@conda_base` to set common libraries required by all
1537
+ steps and use `@conda` to specify step-specific additions.
1538
+
1539
+
1540
+ Parameters
1541
+ ----------
1542
+ packages : Dict[str, str], default {}
1543
+ Packages to use for this flow. The key is the name of the package
1544
+ and the value is the version to use.
1545
+ libraries : Dict[str, str], default {}
1546
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1547
+ python : str, optional, default None
1548
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1549
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1550
+ disabled : bool, default False
1551
+ If set to True, disables Conda.
1552
+ """
1553
+ ...
1554
+
1518
1555
  def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1519
1556
  """
1520
1557
  Specifies what flows belong to the same project.