ob-metaflow-stubs 6.0.3.176rc1__py2.py3-none-any.whl → 6.0.3.176rc2__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 (218) hide show
  1. metaflow-stubs/__init__.pyi +680 -680
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +2 -2
  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 +2 -2
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +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 +92 -92
  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 +2 -2
  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 +2 -2
  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 +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +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 +1 -1
  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 +1 -1
  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 +3 -3
  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/card_utilities/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  89. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  90. metaflow-stubs/multicore_utils.pyi +1 -1
  91. metaflow-stubs/ob_internal.pyi +1 -1
  92. metaflow-stubs/parameters.pyi +2 -2
  93. metaflow-stubs/plugins/__init__.pyi +10 -10
  94. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  96. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  99. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  100. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  101. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  103. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  104. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  105. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  106. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  107. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +1 -1
  108. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  109. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  110. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  111. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  112. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  113. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  114. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  115. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  116. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  117. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  118. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  119. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  121. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  122. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  123. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +1 -1
  124. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  125. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  126. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  127. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  128. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  129. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  130. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  131. metaflow-stubs/plugins/cards/__init__.pyi +5 -5
  132. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  133. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  135. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  140. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  142. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  143. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  144. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  145. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  146. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  148. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  149. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  151. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  152. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  153. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  154. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  155. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  156. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  157. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  158. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  159. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  160. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  161. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  163. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  164. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  165. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  169. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  170. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  171. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  172. metaflow-stubs/plugins/ollama/__init__.pyi +1 -1
  173. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  174. metaflow-stubs/plugins/perimeters.pyi +1 -1
  175. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  177. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  179. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  181. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  182. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  183. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  184. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  185. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  186. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  187. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  188. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  189. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  190. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  191. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  192. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  193. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  194. metaflow-stubs/plugins/uv/uv_environment.pyi +1 -1
  195. metaflow-stubs/profilers/__init__.pyi +1 -1
  196. metaflow-stubs/pylint_wrapper.pyi +1 -1
  197. metaflow-stubs/runner/__init__.pyi +1 -1
  198. metaflow-stubs/runner/deployer.pyi +3 -3
  199. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  200. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  201. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  202. metaflow-stubs/runner/nbrun.pyi +1 -1
  203. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  204. metaflow-stubs/runner/utils.pyi +2 -2
  205. metaflow-stubs/system/__init__.pyi +1 -1
  206. metaflow-stubs/system/system_logger.pyi +2 -2
  207. metaflow-stubs/system/system_monitor.pyi +1 -1
  208. metaflow-stubs/tagging_util.pyi +1 -1
  209. metaflow-stubs/tuple_util.pyi +1 -1
  210. metaflow-stubs/user_configs/__init__.pyi +1 -1
  211. metaflow-stubs/user_configs/config_decorators.pyi +5 -5
  212. metaflow-stubs/user_configs/config_options.pyi +1 -1
  213. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  214. {ob_metaflow_stubs-6.0.3.176rc1.dist-info → ob_metaflow_stubs-6.0.3.176rc2.dist-info}/METADATA +1 -1
  215. ob_metaflow_stubs-6.0.3.176rc2.dist-info/RECORD +218 -0
  216. ob_metaflow_stubs-6.0.3.176rc1.dist-info/RECORD +0 -218
  217. {ob_metaflow_stubs-6.0.3.176rc1.dist-info → ob_metaflow_stubs-6.0.3.176rc2.dist-info}/WHEEL +0 -0
  218. {ob_metaflow_stubs-6.0.3.176rc1.dist-info → ob_metaflow_stubs-6.0.3.176rc2.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.14.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-03T02:57:26.607647 #
4
+ # Generated on 2025-06-03T03:11:47.762585 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -36,17 +36,17 @@ 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
38
  from . import cards as cards
39
+ from . import metaflow_git as metaflow_git
39
40
  from . import tuple_util as tuple_util
40
41
  from . import events as events
41
- from . import metaflow_git as metaflow_git
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
- from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
52
52
  from .client.core import get_namespace as get_namespace
@@ -154,267 +154,268 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
154
154
  """
155
155
  ...
156
156
 
157
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
157
+ @typing.overload
158
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
158
159
  """
159
- Specifies that this step should execute on DGX cloud.
160
-
161
-
162
- Parameters
163
- ----------
164
- gpu : int
165
- Number of GPUs to use.
166
- gpu_type : str
167
- Type of Nvidia GPU to use.
168
- queue_timeout : int
169
- Time to keep the job in NVCF's queue.
160
+ Decorator prototype for all step decorators. This function gets specialized
161
+ and imported for all decorators types by _import_plugin_decorators().
170
162
  """
171
163
  ...
172
164
 
173
165
  @typing.overload
174
- 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]]]:
166
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
167
+ ...
168
+
169
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
175
170
  """
176
- Enables checkpointing for a step.
171
+ Decorator prototype for all step decorators. This function gets specialized
172
+ and imported for all decorators types by _import_plugin_decorators().
173
+ """
174
+ ...
175
+
176
+ @typing.overload
177
+ 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]]]:
178
+ """
179
+ Creates a human-readable report, a Metaflow Card, after this step completes.
177
180
 
181
+ Note that you may add multiple `@card` decorators in a step with different parameters.
178
182
 
179
183
 
180
184
  Parameters
181
185
  ----------
182
- load_policy : str, default: "fresh"
183
- The policy for loading the checkpoint. The following policies are supported:
184
- - "eager": Loads the the latest available checkpoint within the namespace.
185
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
186
- will be loaded at the start of the task.
187
- - "none": Do not load any checkpoint
188
- - "fresh": Loads the lastest checkpoint created within the running Task.
189
- This mode helps loading checkpoints across various retry attempts of the same task.
190
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
191
- created within the task will be loaded when the task is retries execution on failure.
192
-
193
- temp_dir_root : str, default: None
194
- The root directory under which `current.checkpoint.directory` will be created.
186
+ type : str, default 'default'
187
+ Card type.
188
+ id : str, optional, default None
189
+ If multiple cards are present, use this id to identify this card.
190
+ options : Dict[str, Any], default {}
191
+ Options passed to the card. The contents depend on the card type.
192
+ timeout : int, default 45
193
+ Interrupt reporting if it takes more than this many seconds.
195
194
  """
196
195
  ...
197
196
 
198
197
  @typing.overload
199
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
198
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
200
199
  ...
201
200
 
202
201
  @typing.overload
203
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
202
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
204
203
  ...
205
204
 
206
- 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):
205
+ 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):
207
206
  """
208
- Enables checkpointing for a step.
207
+ Creates a human-readable report, a Metaflow Card, after this step completes.
209
208
 
209
+ Note that you may add multiple `@card` decorators in a step with different parameters.
210
210
 
211
211
 
212
212
  Parameters
213
213
  ----------
214
- load_policy : str, default: "fresh"
215
- The policy for loading the checkpoint. The following policies are supported:
216
- - "eager": Loads the the latest available checkpoint within the namespace.
217
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
218
- will be loaded at the start of the task.
219
- - "none": Do not load any checkpoint
220
- - "fresh": Loads the lastest checkpoint created within the running Task.
221
- This mode helps loading checkpoints across various retry attempts of the same task.
222
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
223
- created within the task will be loaded when the task is retries execution on failure.
224
-
225
- temp_dir_root : str, default: None
226
- The root directory under which `current.checkpoint.directory` will be created.
214
+ type : str, default 'default'
215
+ Card type.
216
+ id : str, optional, default None
217
+ If multiple cards are present, use this id to identify this card.
218
+ options : Dict[str, Any], default {}
219
+ Options passed to the card. The contents depend on the card type.
220
+ timeout : int, default 45
221
+ Interrupt reporting if it takes more than this many seconds.
227
222
  """
228
223
  ...
229
224
 
230
225
  @typing.overload
231
- 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]]]:
226
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
232
227
  """
233
- Specifies the number of times the task corresponding
234
- to a step needs to be retried.
228
+ Specifies a timeout for your step.
235
229
 
236
- This decorator is useful for handling transient errors, such as networking issues.
237
- If your task contains operations that can't be retried safely, e.g. database updates,
238
- it is advisable to annotate it with `@retry(times=0)`.
230
+ This decorator is useful if this step may hang indefinitely.
239
231
 
240
- This can be used in conjunction with the `@catch` decorator. The `@catch`
241
- decorator will execute a no-op task after all retries have been exhausted,
242
- ensuring that the flow execution can continue.
232
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
233
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
234
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
235
+
236
+ Note that all the values specified in parameters are added together so if you specify
237
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
243
238
 
244
239
 
245
240
  Parameters
246
241
  ----------
247
- times : int, default 3
248
- Number of times to retry this task.
249
- minutes_between_retries : int, default 2
250
- Number of minutes between retries.
242
+ seconds : int, default 0
243
+ Number of seconds to wait prior to timing out.
244
+ minutes : int, default 0
245
+ Number of minutes to wait prior to timing out.
246
+ hours : int, default 0
247
+ Number of hours to wait prior to timing out.
251
248
  """
252
249
  ...
253
250
 
254
251
  @typing.overload
255
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
252
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
256
253
  ...
257
254
 
258
255
  @typing.overload
259
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
256
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
260
257
  ...
261
258
 
262
- 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):
259
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
263
260
  """
264
- Specifies the number of times the task corresponding
265
- to a step needs to be retried.
261
+ Specifies a timeout for your step.
266
262
 
267
- This decorator is useful for handling transient errors, such as networking issues.
268
- If your task contains operations that can't be retried safely, e.g. database updates,
269
- it is advisable to annotate it with `@retry(times=0)`.
263
+ This decorator is useful if this step may hang indefinitely.
270
264
 
271
- This can be used in conjunction with the `@catch` decorator. The `@catch`
272
- decorator will execute a no-op task after all retries have been exhausted,
273
- ensuring that the flow execution can continue.
265
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
266
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
267
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
268
+
269
+ Note that all the values specified in parameters are added together so if you specify
270
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
274
271
 
275
272
 
276
273
  Parameters
277
274
  ----------
278
- times : int, default 3
279
- Number of times to retry this task.
280
- minutes_between_retries : int, default 2
281
- Number of minutes between retries.
275
+ seconds : int, default 0
276
+ Number of seconds to wait prior to timing out.
277
+ minutes : int, default 0
278
+ Number of minutes to wait prior to timing out.
279
+ hours : int, default 0
280
+ Number of hours to wait prior to timing out.
282
281
  """
283
282
  ...
284
283
 
285
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
284
+ @typing.overload
285
+ 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]]]:
286
286
  """
287
- Specifies that this step should execute on Kubernetes.
287
+ Specifies the resources needed when executing this step.
288
+
289
+ Use `@resources` to specify the resource requirements
290
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
291
+
292
+ You can choose the compute layer on the command line by executing e.g.
293
+ ```
294
+ python myflow.py run --with batch
295
+ ```
296
+ or
297
+ ```
298
+ python myflow.py run --with kubernetes
299
+ ```
300
+ which executes the flow on the desired system using the
301
+ requirements specified in `@resources`.
288
302
 
289
303
 
290
304
  Parameters
291
305
  ----------
292
306
  cpu : int, default 1
293
- Number of CPUs required for this step. If `@resources` is
294
- also present, the maximum value from all decorators is used.
295
- memory : int, default 4096
296
- Memory size (in MB) required for this step. If
297
- `@resources` is also present, the maximum value from all decorators is
298
- used.
299
- disk : int, default 10240
300
- Disk size (in MB) required for this step. If
301
- `@resources` is also present, the maximum value from all decorators is
302
- used.
303
- image : str, optional, default None
304
- Docker image to use when launching on Kubernetes. If not specified, and
305
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
306
- not, a default Docker image mapping to the current version of Python is used.
307
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
308
- If given, the imagePullPolicy to be applied to the Docker image of the step.
309
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
310
- Kubernetes service account to use when launching pod in Kubernetes.
311
- secrets : List[str], optional, default None
312
- Kubernetes secrets to use when launching pod in Kubernetes. These
313
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
314
- in Metaflow configuration.
315
- node_selector: Union[Dict[str,str], str], optional, default None
316
- Kubernetes node selector(s) to apply to the pod running the task.
317
- Can be passed in as a comma separated string of values e.g.
318
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
319
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
320
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
321
- Kubernetes namespace to use when launching pod in Kubernetes.
307
+ Number of CPUs required for this step.
322
308
  gpu : int, optional, default None
323
- Number of GPUs required for this step. A value of zero implies that
324
- the scheduled node should not have GPUs.
325
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
326
- The vendor of the GPUs to be used for this step.
327
- tolerations : List[str], default []
328
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
329
- Kubernetes tolerations to use when launching pod in Kubernetes.
330
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
331
- Kubernetes labels to use when launching pod in Kubernetes.
332
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
333
- Kubernetes annotations to use when launching pod in Kubernetes.
334
- use_tmpfs : bool, default False
335
- This enables an explicit tmpfs mount for this step.
336
- tmpfs_tempdir : bool, default True
337
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
338
- tmpfs_size : int, optional, default: None
339
- The value for the size (in MiB) of the tmpfs mount for this step.
340
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
341
- memory allocated for this step.
342
- tmpfs_path : str, optional, default /metaflow_temp
343
- Path to tmpfs mount for this step.
344
- persistent_volume_claims : Dict[str, str], optional, default None
345
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
346
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
347
- shared_memory: int, optional
348
- Shared memory size (in MiB) required for this step
349
- port: int, optional
350
- Port number to specify in the Kubernetes job object
351
- compute_pool : str, optional, default None
352
- Compute pool to be used for for this step.
353
- If not specified, any accessible compute pool within the perimeter is used.
354
- hostname_resolution_timeout: int, default 10 * 60
355
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
356
- Only applicable when @parallel is used.
357
- qos: str, default: Burstable
358
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
359
-
360
- security_context: Dict[str, Any], optional, default None
361
- Container security context. Applies to the task container. Allows the following keys:
362
- - privileged: bool, optional, default None
363
- - allow_privilege_escalation: bool, optional, default None
364
- - run_as_user: int, optional, default None
365
- - run_as_group: int, optional, default None
366
- - run_as_non_root: bool, optional, default None
309
+ Number of GPUs required for this step.
310
+ disk : int, optional, default None
311
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
312
+ memory : int, default 4096
313
+ Memory size (in MB) required for this step.
314
+ shared_memory : int, optional, default None
315
+ The value for the size (in MiB) of the /dev/shm volume for this step.
316
+ This parameter maps to the `--shm-size` option in Docker.
367
317
  """
368
318
  ...
369
319
 
370
320
  @typing.overload
371
- 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]]]:
372
- """
373
- Specifies that the step will success under all circumstances.
374
-
375
- The decorator will create an optional artifact, specified by `var`, which
376
- contains the exception raised. You can use it to detect the presence
377
- of errors, indicating that all happy-path artifacts produced by the step
378
- are missing.
321
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
322
+ ...
323
+
324
+ @typing.overload
325
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
326
+ ...
327
+
328
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
329
+ """
330
+ Specifies the resources needed when executing this step.
331
+
332
+ Use `@resources` to specify the resource requirements
333
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
334
+
335
+ You can choose the compute layer on the command line by executing e.g.
336
+ ```
337
+ python myflow.py run --with batch
338
+ ```
339
+ or
340
+ ```
341
+ python myflow.py run --with kubernetes
342
+ ```
343
+ which executes the flow on the desired system using the
344
+ requirements specified in `@resources`.
379
345
 
380
346
 
381
347
  Parameters
382
348
  ----------
383
- var : str, optional, default None
384
- Name of the artifact in which to store the caught exception.
385
- If not specified, the exception is not stored.
386
- print_exception : bool, default True
387
- Determines whether or not the exception is printed to
388
- stdout when caught.
349
+ cpu : int, default 1
350
+ Number of CPUs required for this step.
351
+ gpu : int, optional, default None
352
+ Number of GPUs required for this step.
353
+ disk : int, optional, default None
354
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
355
+ memory : int, default 4096
356
+ Memory size (in MB) required for this step.
357
+ shared_memory : int, optional, default None
358
+ The value for the size (in MiB) of the /dev/shm volume for this step.
359
+ This parameter maps to the `--shm-size` option in Docker.
389
360
  """
390
361
  ...
391
362
 
392
363
  @typing.overload
393
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
364
+ def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
365
+ """
366
+ Specifies the Conda environment for the step.
367
+
368
+ Information in this decorator will augment any
369
+ attributes set in the `@conda_base` flow-level decorator. Hence,
370
+ you can use `@conda_base` to set packages required by all
371
+ steps and use `@conda` to specify step-specific overrides.
372
+
373
+
374
+ Parameters
375
+ ----------
376
+ packages : Dict[str, str], default {}
377
+ Packages to use for this step. The key is the name of the package
378
+ and the value is the version to use.
379
+ libraries : Dict[str, str], default {}
380
+ Supported for backward compatibility. When used with packages, packages will take precedence.
381
+ python : str, optional, default None
382
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
383
+ that the version used will correspond to the version of the Python interpreter used to start the run.
384
+ disabled : bool, default False
385
+ If set to True, disables @conda.
386
+ """
394
387
  ...
395
388
 
396
389
  @typing.overload
397
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
390
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
398
391
  ...
399
392
 
400
- 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):
393
+ @typing.overload
394
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
395
+ ...
396
+
397
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
401
398
  """
402
- Specifies that the step will success under all circumstances.
399
+ Specifies the Conda environment for the step.
403
400
 
404
- The decorator will create an optional artifact, specified by `var`, which
405
- contains the exception raised. You can use it to detect the presence
406
- of errors, indicating that all happy-path artifacts produced by the step
407
- are missing.
401
+ Information in this decorator will augment any
402
+ attributes set in the `@conda_base` flow-level decorator. Hence,
403
+ you can use `@conda_base` to set packages required by all
404
+ steps and use `@conda` to specify step-specific overrides.
408
405
 
409
406
 
410
407
  Parameters
411
408
  ----------
412
- var : str, optional, default None
413
- Name of the artifact in which to store the caught exception.
414
- If not specified, the exception is not stored.
415
- print_exception : bool, default True
416
- Determines whether or not the exception is printed to
417
- stdout when caught.
409
+ packages : Dict[str, str], default {}
410
+ Packages to use for this step. The key is the name of the package
411
+ and the value is the version to use.
412
+ libraries : Dict[str, str], default {}
413
+ Supported for backward compatibility. When used with packages, packages will take precedence.
414
+ python : str, optional, default None
415
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
416
+ that the version used will correspond to the version of the Python interpreter used to start the run.
417
+ disabled : bool, default False
418
+ If set to True, disables @conda.
418
419
  """
419
420
  ...
420
421
 
@@ -443,201 +444,71 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
443
444
  """
444
445
  ...
445
446
 
446
- 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]]]:
447
- """
448
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
449
-
450
- User code call
451
- --------------
452
- @ollama(
453
- models=[...],
454
- ...
455
- )
456
-
457
- Valid backend options
458
- ---------------------
459
- - 'local': Run as a separate process on the local task machine.
460
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
461
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
462
-
463
- Valid model options
464
- -------------------
465
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
466
-
467
-
468
- Parameters
469
- ----------
470
- models: list[str]
471
- List of Ollama containers running models in sidecars.
472
- backend: str
473
- Determines where and how to run the Ollama process.
474
- force_pull: bool
475
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
476
- cache_update_policy: str
477
- Cache update policy: "auto", "force", or "never".
478
- force_cache_update: bool
479
- Simple override for "force" cache update policy.
480
- debug: bool
481
- Whether to turn on verbose debugging logs.
482
- circuit_breaker_config: dict
483
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
484
- timeout_config: dict
485
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
486
- """
487
- ...
488
-
489
447
  @typing.overload
490
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
448
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
491
449
  """
492
- Specifies a timeout for your step.
493
-
494
- This decorator is useful if this step may hang indefinitely.
495
-
496
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
497
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
498
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
499
-
500
- Note that all the values specified in parameters are added together so if you specify
501
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
502
-
503
-
504
- Parameters
505
- ----------
506
- seconds : int, default 0
507
- Number of seconds to wait prior to timing out.
508
- minutes : int, default 0
509
- Number of minutes to wait prior to timing out.
510
- hours : int, default 0
511
- Number of hours to wait prior to timing out.
450
+ Internal decorator to support Fast bakery
512
451
  """
513
452
  ...
514
453
 
515
454
  @typing.overload
516
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
517
- ...
518
-
519
- @typing.overload
520
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
455
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
521
456
  ...
522
457
 
523
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
458
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
524
459
  """
525
- Specifies a timeout for your step.
526
-
527
- This decorator is useful if this step may hang indefinitely.
528
-
529
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
530
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
531
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
532
-
533
- Note that all the values specified in parameters are added together so if you specify
534
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
535
-
536
-
537
- Parameters
538
- ----------
539
- seconds : int, default 0
540
- Number of seconds to wait prior to timing out.
541
- minutes : int, default 0
542
- Number of minutes to wait prior to timing out.
543
- hours : int, default 0
544
- Number of hours to wait prior to timing out.
460
+ Internal decorator to support Fast bakery
545
461
  """
546
462
  ...
547
463
 
548
464
  @typing.overload
549
- 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]]]:
465
+ 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]]]:
550
466
  """
551
- Specifies the resources needed when executing this step.
552
-
553
- Use `@resources` to specify the resource requirements
554
- independently of the specific compute layer (`@batch`, `@kubernetes`).
555
-
556
- You can choose the compute layer on the command line by executing e.g.
557
- ```
558
- python myflow.py run --with batch
559
- ```
560
- or
561
- ```
562
- python myflow.py run --with kubernetes
563
- ```
564
- which executes the flow on the desired system using the
565
- requirements specified in `@resources`.
467
+ Specifies secrets to be retrieved and injected as environment variables prior to
468
+ the execution of a step.
566
469
 
567
470
 
568
471
  Parameters
569
472
  ----------
570
- cpu : int, default 1
571
- Number of CPUs required for this step.
572
- gpu : int, optional, default None
573
- Number of GPUs required for this step.
574
- disk : int, optional, default None
575
- Disk size (in MB) required for this step. Only applies on Kubernetes.
576
- memory : int, default 4096
577
- Memory size (in MB) required for this step.
578
- shared_memory : int, optional, default None
579
- The value for the size (in MiB) of the /dev/shm volume for this step.
580
- This parameter maps to the `--shm-size` option in Docker.
473
+ sources : List[Union[str, Dict[str, Any]]], default: []
474
+ List of secret specs, defining how the secrets are to be retrieved
581
475
  """
582
476
  ...
583
477
 
584
478
  @typing.overload
585
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
479
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
586
480
  ...
587
481
 
588
482
  @typing.overload
589
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
483
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
590
484
  ...
591
485
 
592
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
486
+ 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]]] = []):
593
487
  """
594
- Specifies the resources needed when executing this step.
595
-
596
- Use `@resources` to specify the resource requirements
597
- independently of the specific compute layer (`@batch`, `@kubernetes`).
598
-
599
- You can choose the compute layer on the command line by executing e.g.
600
- ```
601
- python myflow.py run --with batch
602
- ```
603
- or
604
- ```
605
- python myflow.py run --with kubernetes
606
- ```
607
- which executes the flow on the desired system using the
608
- requirements specified in `@resources`.
488
+ Specifies secrets to be retrieved and injected as environment variables prior to
489
+ the execution of a step.
609
490
 
610
491
 
611
492
  Parameters
612
493
  ----------
613
- cpu : int, default 1
614
- Number of CPUs required for this step.
615
- gpu : int, optional, default None
616
- Number of GPUs required for this step.
617
- disk : int, optional, default None
618
- Disk size (in MB) required for this step. Only applies on Kubernetes.
619
- memory : int, default 4096
620
- Memory size (in MB) required for this step.
621
- shared_memory : int, optional, default None
622
- The value for the size (in MiB) of the /dev/shm volume for this step.
623
- This parameter maps to the `--shm-size` option in Docker.
624
- """
625
- ...
626
-
627
- @typing.overload
628
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
629
- """
630
- Internal decorator to support Fast bakery
494
+ sources : List[Union[str, Dict[str, Any]]], default: []
495
+ List of secret specs, defining how the secrets are to be retrieved
631
496
  """
632
497
  ...
633
498
 
634
- @typing.overload
635
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
636
- ...
637
-
638
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
499
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
639
500
  """
640
- Internal decorator to support Fast bakery
501
+ Specifies that this step should execute on DGX cloud.
502
+
503
+
504
+ Parameters
505
+ ----------
506
+ gpu : int
507
+ Number of GPUs to use.
508
+ gpu_type : str
509
+ Type of Nvidia GPU to use.
510
+ queue_timeout : int
511
+ Time to keep the job in NVCF's queue.
641
512
  """
642
513
  ...
643
514
 
@@ -674,6 +545,20 @@ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], Non
674
545
  """
675
546
  ...
676
547
 
548
+ 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]]]:
549
+ """
550
+ Specifies that this step should execute on DGX cloud.
551
+
552
+
553
+ Parameters
554
+ ----------
555
+ gpu : int
556
+ Number of GPUs to use.
557
+ gpu_type : str
558
+ Type of Nvidia GPU to use.
559
+ """
560
+ ...
561
+
677
562
  @typing.overload
678
563
  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]]]:
679
564
  """
@@ -726,84 +611,59 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
726
611
  ...
727
612
 
728
613
  @typing.overload
729
- 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]]]:
730
- """
731
- Creates a human-readable report, a Metaflow Card, after this step completes.
732
-
733
- Note that you may add multiple `@card` decorators in a step with different parameters.
734
-
735
-
736
- Parameters
737
- ----------
738
- type : str, default 'default'
739
- Card type.
740
- id : str, optional, default None
741
- If multiple cards are present, use this id to identify this card.
742
- options : Dict[str, Any], default {}
743
- Options passed to the card. The contents depend on the card type.
744
- timeout : int, default 45
745
- Interrupt reporting if it takes more than this many seconds.
746
- """
747
- ...
748
-
749
- @typing.overload
750
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
751
- ...
752
-
753
- @typing.overload
754
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
755
- ...
756
-
757
- 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):
614
+ 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]]]:
758
615
  """
759
- Creates a human-readable report, a Metaflow Card, after this step completes.
616
+ Enables checkpointing for a step.
760
617
 
761
- Note that you may add multiple `@card` decorators in a step with different parameters.
762
618
 
763
619
 
764
620
  Parameters
765
621
  ----------
766
- type : str, default 'default'
767
- Card type.
768
- id : str, optional, default None
769
- If multiple cards are present, use this id to identify this card.
770
- options : Dict[str, Any], default {}
771
- Options passed to the card. The contents depend on the card type.
772
- timeout : int, default 45
773
- Interrupt reporting if it takes more than this many seconds.
622
+ load_policy : str, default: "fresh"
623
+ The policy for loading the checkpoint. The following policies are supported:
624
+ - "eager": Loads the the latest available checkpoint within the namespace.
625
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
626
+ will be loaded at the start of the task.
627
+ - "none": Do not load any checkpoint
628
+ - "fresh": Loads the lastest checkpoint created within the running Task.
629
+ This mode helps loading checkpoints across various retry attempts of the same task.
630
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
631
+ created within the task will be loaded when the task is retries execution on failure.
632
+
633
+ temp_dir_root : str, default: None
634
+ The root directory under which `current.checkpoint.directory` will be created.
774
635
  """
775
636
  ...
776
637
 
777
638
  @typing.overload
778
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
779
- """
780
- Decorator prototype for all step decorators. This function gets specialized
781
- and imported for all decorators types by _import_plugin_decorators().
782
- """
639
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
783
640
  ...
784
641
 
785
642
  @typing.overload
786
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
787
- ...
788
-
789
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
790
- """
791
- Decorator prototype for all step decorators. This function gets specialized
792
- and imported for all decorators types by _import_plugin_decorators().
793
- """
643
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
794
644
  ...
795
645
 
796
- 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]]]:
646
+ 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):
797
647
  """
798
- Specifies that this step should execute on DGX cloud.
648
+ Enables checkpointing for a step.
649
+
799
650
 
800
651
 
801
652
  Parameters
802
653
  ----------
803
- gpu : int
804
- Number of GPUs to use.
805
- gpu_type : str
806
- Type of Nvidia GPU to use.
654
+ load_policy : str, default: "fresh"
655
+ The policy for loading the checkpoint. The following policies are supported:
656
+ - "eager": Loads the the latest available checkpoint within the namespace.
657
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
658
+ will be loaded at the start of the task.
659
+ - "none": Do not load any checkpoint
660
+ - "fresh": Loads the lastest checkpoint created within the running Task.
661
+ This mode helps loading checkpoints across various retry attempts of the same task.
662
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
663
+ created within the task will be loaded when the task is retries execution on failure.
664
+
665
+ temp_dir_root : str, default: None
666
+ The root directory under which `current.checkpoint.directory` will be created.
807
667
  """
808
668
  ...
809
669
 
@@ -880,182 +740,236 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
880
740
  ...
881
741
 
882
742
  @typing.overload
883
- 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]]]:
743
+ 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]]]:
884
744
  """
885
- Specifies secrets to be retrieved and injected as environment variables prior to
886
- the execution of a step.
745
+ Specifies that the step will success under all circumstances.
746
+
747
+ The decorator will create an optional artifact, specified by `var`, which
748
+ contains the exception raised. You can use it to detect the presence
749
+ of errors, indicating that all happy-path artifacts produced by the step
750
+ are missing.
887
751
 
888
752
 
889
753
  Parameters
890
754
  ----------
891
- sources : List[Union[str, Dict[str, Any]]], default: []
892
- List of secret specs, defining how the secrets are to be retrieved
755
+ var : str, optional, default None
756
+ Name of the artifact in which to store the caught exception.
757
+ If not specified, the exception is not stored.
758
+ print_exception : bool, default True
759
+ Determines whether or not the exception is printed to
760
+ stdout when caught.
893
761
  """
894
762
  ...
895
763
 
896
764
  @typing.overload
897
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
765
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
898
766
  ...
899
767
 
900
768
  @typing.overload
901
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
769
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
902
770
  ...
903
771
 
904
- 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]]] = []):
772
+ 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):
905
773
  """
906
- Specifies secrets to be retrieved and injected as environment variables prior to
907
- the execution of a step.
774
+ Specifies that the step will success under all circumstances.
775
+
776
+ The decorator will create an optional artifact, specified by `var`, which
777
+ contains the exception raised. You can use it to detect the presence
778
+ of errors, indicating that all happy-path artifacts produced by the step
779
+ are missing.
908
780
 
909
781
 
910
782
  Parameters
911
783
  ----------
912
- sources : List[Union[str, Dict[str, Any]]], default: []
913
- List of secret specs, defining how the secrets are to be retrieved
784
+ var : str, optional, default None
785
+ Name of the artifact in which to store the caught exception.
786
+ If not specified, the exception is not stored.
787
+ print_exception : bool, default True
788
+ Determines whether or not the exception is printed to
789
+ stdout when caught.
914
790
  """
915
791
  ...
916
792
 
917
793
  @typing.overload
918
- def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
794
+ 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]]]:
919
795
  """
920
- Specifies the Conda environment for the step.
796
+ Specifies the number of times the task corresponding
797
+ to a step needs to be retried.
921
798
 
922
- Information in this decorator will augment any
923
- attributes set in the `@conda_base` flow-level decorator. Hence,
924
- you can use `@conda_base` to set packages required by all
925
- steps and use `@conda` to specify step-specific overrides.
799
+ This decorator is useful for handling transient errors, such as networking issues.
800
+ If your task contains operations that can't be retried safely, e.g. database updates,
801
+ it is advisable to annotate it with `@retry(times=0)`.
802
+
803
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
804
+ decorator will execute a no-op task after all retries have been exhausted,
805
+ ensuring that the flow execution can continue.
926
806
 
927
807
 
928
808
  Parameters
929
809
  ----------
930
- packages : Dict[str, str], default {}
931
- Packages to use for this step. The key is the name of the package
932
- and the value is the version to use.
933
- libraries : Dict[str, str], default {}
934
- Supported for backward compatibility. When used with packages, packages will take precedence.
935
- python : str, optional, default None
936
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
937
- that the version used will correspond to the version of the Python interpreter used to start the run.
938
- disabled : bool, default False
939
- If set to True, disables @conda.
810
+ times : int, default 3
811
+ Number of times to retry this task.
812
+ minutes_between_retries : int, default 2
813
+ Number of minutes between retries.
940
814
  """
941
815
  ...
942
816
 
943
817
  @typing.overload
944
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
818
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
945
819
  ...
946
820
 
947
821
  @typing.overload
948
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
822
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
949
823
  ...
950
824
 
951
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
825
+ 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):
952
826
  """
953
- Specifies the Conda environment for the step.
827
+ Specifies the number of times the task corresponding
828
+ to a step needs to be retried.
954
829
 
955
- Information in this decorator will augment any
956
- attributes set in the `@conda_base` flow-level decorator. Hence,
957
- you can use `@conda_base` to set packages required by all
958
- steps and use `@conda` to specify step-specific overrides.
830
+ This decorator is useful for handling transient errors, such as networking issues.
831
+ If your task contains operations that can't be retried safely, e.g. database updates,
832
+ it is advisable to annotate it with `@retry(times=0)`.
833
+
834
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
835
+ decorator will execute a no-op task after all retries have been exhausted,
836
+ ensuring that the flow execution can continue.
959
837
 
960
838
 
961
839
  Parameters
962
840
  ----------
963
- packages : Dict[str, str], default {}
964
- Packages to use for this step. The key is the name of the package
965
- and the value is the version to use.
966
- libraries : Dict[str, str], default {}
967
- Supported for backward compatibility. When used with packages, packages will take precedence.
968
- python : str, optional, default None
969
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
970
- that the version used will correspond to the version of the Python interpreter used to start the run.
971
- disabled : bool, default False
972
- If set to True, disables @conda.
841
+ times : int, default 3
842
+ Number of times to retry this task.
843
+ minutes_between_retries : int, default 2
844
+ Number of minutes between retries.
973
845
  """
974
846
  ...
975
847
 
976
- 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]]:
848
+ 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]]]:
977
849
  """
978
- 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.
979
- 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.
850
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
851
+
852
+ User code call
853
+ --------------
854
+ @ollama(
855
+ models=[...],
856
+ ...
857
+ )
858
+
859
+ Valid backend options
860
+ ---------------------
861
+ - 'local': Run as a separate process on the local task machine.
862
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
863
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
864
+
865
+ Valid model options
866
+ -------------------
867
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
980
868
 
981
869
 
982
870
  Parameters
983
871
  ----------
984
- timeout : int
985
- Time, in seconds before the task times out and fails. (Default: 3600)
986
- poke_interval : int
987
- Time in seconds that the job should wait in between each try. (Default: 60)
988
- mode : str
989
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
990
- exponential_backoff : bool
991
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
992
- pool : str
993
- the slot pool this task should run in,
994
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
995
- soft_fail : bool
996
- Set to true to mark the task as SKIPPED on failure. (Default: False)
997
- name : str
998
- Name of the sensor on Airflow
999
- description : str
1000
- Description of sensor in the Airflow UI
1001
- external_dag_id : str
1002
- The dag_id that contains the task you want to wait for.
1003
- external_task_ids : List[str]
1004
- The list of task_ids that you want to wait for.
1005
- If None (default value) the sensor waits for the DAG. (Default: None)
1006
- allowed_states : List[str]
1007
- Iterable of allowed states, (Default: ['success'])
1008
- failed_states : List[str]
1009
- Iterable of failed or dis-allowed states. (Default: None)
1010
- execution_delta : datetime.timedelta
1011
- time difference with the previous execution to look at,
1012
- the default is the same logical date as the current task or DAG. (Default: None)
1013
- check_existence: bool
1014
- Set to True to check if the external task exists or check if
1015
- the DAG to wait for exists. (Default: True)
872
+ models: list[str]
873
+ List of Ollama containers running models in sidecars.
874
+ backend: str
875
+ Determines where and how to run the Ollama process.
876
+ force_pull: bool
877
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
878
+ cache_update_policy: str
879
+ Cache update policy: "auto", "force", or "never".
880
+ force_cache_update: bool
881
+ Simple override for "force" cache update policy.
882
+ debug: bool
883
+ Whether to turn on verbose debugging logs.
884
+ circuit_breaker_config: dict
885
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
886
+ timeout_config: dict
887
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
1016
888
  """
1017
889
  ...
1018
890
 
1019
- 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]]:
891
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1020
892
  """
1021
- 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)
1022
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1023
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1024
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1025
- starts only after all sensors finish.
893
+ Specifies that this step should execute on Kubernetes.
1026
894
 
1027
895
 
1028
- Parameters
1029
- ----------
1030
- timeout : int
1031
- Time, in seconds before the task times out and fails. (Default: 3600)
1032
- poke_interval : int
1033
- Time in seconds that the job should wait in between each try. (Default: 60)
1034
- mode : str
1035
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1036
- exponential_backoff : bool
1037
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1038
- pool : str
1039
- the slot pool this task should run in,
1040
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1041
- soft_fail : bool
1042
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1043
- name : str
1044
- Name of the sensor on Airflow
1045
- description : str
1046
- Description of sensor in the Airflow UI
1047
- bucket_key : Union[str, List[str]]
1048
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1049
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1050
- bucket_name : str
1051
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1052
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1053
- wildcard_match : bool
1054
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1055
- aws_conn_id : str
1056
- a reference to the s3 connection on Airflow. (Default: None)
1057
- verify : bool
1058
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
896
+ Parameters
897
+ ----------
898
+ cpu : int, default 1
899
+ Number of CPUs required for this step. If `@resources` is
900
+ also present, the maximum value from all decorators is used.
901
+ memory : int, default 4096
902
+ Memory size (in MB) required for this step. If
903
+ `@resources` is also present, the maximum value from all decorators is
904
+ used.
905
+ disk : int, default 10240
906
+ Disk size (in MB) required for this step. If
907
+ `@resources` is also present, the maximum value from all decorators is
908
+ used.
909
+ image : str, optional, default None
910
+ Docker image to use when launching on Kubernetes. If not specified, and
911
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
912
+ not, a default Docker image mapping to the current version of Python is used.
913
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
914
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
915
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
916
+ Kubernetes service account to use when launching pod in Kubernetes.
917
+ secrets : List[str], optional, default None
918
+ Kubernetes secrets to use when launching pod in Kubernetes. These
919
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
920
+ in Metaflow configuration.
921
+ node_selector: Union[Dict[str,str], str], optional, default None
922
+ Kubernetes node selector(s) to apply to the pod running the task.
923
+ Can be passed in as a comma separated string of values e.g.
924
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
925
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
926
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
927
+ Kubernetes namespace to use when launching pod in Kubernetes.
928
+ gpu : int, optional, default None
929
+ Number of GPUs required for this step. A value of zero implies that
930
+ the scheduled node should not have GPUs.
931
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
932
+ The vendor of the GPUs to be used for this step.
933
+ tolerations : List[str], default []
934
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
935
+ Kubernetes tolerations to use when launching pod in Kubernetes.
936
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
937
+ Kubernetes labels to use when launching pod in Kubernetes.
938
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
939
+ Kubernetes annotations to use when launching pod in Kubernetes.
940
+ use_tmpfs : bool, default False
941
+ This enables an explicit tmpfs mount for this step.
942
+ tmpfs_tempdir : bool, default True
943
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
944
+ tmpfs_size : int, optional, default: None
945
+ The value for the size (in MiB) of the tmpfs mount for this step.
946
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
947
+ memory allocated for this step.
948
+ tmpfs_path : str, optional, default /metaflow_temp
949
+ Path to tmpfs mount for this step.
950
+ persistent_volume_claims : Dict[str, str], optional, default None
951
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
952
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
953
+ shared_memory: int, optional
954
+ Shared memory size (in MiB) required for this step
955
+ port: int, optional
956
+ Port number to specify in the Kubernetes job object
957
+ compute_pool : str, optional, default None
958
+ Compute pool to be used for for this step.
959
+ If not specified, any accessible compute pool within the perimeter is used.
960
+ hostname_resolution_timeout: int, default 10 * 60
961
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
962
+ Only applicable when @parallel is used.
963
+ qos: str, default: Burstable
964
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
965
+
966
+ security_context: Dict[str, Any], optional, default None
967
+ Container security context. Applies to the task container. Allows the following keys:
968
+ - privileged: bool, optional, default None
969
+ - allow_privilege_escalation: bool, optional, default None
970
+ - run_as_user: int, optional, default None
971
+ - run_as_group: int, optional, default None
972
+ - run_as_non_root: bool, optional, default None
1059
973
  """
1060
974
  ...
1061
975
 
@@ -1095,103 +1009,94 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1095
1009
  ...
1096
1010
 
1097
1011
  @typing.overload
1098
- 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]]:
1012
+ 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]]:
1099
1013
  """
1100
- Specifies the flow(s) that this flow depends on.
1101
-
1102
- ```
1103
- @trigger_on_finish(flow='FooFlow')
1104
- ```
1105
- or
1106
- ```
1107
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1108
- ```
1109
- This decorator respects the @project decorator and triggers the flow
1110
- when upstream runs within the same namespace complete successfully
1111
-
1112
- Additionally, you can specify project aware upstream flow dependencies
1113
- by specifying the fully qualified project_flow_name.
1114
- ```
1115
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1116
- ```
1117
- or
1118
- ```
1119
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1120
- ```
1121
-
1122
- You can also specify just the project or project branch (other values will be
1123
- inferred from the current project or project branch):
1124
- ```
1125
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1126
- ```
1014
+ Specifies the Conda environment for all steps of the flow.
1127
1015
 
1128
- Note that `branch` is typically one of:
1129
- - `prod`
1130
- - `user.bob`
1131
- - `test.my_experiment`
1132
- - `prod.staging`
1016
+ Use `@conda_base` to set common libraries required by all
1017
+ steps and use `@conda` to specify step-specific additions.
1133
1018
 
1134
1019
 
1135
1020
  Parameters
1136
1021
  ----------
1137
- flow : Union[str, Dict[str, str]], optional, default None
1138
- Upstream flow dependency for this flow.
1139
- flows : List[Union[str, Dict[str, str]]], default []
1140
- Upstream flow dependencies for this flow.
1141
- options : Dict[str, Any], default {}
1142
- Backend-specific configuration for tuning eventing behavior.
1022
+ packages : Dict[str, str], default {}
1023
+ Packages to use for this flow. The key is the name of the package
1024
+ and the value is the version to use.
1025
+ libraries : Dict[str, str], default {}
1026
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1027
+ python : str, optional, default None
1028
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1029
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1030
+ disabled : bool, default False
1031
+ If set to True, disables Conda.
1143
1032
  """
1144
1033
  ...
1145
1034
 
1146
1035
  @typing.overload
1147
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1036
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1148
1037
  ...
1149
1038
 
1150
- 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] = {}):
1039
+ 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):
1151
1040
  """
1152
- Specifies the flow(s) that this flow depends on.
1041
+ Specifies the Conda environment for all steps of the flow.
1153
1042
 
1154
- ```
1155
- @trigger_on_finish(flow='FooFlow')
1156
- ```
1157
- or
1158
- ```
1159
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1160
- ```
1161
- This decorator respects the @project decorator and triggers the flow
1162
- when upstream runs within the same namespace complete successfully
1043
+ Use `@conda_base` to set common libraries required by all
1044
+ steps and use `@conda` to specify step-specific additions.
1163
1045
 
1164
- Additionally, you can specify project aware upstream flow dependencies
1165
- by specifying the fully qualified project_flow_name.
1166
- ```
1167
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1168
- ```
1169
- or
1170
- ```
1171
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1172
- ```
1173
1046
 
1174
- You can also specify just the project or project branch (other values will be
1175
- inferred from the current project or project branch):
1176
- ```
1177
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1178
- ```
1047
+ Parameters
1048
+ ----------
1049
+ packages : Dict[str, str], default {}
1050
+ Packages to use for this flow. The key is the name of the package
1051
+ and the value is the version to use.
1052
+ libraries : Dict[str, str], default {}
1053
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1054
+ python : str, optional, default None
1055
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1056
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1057
+ disabled : bool, default False
1058
+ If set to True, disables Conda.
1059
+ """
1060
+ ...
1061
+
1062
+ @typing.overload
1063
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1064
+ """
1065
+ Specifies the PyPI packages for all steps of the flow.
1179
1066
 
1180
- Note that `branch` is typically one of:
1181
- - `prod`
1182
- - `user.bob`
1183
- - `test.my_experiment`
1184
- - `prod.staging`
1067
+ Use `@pypi_base` to set common packages required by all
1068
+ steps and use `@pypi` to specify step-specific overrides.
1069
+
1070
+ Parameters
1071
+ ----------
1072
+ packages : Dict[str, str], default: {}
1073
+ Packages to use for this flow. The key is the name of the package
1074
+ and the value is the version to use.
1075
+ python : str, optional, default: None
1076
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1077
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1078
+ """
1079
+ ...
1080
+
1081
+ @typing.overload
1082
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1083
+ ...
1084
+
1085
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1086
+ """
1087
+ Specifies the PyPI packages for all steps of the flow.
1185
1088
 
1089
+ Use `@pypi_base` to set common packages required by all
1090
+ steps and use `@pypi` to specify step-specific overrides.
1186
1091
 
1187
1092
  Parameters
1188
1093
  ----------
1189
- flow : Union[str, Dict[str, str]], optional, default None
1190
- Upstream flow dependency for this flow.
1191
- flows : List[Union[str, Dict[str, str]]], default []
1192
- Upstream flow dependencies for this flow.
1193
- options : Dict[str, Any], default {}
1194
- Backend-specific configuration for tuning eventing behavior.
1094
+ packages : Dict[str, str], default: {}
1095
+ Packages to use for this flow. The key is the name of the package
1096
+ and the value is the version to use.
1097
+ python : str, optional, default: None
1098
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1099
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1195
1100
  """
1196
1101
  ...
1197
1102
 
@@ -1279,53 +1184,55 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1279
1184
 
1280
1185
  Parameters
1281
1186
  ----------
1282
- event : Union[str, Dict[str, Any]], optional, default None
1283
- Event dependency for this flow.
1284
- events : List[Union[str, Dict[str, Any]]], default []
1285
- Events dependency for this flow.
1286
- options : Dict[str, Any], default {}
1287
- Backend-specific configuration for tuning eventing behavior.
1288
- """
1289
- ...
1290
-
1291
- @typing.overload
1292
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1293
- """
1294
- Specifies the PyPI packages for all steps of the flow.
1295
-
1296
- Use `@pypi_base` to set common packages required by all
1297
- steps and use `@pypi` to specify step-specific overrides.
1298
-
1299
- Parameters
1300
- ----------
1301
- packages : Dict[str, str], default: {}
1302
- Packages to use for this flow. The key is the name of the package
1303
- and the value is the version to use.
1304
- python : str, optional, default: None
1305
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1306
- that the version used will correspond to the version of the Python interpreter used to start the run.
1187
+ event : Union[str, Dict[str, Any]], optional, default None
1188
+ Event dependency for this flow.
1189
+ events : List[Union[str, Dict[str, Any]]], default []
1190
+ Events dependency for this flow.
1191
+ options : Dict[str, Any], default {}
1192
+ Backend-specific configuration for tuning eventing behavior.
1307
1193
  """
1308
1194
  ...
1309
1195
 
1310
- @typing.overload
1311
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1312
- ...
1313
-
1314
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1196
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1315
1197
  """
1316
- Specifies the PyPI packages for all steps of the flow.
1198
+ 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)
1199
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1200
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1201
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1202
+ starts only after all sensors finish.
1317
1203
 
1318
- Use `@pypi_base` to set common packages required by all
1319
- steps and use `@pypi` to specify step-specific overrides.
1320
1204
 
1321
1205
  Parameters
1322
1206
  ----------
1323
- packages : Dict[str, str], default: {}
1324
- Packages to use for this flow. The key is the name of the package
1325
- and the value is the version to use.
1326
- python : str, optional, default: None
1327
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1328
- that the version used will correspond to the version of the Python interpreter used to start the run.
1207
+ timeout : int
1208
+ Time, in seconds before the task times out and fails. (Default: 3600)
1209
+ poke_interval : int
1210
+ Time in seconds that the job should wait in between each try. (Default: 60)
1211
+ mode : str
1212
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1213
+ exponential_backoff : bool
1214
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1215
+ pool : str
1216
+ the slot pool this task should run in,
1217
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1218
+ soft_fail : bool
1219
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1220
+ name : str
1221
+ Name of the sensor on Airflow
1222
+ description : str
1223
+ Description of sensor in the Airflow UI
1224
+ bucket_key : Union[str, List[str]]
1225
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1226
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1227
+ bucket_name : str
1228
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1229
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1230
+ wildcard_match : bool
1231
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1232
+ aws_conn_id : str
1233
+ a reference to the s3 connection on Airflow. (Default: None)
1234
+ verify : bool
1235
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1329
1236
  """
1330
1237
  ...
1331
1238
 
@@ -1380,57 +1287,6 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1380
1287
  """
1381
1288
  ...
1382
1289
 
1383
- @typing.overload
1384
- 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]]:
1385
- """
1386
- Specifies the Conda environment for all steps of the flow.
1387
-
1388
- Use `@conda_base` to set common libraries required by all
1389
- steps and use `@conda` to specify step-specific additions.
1390
-
1391
-
1392
- Parameters
1393
- ----------
1394
- packages : Dict[str, str], default {}
1395
- Packages to use for this flow. The key is the name of the package
1396
- and the value is the version to use.
1397
- libraries : Dict[str, str], default {}
1398
- Supported for backward compatibility. When used with packages, packages will take precedence.
1399
- python : str, optional, default None
1400
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1401
- that the version used will correspond to the version of the Python interpreter used to start the run.
1402
- disabled : bool, default False
1403
- If set to True, disables Conda.
1404
- """
1405
- ...
1406
-
1407
- @typing.overload
1408
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1409
- ...
1410
-
1411
- 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):
1412
- """
1413
- Specifies the Conda environment for all steps of the flow.
1414
-
1415
- Use `@conda_base` to set common libraries required by all
1416
- steps and use `@conda` to specify step-specific additions.
1417
-
1418
-
1419
- Parameters
1420
- ----------
1421
- packages : Dict[str, str], default {}
1422
- Packages to use for this flow. The key is the name of the package
1423
- and the value is the version to use.
1424
- libraries : Dict[str, str], default {}
1425
- Supported for backward compatibility. When used with packages, packages will take precedence.
1426
- python : str, optional, default None
1427
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1428
- that the version used will correspond to the version of the Python interpreter used to start the run.
1429
- disabled : bool, default False
1430
- If set to True, disables Conda.
1431
- """
1432
- ...
1433
-
1434
1290
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1435
1291
  """
1436
1292
  Allows setting external datastores to save data for the
@@ -1545,5 +1401,149 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1545
1401
  """
1546
1402
  ...
1547
1403
 
1404
+ 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]]:
1405
+ """
1406
+ 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.
1407
+ 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.
1408
+
1409
+
1410
+ Parameters
1411
+ ----------
1412
+ timeout : int
1413
+ Time, in seconds before the task times out and fails. (Default: 3600)
1414
+ poke_interval : int
1415
+ Time in seconds that the job should wait in between each try. (Default: 60)
1416
+ mode : str
1417
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1418
+ exponential_backoff : bool
1419
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1420
+ pool : str
1421
+ the slot pool this task should run in,
1422
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1423
+ soft_fail : bool
1424
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1425
+ name : str
1426
+ Name of the sensor on Airflow
1427
+ description : str
1428
+ Description of sensor in the Airflow UI
1429
+ external_dag_id : str
1430
+ The dag_id that contains the task you want to wait for.
1431
+ external_task_ids : List[str]
1432
+ The list of task_ids that you want to wait for.
1433
+ If None (default value) the sensor waits for the DAG. (Default: None)
1434
+ allowed_states : List[str]
1435
+ Iterable of allowed states, (Default: ['success'])
1436
+ failed_states : List[str]
1437
+ Iterable of failed or dis-allowed states. (Default: None)
1438
+ execution_delta : datetime.timedelta
1439
+ time difference with the previous execution to look at,
1440
+ the default is the same logical date as the current task or DAG. (Default: None)
1441
+ check_existence: bool
1442
+ Set to True to check if the external task exists or check if
1443
+ the DAG to wait for exists. (Default: True)
1444
+ """
1445
+ ...
1446
+
1447
+ @typing.overload
1448
+ 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]]:
1449
+ """
1450
+ Specifies the flow(s) that this flow depends on.
1451
+
1452
+ ```
1453
+ @trigger_on_finish(flow='FooFlow')
1454
+ ```
1455
+ or
1456
+ ```
1457
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1458
+ ```
1459
+ This decorator respects the @project decorator and triggers the flow
1460
+ when upstream runs within the same namespace complete successfully
1461
+
1462
+ Additionally, you can specify project aware upstream flow dependencies
1463
+ by specifying the fully qualified project_flow_name.
1464
+ ```
1465
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1466
+ ```
1467
+ or
1468
+ ```
1469
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1470
+ ```
1471
+
1472
+ You can also specify just the project or project branch (other values will be
1473
+ inferred from the current project or project branch):
1474
+ ```
1475
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1476
+ ```
1477
+
1478
+ Note that `branch` is typically one of:
1479
+ - `prod`
1480
+ - `user.bob`
1481
+ - `test.my_experiment`
1482
+ - `prod.staging`
1483
+
1484
+
1485
+ Parameters
1486
+ ----------
1487
+ flow : Union[str, Dict[str, str]], optional, default None
1488
+ Upstream flow dependency for this flow.
1489
+ flows : List[Union[str, Dict[str, str]]], default []
1490
+ Upstream flow dependencies for this flow.
1491
+ options : Dict[str, Any], default {}
1492
+ Backend-specific configuration for tuning eventing behavior.
1493
+ """
1494
+ ...
1495
+
1496
+ @typing.overload
1497
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1498
+ ...
1499
+
1500
+ 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] = {}):
1501
+ """
1502
+ Specifies the flow(s) that this flow depends on.
1503
+
1504
+ ```
1505
+ @trigger_on_finish(flow='FooFlow')
1506
+ ```
1507
+ or
1508
+ ```
1509
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1510
+ ```
1511
+ This decorator respects the @project decorator and triggers the flow
1512
+ when upstream runs within the same namespace complete successfully
1513
+
1514
+ Additionally, you can specify project aware upstream flow dependencies
1515
+ by specifying the fully qualified project_flow_name.
1516
+ ```
1517
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1518
+ ```
1519
+ or
1520
+ ```
1521
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1522
+ ```
1523
+
1524
+ You can also specify just the project or project branch (other values will be
1525
+ inferred from the current project or project branch):
1526
+ ```
1527
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1528
+ ```
1529
+
1530
+ Note that `branch` is typically one of:
1531
+ - `prod`
1532
+ - `user.bob`
1533
+ - `test.my_experiment`
1534
+ - `prod.staging`
1535
+
1536
+
1537
+ Parameters
1538
+ ----------
1539
+ flow : Union[str, Dict[str, str]], optional, default None
1540
+ Upstream flow dependency for this flow.
1541
+ flows : List[Union[str, Dict[str, str]]], default []
1542
+ Upstream flow dependencies for this flow.
1543
+ options : Dict[str, Any], default {}
1544
+ Backend-specific configuration for tuning eventing behavior.
1545
+ """
1546
+ ...
1547
+
1548
1548
  pkg_name: str
1549
1549