ob-metaflow-stubs 6.0.3.186__py2.py3-none-any.whl → 6.0.3.188__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 (222) hide show
  1. metaflow-stubs/__init__.pyi +878 -862
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +6 -6
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +4 -4
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +3 -3
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +39 -39
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  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 +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +4 -4
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +5 -5
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +4 -4
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +3 -3
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +3 -3
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +3 -3
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +3 -3
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  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 +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +5 -5
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  86. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  88. metaflow-stubs/multicore_utils.pyi +2 -2
  89. metaflow-stubs/ob_internal.pyi +2 -2
  90. metaflow-stubs/parameters.pyi +4 -4
  91. metaflow-stubs/plugins/__init__.pyi +14 -12
  92. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  93. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  94. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  95. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  96. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  97. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  98. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  99. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  100. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  101. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  102. metaflow-stubs/plugins/argo/argo_workflows.pyi +6 -33
  103. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  104. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  105. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +4 -4
  106. metaflow-stubs/plugins/argo/exit_hooks.pyi +45 -0
  107. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  108. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  109. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  110. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  111. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  112. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  113. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  114. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  115. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  116. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  117. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  118. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  119. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  120. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  121. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  122. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  123. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  124. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  125. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  126. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  127. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  128. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  129. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  130. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  132. metaflow-stubs/plugins/cards/card_creator.pyi +4 -3
  133. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  134. metaflow-stubs/plugins/cards/card_decorator.pyi +13 -2
  135. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  136. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  137. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  138. metaflow-stubs/plugins/cards/card_modules/components.pyi +5 -5
  139. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  140. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  141. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +10 -2
  142. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  143. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  144. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  145. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  146. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  147. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  148. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  149. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  150. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  151. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  152. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  153. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  154. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  155. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  156. metaflow-stubs/plugins/exit_hook/__init__.pyi +11 -0
  157. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +20 -0
  158. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  159. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  160. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  161. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  162. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  163. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  164. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  165. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  166. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  167. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  168. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  169. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  170. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  171. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  172. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  173. metaflow-stubs/plugins/ollama/__init__.pyi +3 -3
  174. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  175. metaflow-stubs/plugins/perimeters.pyi +2 -2
  176. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  177. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  178. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  179. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  180. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  181. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  182. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  183. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  184. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  185. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  186. metaflow-stubs/plugins/secrets/__init__.pyi +7 -3
  187. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  188. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +9 -49
  189. metaflow-stubs/plugins/secrets/secrets_func.pyi +31 -0
  190. metaflow-stubs/plugins/secrets/secrets_spec.pyi +42 -0
  191. metaflow-stubs/plugins/secrets/utils.pyi +28 -0
  192. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  193. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  194. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  195. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  196. metaflow-stubs/plugins/torchtune/__init__.pyi +2 -2
  197. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  198. metaflow-stubs/plugins/uv/uv_environment.pyi +3 -3
  199. metaflow-stubs/profilers/__init__.pyi +2 -2
  200. metaflow-stubs/pylint_wrapper.pyi +2 -2
  201. metaflow-stubs/runner/__init__.pyi +2 -2
  202. metaflow-stubs/runner/deployer.pyi +29 -29
  203. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  204. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  205. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  206. metaflow-stubs/runner/nbrun.pyi +2 -2
  207. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  208. metaflow-stubs/runner/utils.pyi +3 -3
  209. metaflow-stubs/system/__init__.pyi +2 -2
  210. metaflow-stubs/system/system_logger.pyi +2 -2
  211. metaflow-stubs/system/system_monitor.pyi +2 -2
  212. metaflow-stubs/tagging_util.pyi +2 -2
  213. metaflow-stubs/tuple_util.pyi +2 -2
  214. metaflow-stubs/user_configs/__init__.pyi +2 -2
  215. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  216. metaflow-stubs/user_configs/config_options.pyi +3 -3
  217. metaflow-stubs/user_configs/config_parameters.pyi +7 -7
  218. {ob_metaflow_stubs-6.0.3.186.dist-info → ob_metaflow_stubs-6.0.3.188.dist-info}/METADATA +1 -1
  219. ob_metaflow_stubs-6.0.3.188.dist-info/RECORD +222 -0
  220. ob_metaflow_stubs-6.0.3.186.dist-info/RECORD +0 -216
  221. {ob_metaflow_stubs-6.0.3.186.dist-info → ob_metaflow_stubs-6.0.3.188.dist-info}/WHEEL +0 -0
  222. {ob_metaflow_stubs-6.0.3.186.dist-info → ob_metaflow_stubs-6.0.3.188.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.18.1+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-07-01T08:40:33.108391 #
3
+ # MF version: 2.15.21.1+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-07-11T16:11:55.172406 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -44,9 +44,9 @@ 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 pyproject_toml_parser as pyproject_toml_parser
49
47
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
52
52
  from .client.core import get_namespace as get_namespace
@@ -156,324 +156,419 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
156
156
  ...
157
157
 
158
158
  @typing.overload
159
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
159
+ def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
160
160
  """
161
- Specifies the PyPI packages for the step.
162
-
163
- Information in this decorator will augment any
164
- attributes set in the `@pyi_base` flow-level decorator. Hence,
165
- you can use `@pypi_base` to set packages required by all
166
- steps and use `@pypi` to specify step-specific overrides.
161
+ Specifies environment variables to be set prior to the execution of a step.
167
162
 
168
163
 
169
164
  Parameters
170
165
  ----------
171
- packages : Dict[str, str], default: {}
172
- Packages to use for this step. The key is the name of the package
173
- and the value is the version to use.
174
- python : str, optional, default: None
175
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
176
- that the version used will correspond to the version of the Python interpreter used to start the run.
166
+ vars : Dict[str, str], default {}
167
+ Dictionary of environment variables to set.
177
168
  """
178
169
  ...
179
170
 
180
171
  @typing.overload
181
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
172
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
182
173
  ...
183
174
 
184
175
  @typing.overload
185
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
176
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
186
177
  ...
187
178
 
188
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
179
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
189
180
  """
190
- Specifies the PyPI packages for the step.
191
-
192
- Information in this decorator will augment any
193
- attributes set in the `@pyi_base` flow-level decorator. Hence,
194
- you can use `@pypi_base` to set packages required by all
195
- steps and use `@pypi` to specify step-specific overrides.
181
+ Specifies environment variables to be set prior to the execution of a step.
196
182
 
197
183
 
198
184
  Parameters
199
185
  ----------
200
- packages : Dict[str, str], default: {}
201
- Packages to use for this step. The key is the name of the package
202
- and the value is the version to use.
203
- python : str, optional, default: None
204
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
205
- that the version used will correspond to the version of the Python interpreter used to start the run.
206
- """
207
- ...
208
-
209
- @typing.overload
210
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
211
- """
212
- Internal decorator to support Fast bakery
186
+ vars : Dict[str, str], default {}
187
+ Dictionary of environment variables to set.
213
188
  """
214
189
  ...
215
190
 
216
- @typing.overload
217
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
218
- ...
219
-
220
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
191
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
221
192
  """
222
- Internal decorator to support Fast bakery
193
+ Specifies that this step is used to deploy an instance of the app.
194
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
195
+
196
+
197
+ Parameters
198
+ ----------
199
+ app_port : int
200
+ Number of GPUs to use.
201
+ app_name : str
202
+ Name of the app to deploy.
223
203
  """
224
204
  ...
225
205
 
226
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
206
+ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card_refresh_interval: int, max_retries: int, retry_alert_frequency: int, engine_args: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
227
207
  """
228
- Decorator that helps cache, version and store models/datasets from huggingface hub.
229
-
230
- > Examples
231
-
232
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
233
- ```python
234
- @huggingface_hub
235
- @step
236
- def pull_model_from_huggingface(self):
237
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
238
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
239
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
240
- # value of the function is a reference to the model in the backend storage.
241
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
208
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
242
209
 
243
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
244
- self.llama_model = current.huggingface_hub.snapshot_download(
245
- repo_id=self.model_id,
246
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
247
- )
248
- self.next(self.train)
249
- ```
210
+ User code call
211
+ --------------
212
+ @vllm(
213
+ model="...",
214
+ ...
215
+ )
250
216
 
251
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
252
- ```python
253
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
254
- @step
255
- def pull_model_from_huggingface(self):
256
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
257
- ```
217
+ Valid backend options
218
+ ---------------------
219
+ - 'local': Run as a separate process on the local task machine.
258
220
 
259
- ```python
260
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
261
- @step
262
- def finetune_model(self):
263
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
264
- # path_to_model will be /my-directory
265
- ```
221
+ Valid model options
222
+ -------------------
223
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
266
224
 
267
- ```python
268
- # Takes all the arguments passed to `snapshot_download`
269
- # except for `local_dir`
270
- @huggingface_hub(load=[
271
- {
272
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
273
- },
274
- {
275
- "repo_id": "myorg/mistral-lora",
276
- "repo_type": "model",
277
- },
278
- ])
279
- @step
280
- def finetune_model(self):
281
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
282
- # path_to_model will be /my-directory
283
- ```
225
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
226
+ If you need multiple models, you must create multiple @vllm decorators.
284
227
 
285
228
 
286
229
  Parameters
287
230
  ----------
288
- temp_dir_root : str, optional
289
- The root directory that will hold the temporary directory where objects will be downloaded.
290
-
291
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
292
- The list of repos (models/datasets) to load.
293
-
294
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
295
-
296
- - If repo (model/dataset) is not found in the datastore:
297
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
298
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
299
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
300
-
301
- - If repo is found in the datastore:
302
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
231
+ model: str
232
+ HuggingFace model identifier to be served by vLLM.
233
+ backend: str
234
+ Determines where and how to run the vLLM process.
235
+ openai_api_server: bool
236
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
237
+ Default is False (uses native engine).
238
+ Set to True for backward compatibility with existing code.
239
+ debug: bool
240
+ Whether to turn on verbose debugging logs.
241
+ card_refresh_interval: int
242
+ Interval in seconds for refreshing the vLLM status card.
243
+ Only used when openai_api_server=True.
244
+ max_retries: int
245
+ Maximum number of retries checking for vLLM server startup.
246
+ Only used when openai_api_server=True.
247
+ retry_alert_frequency: int
248
+ Frequency of alert logs for vLLM server startup retries.
249
+ Only used when openai_api_server=True.
250
+ engine_args : dict
251
+ Additional keyword arguments to pass to the vLLM engine.
252
+ For example, `tensor_parallel_size=2`.
303
253
  """
304
254
  ...
305
255
 
306
256
  @typing.overload
307
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
257
+ 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]]]:
308
258
  """
309
- Specifies secrets to be retrieved and injected as environment variables prior to
310
- the execution of a step.
259
+ Specifies the Conda environment for the step.
260
+
261
+ Information in this decorator will augment any
262
+ attributes set in the `@conda_base` flow-level decorator. Hence,
263
+ you can use `@conda_base` to set packages required by all
264
+ steps and use `@conda` to specify step-specific overrides.
311
265
 
312
266
 
313
267
  Parameters
314
268
  ----------
315
- sources : List[Union[str, Dict[str, Any]]], default: []
316
- List of secret specs, defining how the secrets are to be retrieved
269
+ packages : Dict[str, str], default {}
270
+ Packages to use for this step. The key is the name of the package
271
+ and the value is the version to use.
272
+ libraries : Dict[str, str], default {}
273
+ Supported for backward compatibility. When used with packages, packages will take precedence.
274
+ python : str, optional, default None
275
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
276
+ that the version used will correspond to the version of the Python interpreter used to start the run.
277
+ disabled : bool, default False
278
+ If set to True, disables @conda.
317
279
  """
318
280
  ...
319
281
 
320
282
  @typing.overload
321
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
283
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
322
284
  ...
323
285
 
324
286
  @typing.overload
325
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
287
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
326
288
  ...
327
289
 
328
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
290
+ 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):
329
291
  """
330
- Specifies secrets to be retrieved and injected as environment variables prior to
331
- the execution of a step.
292
+ Specifies the Conda environment for the step.
293
+
294
+ Information in this decorator will augment any
295
+ attributes set in the `@conda_base` flow-level decorator. Hence,
296
+ you can use `@conda_base` to set packages required by all
297
+ steps and use `@conda` to specify step-specific overrides.
332
298
 
333
299
 
334
300
  Parameters
335
301
  ----------
336
- sources : List[Union[str, Dict[str, Any]]], default: []
337
- List of secret specs, defining how the secrets are to be retrieved
302
+ packages : Dict[str, str], default {}
303
+ Packages to use for this step. The key is the name of the package
304
+ and the value is the version to use.
305
+ libraries : Dict[str, str], default {}
306
+ Supported for backward compatibility. When used with packages, packages will take precedence.
307
+ python : str, optional, default None
308
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
309
+ that the version used will correspond to the version of the Python interpreter used to start the run.
310
+ disabled : bool, default False
311
+ If set to True, disables @conda.
338
312
  """
339
313
  ...
340
314
 
341
315
  @typing.overload
342
- 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]]]:
316
+ 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]]]:
343
317
  """
344
- Specifies that the step will success under all circumstances.
318
+ Creates a human-readable report, a Metaflow Card, after this step completes.
345
319
 
346
- The decorator will create an optional artifact, specified by `var`, which
347
- contains the exception raised. You can use it to detect the presence
348
- of errors, indicating that all happy-path artifacts produced by the step
349
- are missing.
320
+ Note that you may add multiple `@card` decorators in a step with different parameters.
350
321
 
351
322
 
352
323
  Parameters
353
324
  ----------
354
- var : str, optional, default None
355
- Name of the artifact in which to store the caught exception.
356
- If not specified, the exception is not stored.
357
- print_exception : bool, default True
358
- Determines whether or not the exception is printed to
359
- stdout when caught.
325
+ type : str, default 'default'
326
+ Card type.
327
+ id : str, optional, default None
328
+ If multiple cards are present, use this id to identify this card.
329
+ options : Dict[str, Any], default {}
330
+ Options passed to the card. The contents depend on the card type.
331
+ timeout : int, default 45
332
+ Interrupt reporting if it takes more than this many seconds.
360
333
  """
361
334
  ...
362
335
 
363
336
  @typing.overload
364
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
337
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
365
338
  ...
366
339
 
367
340
  @typing.overload
368
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
341
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
369
342
  ...
370
343
 
371
- 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):
344
+ 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):
372
345
  """
373
- Specifies that the step will success under all circumstances.
346
+ Creates a human-readable report, a Metaflow Card, after this step completes.
374
347
 
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.
348
+ Note that you may add multiple `@card` decorators in a step with different parameters.
379
349
 
380
350
 
381
351
  Parameters
382
352
  ----------
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.
353
+ type : str, default 'default'
354
+ Card type.
355
+ id : str, optional, default None
356
+ If multiple cards are present, use this id to identify this card.
357
+ options : Dict[str, Any], default {}
358
+ Options passed to the card. The contents depend on the card type.
359
+ timeout : int, default 45
360
+ Interrupt reporting if it takes more than this many seconds.
389
361
  """
390
362
  ...
391
363
 
392
- @typing.overload
393
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
364
+ 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]]]:
394
365
  """
395
- Specifies a timeout for your step.
366
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
396
367
 
397
- This decorator is useful if this step may hang indefinitely.
368
+ User code call
369
+ --------------
370
+ @ollama(
371
+ models=[...],
372
+ ...
373
+ )
398
374
 
399
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
400
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
401
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
375
+ Valid backend options
376
+ ---------------------
377
+ - 'local': Run as a separate process on the local task machine.
378
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
379
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
402
380
 
403
- Note that all the values specified in parameters are added together so if you specify
404
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
381
+ Valid model options
382
+ -------------------
383
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
405
384
 
406
385
 
407
386
  Parameters
408
387
  ----------
409
- seconds : int, default 0
410
- Number of seconds to wait prior to timing out.
411
- minutes : int, default 0
412
- Number of minutes to wait prior to timing out.
413
- hours : int, default 0
414
- Number of hours to wait prior to timing out.
388
+ models: list[str]
389
+ List of Ollama containers running models in sidecars.
390
+ backend: str
391
+ Determines where and how to run the Ollama process.
392
+ force_pull: bool
393
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
394
+ cache_update_policy: str
395
+ Cache update policy: "auto", "force", or "never".
396
+ force_cache_update: bool
397
+ Simple override for "force" cache update policy.
398
+ debug: bool
399
+ Whether to turn on verbose debugging logs.
400
+ circuit_breaker_config: dict
401
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
402
+ timeout_config: dict
403
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
415
404
  """
416
405
  ...
417
406
 
418
407
  @typing.overload
419
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
420
- ...
421
-
422
- @typing.overload
423
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
424
- ...
425
-
426
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
408
+ 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]]]:
427
409
  """
428
- Specifies a timeout for your step.
429
-
430
- This decorator is useful if this step may hang indefinitely.
410
+ Specifies the resources needed when executing this step.
431
411
 
432
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
433
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
434
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
412
+ Use `@resources` to specify the resource requirements
413
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
435
414
 
436
- Note that all the values specified in parameters are added together so if you specify
437
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
415
+ You can choose the compute layer on the command line by executing e.g.
416
+ ```
417
+ python myflow.py run --with batch
418
+ ```
419
+ or
420
+ ```
421
+ python myflow.py run --with kubernetes
422
+ ```
423
+ which executes the flow on the desired system using the
424
+ requirements specified in `@resources`.
438
425
 
439
426
 
440
427
  Parameters
441
428
  ----------
442
- seconds : int, default 0
443
- Number of seconds to wait prior to timing out.
444
- minutes : int, default 0
445
- Number of minutes to wait prior to timing out.
446
- hours : int, default 0
447
- Number of hours to wait prior to timing out.
429
+ cpu : int, default 1
430
+ Number of CPUs required for this step.
431
+ gpu : int, optional, default None
432
+ Number of GPUs required for this step.
433
+ disk : int, optional, default None
434
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
435
+ memory : int, default 4096
436
+ Memory size (in MB) required for this step.
437
+ shared_memory : int, optional, default None
438
+ The value for the size (in MiB) of the /dev/shm volume for this step.
439
+ This parameter maps to the `--shm-size` option in Docker.
448
440
  """
449
441
  ...
450
442
 
451
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
443
+ @typing.overload
444
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
445
+ ...
446
+
447
+ @typing.overload
448
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
449
+ ...
450
+
451
+ 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):
452
452
  """
453
- Specifies that this step should execute on DGX cloud.
453
+ Specifies the resources needed when executing this step.
454
+
455
+ Use `@resources` to specify the resource requirements
456
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
457
+
458
+ You can choose the compute layer on the command line by executing e.g.
459
+ ```
460
+ python myflow.py run --with batch
461
+ ```
462
+ or
463
+ ```
464
+ python myflow.py run --with kubernetes
465
+ ```
466
+ which executes the flow on the desired system using the
467
+ requirements specified in `@resources`.
454
468
 
455
469
 
456
470
  Parameters
457
471
  ----------
458
- gpu : int
459
- Number of GPUs to use.
460
- gpu_type : str
461
- Type of Nvidia GPU to use.
472
+ cpu : int, default 1
473
+ Number of CPUs required for this step.
474
+ gpu : int, optional, default None
475
+ Number of GPUs required for this step.
476
+ disk : int, optional, default None
477
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
478
+ memory : int, default 4096
479
+ Memory size (in MB) required for this step.
480
+ shared_memory : int, optional, default None
481
+ The value for the size (in MiB) of the /dev/shm volume for this step.
482
+ This parameter maps to the `--shm-size` option in Docker.
462
483
  """
463
484
  ...
464
485
 
465
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
486
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
466
487
  """
467
- Specifies that this step is used to deploy an instance of the app.
468
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
488
+ Specifies that this step should execute on Kubernetes.
469
489
 
470
490
 
471
491
  Parameters
472
492
  ----------
473
- app_port : int
474
- Number of GPUs to use.
475
- app_name : str
476
- Name of the app to deploy.
493
+ cpu : int, default 1
494
+ Number of CPUs required for this step. If `@resources` is
495
+ also present, the maximum value from all decorators is used.
496
+ memory : int, default 4096
497
+ Memory size (in MB) required for this step. If
498
+ `@resources` is also present, the maximum value from all decorators is
499
+ used.
500
+ disk : int, default 10240
501
+ Disk size (in MB) required for this step. If
502
+ `@resources` is also present, the maximum value from all decorators is
503
+ used.
504
+ image : str, optional, default None
505
+ Docker image to use when launching on Kubernetes. If not specified, and
506
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
507
+ not, a default Docker image mapping to the current version of Python is used.
508
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
509
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
510
+ image_pull_secrets: List[str], default []
511
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
512
+ Kubernetes image pull secrets to use when pulling container images
513
+ in Kubernetes.
514
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
515
+ Kubernetes service account to use when launching pod in Kubernetes.
516
+ secrets : List[str], optional, default None
517
+ Kubernetes secrets to use when launching pod in Kubernetes. These
518
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
519
+ in Metaflow configuration.
520
+ node_selector: Union[Dict[str,str], str], optional, default None
521
+ Kubernetes node selector(s) to apply to the pod running the task.
522
+ Can be passed in as a comma separated string of values e.g.
523
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
524
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
525
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
526
+ Kubernetes namespace to use when launching pod in Kubernetes.
527
+ gpu : int, optional, default None
528
+ Number of GPUs required for this step. A value of zero implies that
529
+ the scheduled node should not have GPUs.
530
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
531
+ The vendor of the GPUs to be used for this step.
532
+ tolerations : List[str], default []
533
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
534
+ Kubernetes tolerations to use when launching pod in Kubernetes.
535
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
536
+ Kubernetes labels to use when launching pod in Kubernetes.
537
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
538
+ Kubernetes annotations to use when launching pod in Kubernetes.
539
+ use_tmpfs : bool, default False
540
+ This enables an explicit tmpfs mount for this step.
541
+ tmpfs_tempdir : bool, default True
542
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
543
+ tmpfs_size : int, optional, default: None
544
+ The value for the size (in MiB) of the tmpfs mount for this step.
545
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
546
+ memory allocated for this step.
547
+ tmpfs_path : str, optional, default /metaflow_temp
548
+ Path to tmpfs mount for this step.
549
+ persistent_volume_claims : Dict[str, str], optional, default None
550
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
551
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
552
+ shared_memory: int, optional
553
+ Shared memory size (in MiB) required for this step
554
+ port: int, optional
555
+ Port number to specify in the Kubernetes job object
556
+ compute_pool : str, optional, default None
557
+ Compute pool to be used for for this step.
558
+ If not specified, any accessible compute pool within the perimeter is used.
559
+ hostname_resolution_timeout: int, default 10 * 60
560
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
561
+ Only applicable when @parallel is used.
562
+ qos: str, default: Burstable
563
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
564
+
565
+ security_context: Dict[str, Any], optional, default None
566
+ Container security context. Applies to the task container. Allows the following keys:
567
+ - privileged: bool, optional, default None
568
+ - allow_privilege_escalation: bool, optional, default None
569
+ - run_as_user: int, optional, default None
570
+ - run_as_group: int, optional, default None
571
+ - run_as_non_root: bool, optional, default None
477
572
  """
478
573
  ...
479
574
 
@@ -493,46 +588,83 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
493
588
  """
494
589
  ...
495
590
 
496
- 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]]]:
591
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
497
592
  """
498
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
593
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
499
594
 
500
- User code call
501
- --------------
502
- @ollama(
503
- models=[...],
504
- ...
505
- )
595
+ > Examples
506
596
 
507
- Valid backend options
508
- ---------------------
509
- - 'local': Run as a separate process on the local task machine.
510
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
511
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
597
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
598
+ ```python
599
+ @huggingface_hub
600
+ @step
601
+ def pull_model_from_huggingface(self):
602
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
603
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
604
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
605
+ # value of the function is a reference to the model in the backend storage.
606
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
512
607
 
513
- Valid model options
514
- -------------------
515
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
608
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
609
+ self.llama_model = current.huggingface_hub.snapshot_download(
610
+ repo_id=self.model_id,
611
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
612
+ )
613
+ self.next(self.train)
614
+ ```
615
+
616
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
617
+ ```python
618
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
619
+ @step
620
+ def pull_model_from_huggingface(self):
621
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
622
+ ```
623
+
624
+ ```python
625
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
626
+ @step
627
+ def finetune_model(self):
628
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
629
+ # path_to_model will be /my-directory
630
+ ```
631
+
632
+ ```python
633
+ # Takes all the arguments passed to `snapshot_download`
634
+ # except for `local_dir`
635
+ @huggingface_hub(load=[
636
+ {
637
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
638
+ },
639
+ {
640
+ "repo_id": "myorg/mistral-lora",
641
+ "repo_type": "model",
642
+ },
643
+ ])
644
+ @step
645
+ def finetune_model(self):
646
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
647
+ # path_to_model will be /my-directory
648
+ ```
516
649
 
517
650
 
518
651
  Parameters
519
652
  ----------
520
- models: list[str]
521
- List of Ollama containers running models in sidecars.
522
- backend: str
523
- Determines where and how to run the Ollama process.
524
- force_pull: bool
525
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
526
- cache_update_policy: str
527
- Cache update policy: "auto", "force", or "never".
528
- force_cache_update: bool
529
- Simple override for "force" cache update policy.
530
- debug: bool
531
- Whether to turn on verbose debugging logs.
532
- circuit_breaker_config: dict
533
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
534
- timeout_config: dict
535
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
653
+ temp_dir_root : str, optional
654
+ The root directory that will hold the temporary directory where objects will be downloaded.
655
+
656
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
657
+ The list of repos (models/datasets) to load.
658
+
659
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
660
+
661
+ - If repo (model/dataset) is not found in the datastore:
662
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
663
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
664
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
665
+
666
+ - If repo is found in the datastore:
667
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
536
668
  """
537
669
  ...
538
670
 
@@ -683,391 +815,234 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
683
815
  """
684
816
  ...
685
817
 
686
- def vllm(*, model: str, backend: str, debug: bool, kwargs: typing.Any) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
687
- """
688
- This decorator is used to run vllm APIs as Metaflow task sidecars.
689
-
690
- User code call
691
- --------------
692
- @vllm(
693
- model="...",
694
- ...
695
- )
696
-
697
- Valid backend options
698
- ---------------------
699
- - 'local': Run as a separate process on the local task machine.
700
-
701
- Valid model options
702
- -------------------
703
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
704
-
705
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
706
- If you need multiple models, you must create multiple @vllm decorators.
707
-
708
-
709
- Parameters
710
- ----------
711
- model: str
712
- HuggingFace model identifier to be served by vLLM.
713
- backend: str
714
- Determines where and how to run the vLLM process.
715
- debug: bool
716
- Whether to turn on verbose debugging logs.
717
- kwargs : Any
718
- Any other keyword arguments are passed directly to the vLLM engine.
719
- This allows for flexible configuration of vLLM server settings.
720
- For example, `tensor_parallel_size=2`.
721
- """
722
- ...
723
-
724
818
  @typing.overload
725
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
819
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
726
820
  """
727
- Specifies the resources needed when executing this step.
728
-
729
- Use `@resources` to specify the resource requirements
730
- independently of the specific compute layer (`@batch`, `@kubernetes`).
731
-
732
- You can choose the compute layer on the command line by executing e.g.
733
- ```
734
- python myflow.py run --with batch
735
- ```
736
- or
737
- ```
738
- python myflow.py run --with kubernetes
739
- ```
740
- which executes the flow on the desired system using the
741
- requirements specified in `@resources`.
742
-
743
-
744
- Parameters
745
- ----------
746
- cpu : int, default 1
747
- Number of CPUs required for this step.
748
- gpu : int, optional, default None
749
- Number of GPUs required for this step.
750
- disk : int, optional, default None
751
- Disk size (in MB) required for this step. Only applies on Kubernetes.
752
- memory : int, default 4096
753
- Memory size (in MB) required for this step.
754
- shared_memory : int, optional, default None
755
- The value for the size (in MiB) of the /dev/shm volume for this step.
756
- This parameter maps to the `--shm-size` option in Docker.
821
+ Internal decorator to support Fast bakery
757
822
  """
758
823
  ...
759
824
 
760
825
  @typing.overload
761
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
762
- ...
763
-
764
- @typing.overload
765
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
826
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
766
827
  ...
767
-
768
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
769
- """
770
- Specifies the resources needed when executing this step.
771
-
772
- Use `@resources` to specify the resource requirements
773
- independently of the specific compute layer (`@batch`, `@kubernetes`).
774
-
775
- You can choose the compute layer on the command line by executing e.g.
776
- ```
777
- python myflow.py run --with batch
778
- ```
779
- or
780
- ```
781
- python myflow.py run --with kubernetes
782
- ```
783
- which executes the flow on the desired system using the
784
- requirements specified in `@resources`.
785
-
786
-
787
- Parameters
788
- ----------
789
- cpu : int, default 1
790
- Number of CPUs required for this step.
791
- gpu : int, optional, default None
792
- Number of GPUs required for this step.
793
- disk : int, optional, default None
794
- Disk size (in MB) required for this step. Only applies on Kubernetes.
795
- memory : int, default 4096
796
- Memory size (in MB) required for this step.
797
- shared_memory : int, optional, default None
798
- The value for the size (in MiB) of the /dev/shm volume for this step.
799
- This parameter maps to the `--shm-size` option in Docker.
828
+
829
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
830
+ """
831
+ Internal decorator to support Fast bakery
800
832
  """
801
833
  ...
802
834
 
803
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
835
+ 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]]]:
804
836
  """
805
- Specifies that this step should execute on Kubernetes.
837
+ Specifies that this step should execute on DGX cloud.
806
838
 
807
839
 
808
840
  Parameters
809
841
  ----------
810
- cpu : int, default 1
811
- Number of CPUs required for this step. If `@resources` is
812
- also present, the maximum value from all decorators is used.
813
- memory : int, default 4096
814
- Memory size (in MB) required for this step. If
815
- `@resources` is also present, the maximum value from all decorators is
816
- used.
817
- disk : int, default 10240
818
- Disk size (in MB) required for this step. If
819
- `@resources` is also present, the maximum value from all decorators is
820
- used.
821
- image : str, optional, default None
822
- Docker image to use when launching on Kubernetes. If not specified, and
823
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
824
- not, a default Docker image mapping to the current version of Python is used.
825
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
826
- If given, the imagePullPolicy to be applied to the Docker image of the step.
827
- image_pull_secrets: List[str], default []
828
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
829
- Kubernetes image pull secrets to use when pulling container images
830
- in Kubernetes.
831
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
832
- Kubernetes service account to use when launching pod in Kubernetes.
833
- secrets : List[str], optional, default None
834
- Kubernetes secrets to use when launching pod in Kubernetes. These
835
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
836
- in Metaflow configuration.
837
- node_selector: Union[Dict[str,str], str], optional, default None
838
- Kubernetes node selector(s) to apply to the pod running the task.
839
- Can be passed in as a comma separated string of values e.g.
840
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
841
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
842
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
843
- Kubernetes namespace to use when launching pod in Kubernetes.
844
- gpu : int, optional, default None
845
- Number of GPUs required for this step. A value of zero implies that
846
- the scheduled node should not have GPUs.
847
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
848
- The vendor of the GPUs to be used for this step.
849
- tolerations : List[str], default []
850
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
851
- Kubernetes tolerations to use when launching pod in Kubernetes.
852
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
853
- Kubernetes labels to use when launching pod in Kubernetes.
854
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
855
- Kubernetes annotations to use when launching pod in Kubernetes.
856
- use_tmpfs : bool, default False
857
- This enables an explicit tmpfs mount for this step.
858
- tmpfs_tempdir : bool, default True
859
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
860
- tmpfs_size : int, optional, default: None
861
- The value for the size (in MiB) of the tmpfs mount for this step.
862
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
863
- memory allocated for this step.
864
- tmpfs_path : str, optional, default /metaflow_temp
865
- Path to tmpfs mount for this step.
866
- persistent_volume_claims : Dict[str, str], optional, default None
867
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
868
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
869
- shared_memory: int, optional
870
- Shared memory size (in MiB) required for this step
871
- port: int, optional
872
- Port number to specify in the Kubernetes job object
873
- compute_pool : str, optional, default None
874
- Compute pool to be used for for this step.
875
- If not specified, any accessible compute pool within the perimeter is used.
876
- hostname_resolution_timeout: int, default 10 * 60
877
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
878
- Only applicable when @parallel is used.
879
- qos: str, default: Burstable
880
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
881
-
882
- security_context: Dict[str, Any], optional, default None
883
- Container security context. Applies to the task container. Allows the following keys:
884
- - privileged: bool, optional, default None
885
- - allow_privilege_escalation: bool, optional, default None
886
- - run_as_user: int, optional, default None
887
- - run_as_group: int, optional, default None
888
- - run_as_non_root: bool, optional, default None
842
+ gpu : int
843
+ Number of GPUs to use.
844
+ gpu_type : str
845
+ Type of Nvidia GPU to use.
889
846
  """
890
847
  ...
891
848
 
892
849
  @typing.overload
893
- def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
850
+ 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]]]:
894
851
  """
895
- Creates a human-readable report, a Metaflow Card, after this step completes.
852
+ Specifies the PyPI packages for the step.
896
853
 
897
- Note that you may add multiple `@card` decorators in a step with different parameters.
854
+ Information in this decorator will augment any
855
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
856
+ you can use `@pypi_base` to set packages required by all
857
+ steps and use `@pypi` to specify step-specific overrides.
898
858
 
899
859
 
900
860
  Parameters
901
861
  ----------
902
- type : str, default 'default'
903
- Card type.
904
- id : str, optional, default None
905
- If multiple cards are present, use this id to identify this card.
906
- options : Dict[str, Any], default {}
907
- Options passed to the card. The contents depend on the card type.
908
- timeout : int, default 45
909
- Interrupt reporting if it takes more than this many seconds.
862
+ packages : Dict[str, str], default: {}
863
+ Packages to use for this step. The key is the name of the package
864
+ and the value is the version to use.
865
+ python : str, optional, default: None
866
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
867
+ that the version used will correspond to the version of the Python interpreter used to start the run.
910
868
  """
911
869
  ...
912
870
 
913
871
  @typing.overload
914
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
872
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
915
873
  ...
916
874
 
917
875
  @typing.overload
918
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
876
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
919
877
  ...
920
878
 
921
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
879
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
922
880
  """
923
- Creates a human-readable report, a Metaflow Card, after this step completes.
881
+ Specifies the PyPI packages for the step.
924
882
 
925
- Note that you may add multiple `@card` decorators in a step with different parameters.
883
+ Information in this decorator will augment any
884
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
885
+ you can use `@pypi_base` to set packages required by all
886
+ steps and use `@pypi` to specify step-specific overrides.
926
887
 
927
888
 
928
889
  Parameters
929
890
  ----------
930
- type : str, default 'default'
931
- Card type.
932
- id : str, optional, default None
933
- If multiple cards are present, use this id to identify this card.
934
- options : Dict[str, Any], default {}
935
- Options passed to the card. The contents depend on the card type.
936
- timeout : int, default 45
937
- Interrupt reporting if it takes more than this many seconds.
891
+ packages : Dict[str, str], default: {}
892
+ Packages to use for this step. The key is the name of the package
893
+ and the value is the version to use.
894
+ python : str, optional, default: None
895
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
896
+ that the version used will correspond to the version of the Python interpreter used to start the run.
938
897
  """
939
898
  ...
940
899
 
941
900
  @typing.overload
942
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
901
+ 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]]]:
943
902
  """
944
- Decorator prototype for all step decorators. This function gets specialized
945
- and imported for all decorators types by _import_plugin_decorators().
903
+ Specifies that the step will success under all circumstances.
904
+
905
+ The decorator will create an optional artifact, specified by `var`, which
906
+ contains the exception raised. You can use it to detect the presence
907
+ of errors, indicating that all happy-path artifacts produced by the step
908
+ are missing.
909
+
910
+
911
+ Parameters
912
+ ----------
913
+ var : str, optional, default None
914
+ Name of the artifact in which to store the caught exception.
915
+ If not specified, the exception is not stored.
916
+ print_exception : bool, default True
917
+ Determines whether or not the exception is printed to
918
+ stdout when caught.
946
919
  """
947
920
  ...
948
921
 
949
922
  @typing.overload
950
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
923
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
951
924
  ...
952
925
 
953
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
926
+ @typing.overload
927
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
928
+ ...
929
+
930
+ 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):
954
931
  """
955
- Decorator prototype for all step decorators. This function gets specialized
956
- and imported for all decorators types by _import_plugin_decorators().
932
+ Specifies that the step will success under all circumstances.
933
+
934
+ The decorator will create an optional artifact, specified by `var`, which
935
+ contains the exception raised. You can use it to detect the presence
936
+ of errors, indicating that all happy-path artifacts produced by the step
937
+ are missing.
938
+
939
+
940
+ Parameters
941
+ ----------
942
+ var : str, optional, default None
943
+ Name of the artifact in which to store the caught exception.
944
+ If not specified, the exception is not stored.
945
+ print_exception : bool, default True
946
+ Determines whether or not the exception is printed to
947
+ stdout when caught.
957
948
  """
958
949
  ...
959
950
 
960
951
  @typing.overload
961
- 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]]]:
952
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: 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]]]:
962
953
  """
963
- Specifies the number of times the task corresponding
964
- to a step needs to be retried.
965
-
966
- This decorator is useful for handling transient errors, such as networking issues.
967
- If your task contains operations that can't be retried safely, e.g. database updates,
968
- it is advisable to annotate it with `@retry(times=0)`.
969
-
970
- This can be used in conjunction with the `@catch` decorator. The `@catch`
971
- decorator will execute a no-op task after all retries have been exhausted,
972
- ensuring that the flow execution can continue.
954
+ Specifies secrets to be retrieved and injected as environment variables prior to
955
+ the execution of a step.
973
956
 
974
957
 
975
958
  Parameters
976
959
  ----------
977
- times : int, default 3
978
- Number of times to retry this task.
979
- minutes_between_retries : int, default 2
980
- Number of minutes between retries.
960
+ sources : List[Union[str, Dict[str, Any]]], default: []
961
+ List of secret specs, defining how the secrets are to be retrieved
962
+ role : str, optional, default: None
963
+ Role to use for fetching secrets
981
964
  """
982
965
  ...
983
966
 
984
967
  @typing.overload
985
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
968
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
986
969
  ...
987
970
 
988
971
  @typing.overload
989
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
972
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
990
973
  ...
991
974
 
992
- 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):
975
+ 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]]] = [], role: typing.Optional[str] = None):
993
976
  """
994
- Specifies the number of times the task corresponding
995
- to a step needs to be retried.
996
-
997
- This decorator is useful for handling transient errors, such as networking issues.
998
- If your task contains operations that can't be retried safely, e.g. database updates,
999
- it is advisable to annotate it with `@retry(times=0)`.
1000
-
1001
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1002
- decorator will execute a no-op task after all retries have been exhausted,
1003
- ensuring that the flow execution can continue.
977
+ Specifies secrets to be retrieved and injected as environment variables prior to
978
+ the execution of a step.
1004
979
 
1005
980
 
1006
981
  Parameters
1007
982
  ----------
1008
- times : int, default 3
1009
- Number of times to retry this task.
1010
- minutes_between_retries : int, default 2
1011
- Number of minutes between retries.
983
+ sources : List[Union[str, Dict[str, Any]]], default: []
984
+ List of secret specs, defining how the secrets are to be retrieved
985
+ role : str, optional, default: None
986
+ Role to use for fetching secrets
1012
987
  """
1013
988
  ...
1014
989
 
1015
990
  @typing.overload
1016
- 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]]]:
991
+ 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]]]:
1017
992
  """
1018
- Specifies the Conda environment for the step.
993
+ Specifies a timeout for your step.
1019
994
 
1020
- Information in this decorator will augment any
1021
- attributes set in the `@conda_base` flow-level decorator. Hence,
1022
- you can use `@conda_base` to set packages required by all
1023
- steps and use `@conda` to specify step-specific overrides.
995
+ This decorator is useful if this step may hang indefinitely.
996
+
997
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
998
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
999
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1000
+
1001
+ Note that all the values specified in parameters are added together so if you specify
1002
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1024
1003
 
1025
1004
 
1026
1005
  Parameters
1027
1006
  ----------
1028
- packages : Dict[str, str], default {}
1029
- Packages to use for this step. The key is the name of the package
1030
- and the value is the version to use.
1031
- libraries : Dict[str, str], default {}
1032
- Supported for backward compatibility. When used with packages, packages will take precedence.
1033
- python : str, optional, default None
1034
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1035
- that the version used will correspond to the version of the Python interpreter used to start the run.
1036
- disabled : bool, default False
1037
- If set to True, disables @conda.
1007
+ seconds : int, default 0
1008
+ Number of seconds to wait prior to timing out.
1009
+ minutes : int, default 0
1010
+ Number of minutes to wait prior to timing out.
1011
+ hours : int, default 0
1012
+ Number of hours to wait prior to timing out.
1038
1013
  """
1039
1014
  ...
1040
1015
 
1041
1016
  @typing.overload
1042
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1017
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1043
1018
  ...
1044
1019
 
1045
1020
  @typing.overload
1046
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1021
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1047
1022
  ...
1048
1023
 
1049
- 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):
1024
+ 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):
1050
1025
  """
1051
- Specifies the Conda environment for the step.
1026
+ Specifies a timeout for your step.
1052
1027
 
1053
- Information in this decorator will augment any
1054
- attributes set in the `@conda_base` flow-level decorator. Hence,
1055
- you can use `@conda_base` to set packages required by all
1056
- steps and use `@conda` to specify step-specific overrides.
1028
+ This decorator is useful if this step may hang indefinitely.
1029
+
1030
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1031
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1032
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1033
+
1034
+ Note that all the values specified in parameters are added together so if you specify
1035
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1057
1036
 
1058
1037
 
1059
1038
  Parameters
1060
1039
  ----------
1061
- packages : Dict[str, str], default {}
1062
- Packages to use for this step. The key is the name of the package
1063
- and the value is the version to use.
1064
- libraries : Dict[str, str], default {}
1065
- Supported for backward compatibility. When used with packages, packages will take precedence.
1066
- python : str, optional, default None
1067
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1068
- that the version used will correspond to the version of the Python interpreter used to start the run.
1069
- disabled : bool, default False
1070
- If set to True, disables @conda.
1040
+ seconds : int, default 0
1041
+ Number of seconds to wait prior to timing out.
1042
+ minutes : int, default 0
1043
+ Number of minutes to wait prior to timing out.
1044
+ hours : int, default 0
1045
+ Number of hours to wait prior to timing out.
1071
1046
  """
1072
1047
  ...
1073
1048
 
@@ -1166,325 +1141,111 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1166
1141
  @step
1167
1142
  def test(self):
1168
1143
  # `current.model.loaded` returns a dictionary of the loaded models
1169
- # where the key is the name of the artifact and the value is the path to the model
1170
- print(os.listdir(current.model.loaded["my_model"]))
1171
- self.next(self.end)
1172
- ```
1173
-
1174
- - Loading models
1175
- ```python
1176
- @step
1177
- def train(self):
1178
- # current.model.load returns the path to the model loaded
1179
- checkpoint_path = current.model.load(
1180
- self.checkpoint_key,
1181
- )
1182
- model_path = current.model.load(
1183
- self.model,
1184
- )
1185
- self.next(self.test)
1186
- ```
1187
-
1188
-
1189
- Parameters
1190
- ----------
1191
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1192
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1193
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1194
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1195
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1196
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1197
-
1198
- temp_dir_root : str, default: None
1199
- The root directory under which `current.model.loaded` will store loaded models
1200
- """
1201
- ...
1202
-
1203
- @typing.overload
1204
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1205
- """
1206
- Specifies environment variables to be set prior to the execution of a step.
1207
-
1208
-
1209
- Parameters
1210
- ----------
1211
- vars : Dict[str, str], default {}
1212
- Dictionary of environment variables to set.
1213
- """
1214
- ...
1215
-
1216
- @typing.overload
1217
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1218
- ...
1219
-
1220
- @typing.overload
1221
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1222
- ...
1223
-
1224
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1225
- """
1226
- Specifies environment variables to be set prior to the execution of a step.
1227
-
1228
-
1229
- Parameters
1230
- ----------
1231
- vars : Dict[str, str], default {}
1232
- Dictionary of environment variables to set.
1233
- """
1234
- ...
1235
-
1236
- @typing.overload
1237
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1238
- """
1239
- Specifies the PyPI packages for all steps of the flow.
1240
-
1241
- Use `@pypi_base` to set common packages required by all
1242
- steps and use `@pypi` to specify step-specific overrides.
1243
-
1244
- Parameters
1245
- ----------
1246
- packages : Dict[str, str], default: {}
1247
- Packages to use for this flow. The key is the name of the package
1248
- and the value is the version to use.
1249
- python : str, optional, default: None
1250
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1251
- that the version used will correspond to the version of the Python interpreter used to start the run.
1252
- """
1253
- ...
1254
-
1255
- @typing.overload
1256
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1257
- ...
1258
-
1259
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1260
- """
1261
- Specifies the PyPI packages for all steps of the flow.
1262
-
1263
- Use `@pypi_base` to set common packages required by all
1264
- steps and use `@pypi` to specify step-specific overrides.
1265
-
1266
- Parameters
1267
- ----------
1268
- packages : Dict[str, str], default: {}
1269
- Packages to use for this flow. The key is the name of the package
1270
- and the value is the version to use.
1271
- python : str, optional, default: None
1272
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1273
- that the version used will correspond to the version of the Python interpreter used to start the run.
1274
- """
1275
- ...
1276
-
1277
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1278
- """
1279
- Specifies what flows belong to the same project.
1280
-
1281
- A project-specific namespace is created for all flows that
1282
- use the same `@project(name)`.
1283
-
1284
-
1285
- Parameters
1286
- ----------
1287
- name : str
1288
- Project name. Make sure that the name is unique amongst all
1289
- projects that use the same production scheduler. The name may
1290
- contain only lowercase alphanumeric characters and underscores.
1291
-
1292
- branch : Optional[str], default None
1293
- The branch to use. If not specified, the branch is set to
1294
- `user.<username>` unless `production` is set to `True`. This can
1295
- also be set on the command line using `--branch` as a top-level option.
1296
- It is an error to specify `branch` in the decorator and on the command line.
1297
-
1298
- production : bool, default False
1299
- Whether or not the branch is the production branch. This can also be set on the
1300
- command line using `--production` as a top-level option. It is an error to specify
1301
- `production` in the decorator and on the command line.
1302
- The project branch name will be:
1303
- - if `branch` is specified:
1304
- - if `production` is True: `prod.<branch>`
1305
- - if `production` is False: `test.<branch>`
1306
- - if `branch` is not specified:
1307
- - if `production` is True: `prod`
1308
- - if `production` is False: `user.<username>`
1309
- """
1310
- ...
1311
-
1312
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1313
- """
1314
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1315
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1316
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1317
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1318
- starts only after all sensors finish.
1319
-
1320
-
1321
- Parameters
1322
- ----------
1323
- timeout : int
1324
- Time, in seconds before the task times out and fails. (Default: 3600)
1325
- poke_interval : int
1326
- Time in seconds that the job should wait in between each try. (Default: 60)
1327
- mode : str
1328
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1329
- exponential_backoff : bool
1330
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1331
- pool : str
1332
- the slot pool this task should run in,
1333
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1334
- soft_fail : bool
1335
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1336
- name : str
1337
- Name of the sensor on Airflow
1338
- description : str
1339
- Description of sensor in the Airflow UI
1340
- bucket_key : Union[str, List[str]]
1341
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1342
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1343
- bucket_name : str
1344
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1345
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1346
- wildcard_match : bool
1347
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1348
- aws_conn_id : str
1349
- a reference to the s3 connection on Airflow. (Default: None)
1350
- verify : bool
1351
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1352
- """
1353
- ...
1354
-
1355
- @typing.overload
1356
- def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1357
- """
1358
- Specifies the event(s) that this flow depends on.
1359
-
1360
- ```
1361
- @trigger(event='foo')
1362
- ```
1363
- or
1364
- ```
1365
- @trigger(events=['foo', 'bar'])
1366
- ```
1367
-
1368
- Additionally, you can specify the parameter mappings
1369
- to map event payload to Metaflow parameters for the flow.
1370
- ```
1371
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1372
- ```
1373
- or
1374
- ```
1375
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1376
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1377
- ```
1378
-
1379
- 'parameters' can also be a list of strings and tuples like so:
1380
- ```
1381
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1382
- ```
1383
- This is equivalent to:
1144
+ # where the key is the name of the artifact and the value is the path to the model
1145
+ print(os.listdir(current.model.loaded["my_model"]))
1146
+ self.next(self.end)
1384
1147
  ```
1385
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1148
+
1149
+ - Loading models
1150
+ ```python
1151
+ @step
1152
+ def train(self):
1153
+ # current.model.load returns the path to the model loaded
1154
+ checkpoint_path = current.model.load(
1155
+ self.checkpoint_key,
1156
+ )
1157
+ model_path = current.model.load(
1158
+ self.model,
1159
+ )
1160
+ self.next(self.test)
1386
1161
  ```
1387
1162
 
1388
1163
 
1389
1164
  Parameters
1390
1165
  ----------
1391
- event : Union[str, Dict[str, Any]], optional, default None
1392
- Event dependency for this flow.
1393
- events : List[Union[str, Dict[str, Any]]], default []
1394
- Events dependency for this flow.
1395
- options : Dict[str, Any], default {}
1396
- Backend-specific configuration for tuning eventing behavior.
1166
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1167
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1168
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1169
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1170
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1171
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1172
+
1173
+ temp_dir_root : str, default: None
1174
+ The root directory under which `current.model.loaded` will store loaded models
1397
1175
  """
1398
1176
  ...
1399
1177
 
1400
1178
  @typing.overload
1401
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1179
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1180
+ """
1181
+ Decorator prototype for all step decorators. This function gets specialized
1182
+ and imported for all decorators types by _import_plugin_decorators().
1183
+ """
1402
1184
  ...
1403
1185
 
1404
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1186
+ @typing.overload
1187
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1188
+ ...
1189
+
1190
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1405
1191
  """
1406
- Specifies the event(s) that this flow depends on.
1407
-
1408
- ```
1409
- @trigger(event='foo')
1410
- ```
1411
- or
1412
- ```
1413
- @trigger(events=['foo', 'bar'])
1414
- ```
1192
+ Decorator prototype for all step decorators. This function gets specialized
1193
+ and imported for all decorators types by _import_plugin_decorators().
1194
+ """
1195
+ ...
1196
+
1197
+ @typing.overload
1198
+ 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]]]:
1199
+ """
1200
+ Specifies the number of times the task corresponding
1201
+ to a step needs to be retried.
1415
1202
 
1416
- Additionally, you can specify the parameter mappings
1417
- to map event payload to Metaflow parameters for the flow.
1418
- ```
1419
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1420
- ```
1421
- or
1422
- ```
1423
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1424
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1425
- ```
1203
+ This decorator is useful for handling transient errors, such as networking issues.
1204
+ If your task contains operations that can't be retried safely, e.g. database updates,
1205
+ it is advisable to annotate it with `@retry(times=0)`.
1426
1206
 
1427
- 'parameters' can also be a list of strings and tuples like so:
1428
- ```
1429
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1430
- ```
1431
- This is equivalent to:
1432
- ```
1433
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1434
- ```
1207
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1208
+ decorator will execute a no-op task after all retries have been exhausted,
1209
+ ensuring that the flow execution can continue.
1435
1210
 
1436
1211
 
1437
1212
  Parameters
1438
1213
  ----------
1439
- event : Union[str, Dict[str, Any]], optional, default None
1440
- Event dependency for this flow.
1441
- events : List[Union[str, Dict[str, Any]]], default []
1442
- Events dependency for this flow.
1443
- options : Dict[str, Any], default {}
1444
- Backend-specific configuration for tuning eventing behavior.
1214
+ times : int, default 3
1215
+ Number of times to retry this task.
1216
+ minutes_between_retries : int, default 2
1217
+ Number of minutes between retries.
1445
1218
  """
1446
1219
  ...
1447
1220
 
1448
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1221
+ @typing.overload
1222
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1223
+ ...
1224
+
1225
+ @typing.overload
1226
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1227
+ ...
1228
+
1229
+ 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):
1449
1230
  """
1450
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1451
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1231
+ Specifies the number of times the task corresponding
1232
+ to a step needs to be retried.
1233
+
1234
+ This decorator is useful for handling transient errors, such as networking issues.
1235
+ If your task contains operations that can't be retried safely, e.g. database updates,
1236
+ it is advisable to annotate it with `@retry(times=0)`.
1237
+
1238
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1239
+ decorator will execute a no-op task after all retries have been exhausted,
1240
+ ensuring that the flow execution can continue.
1452
1241
 
1453
1242
 
1454
1243
  Parameters
1455
1244
  ----------
1456
- timeout : int
1457
- Time, in seconds before the task times out and fails. (Default: 3600)
1458
- poke_interval : int
1459
- Time in seconds that the job should wait in between each try. (Default: 60)
1460
- mode : str
1461
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1462
- exponential_backoff : bool
1463
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1464
- pool : str
1465
- the slot pool this task should run in,
1466
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1467
- soft_fail : bool
1468
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1469
- name : str
1470
- Name of the sensor on Airflow
1471
- description : str
1472
- Description of sensor in the Airflow UI
1473
- external_dag_id : str
1474
- The dag_id that contains the task you want to wait for.
1475
- external_task_ids : List[str]
1476
- The list of task_ids that you want to wait for.
1477
- If None (default value) the sensor waits for the DAG. (Default: None)
1478
- allowed_states : List[str]
1479
- Iterable of allowed states, (Default: ['success'])
1480
- failed_states : List[str]
1481
- Iterable of failed or dis-allowed states. (Default: None)
1482
- execution_delta : datetime.timedelta
1483
- time difference with the previous execution to look at,
1484
- the default is the same logical date as the current task or DAG. (Default: None)
1485
- check_existence: bool
1486
- Set to True to check if the external task exists or check if
1487
- the DAG to wait for exists. (Default: True)
1245
+ times : int, default 3
1246
+ Number of times to retry this task.
1247
+ minutes_between_retries : int, default 2
1248
+ Number of minutes between retries.
1488
1249
  """
1489
1250
  ...
1490
1251
 
@@ -1589,57 +1350,6 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1589
1350
  """
1590
1351
  ...
1591
1352
 
1592
- @typing.overload
1593
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1594
- """
1595
- Specifies the times when the flow should be run when running on a
1596
- production scheduler.
1597
-
1598
-
1599
- Parameters
1600
- ----------
1601
- hourly : bool, default False
1602
- Run the workflow hourly.
1603
- daily : bool, default True
1604
- Run the workflow daily.
1605
- weekly : bool, default False
1606
- Run the workflow weekly.
1607
- cron : str, optional, default None
1608
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1609
- specified by this expression.
1610
- timezone : str, optional, default None
1611
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1612
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1613
- """
1614
- ...
1615
-
1616
- @typing.overload
1617
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1618
- ...
1619
-
1620
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1621
- """
1622
- Specifies the times when the flow should be run when running on a
1623
- production scheduler.
1624
-
1625
-
1626
- Parameters
1627
- ----------
1628
- hourly : bool, default False
1629
- Run the workflow hourly.
1630
- daily : bool, default True
1631
- Run the workflow daily.
1632
- weekly : bool, default False
1633
- Run the workflow weekly.
1634
- cron : str, optional, default None
1635
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1636
- specified by this expression.
1637
- timezone : str, optional, default None
1638
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1639
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1640
- """
1641
- ...
1642
-
1643
1353
  @typing.overload
1644
1354
  def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1645
1355
  """
@@ -1691,6 +1401,57 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1691
1401
  """
1692
1402
  ...
1693
1403
 
1404
+ @typing.overload
1405
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1406
+ """
1407
+ Specifies the times when the flow should be run when running on a
1408
+ production scheduler.
1409
+
1410
+
1411
+ Parameters
1412
+ ----------
1413
+ hourly : bool, default False
1414
+ Run the workflow hourly.
1415
+ daily : bool, default True
1416
+ Run the workflow daily.
1417
+ weekly : bool, default False
1418
+ Run the workflow weekly.
1419
+ cron : str, optional, default None
1420
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1421
+ specified by this expression.
1422
+ timezone : str, optional, default None
1423
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1424
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1425
+ """
1426
+ ...
1427
+
1428
+ @typing.overload
1429
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1430
+ ...
1431
+
1432
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1433
+ """
1434
+ Specifies the times when the flow should be run when running on a
1435
+ production scheduler.
1436
+
1437
+
1438
+ Parameters
1439
+ ----------
1440
+ hourly : bool, default False
1441
+ Run the workflow hourly.
1442
+ daily : bool, default True
1443
+ Run the workflow daily.
1444
+ weekly : bool, default False
1445
+ Run the workflow weekly.
1446
+ cron : str, optional, default None
1447
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1448
+ specified by this expression.
1449
+ timezone : str, optional, default None
1450
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1451
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1452
+ """
1453
+ ...
1454
+
1694
1455
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1695
1456
  """
1696
1457
  Allows setting external datastores to save data for the
@@ -1805,5 +1566,260 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1805
1566
  """
1806
1567
  ...
1807
1568
 
1569
+ @typing.overload
1570
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1571
+ """
1572
+ Specifies the PyPI packages for all steps of the flow.
1573
+
1574
+ Use `@pypi_base` to set common packages required by all
1575
+ steps and use `@pypi` to specify step-specific overrides.
1576
+
1577
+ Parameters
1578
+ ----------
1579
+ packages : Dict[str, str], default: {}
1580
+ Packages to use for this flow. The key is the name of the package
1581
+ and the value is the version to use.
1582
+ python : str, optional, default: None
1583
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1584
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1585
+ """
1586
+ ...
1587
+
1588
+ @typing.overload
1589
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1590
+ ...
1591
+
1592
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1593
+ """
1594
+ Specifies the PyPI packages for all steps of the flow.
1595
+
1596
+ Use `@pypi_base` to set common packages required by all
1597
+ steps and use `@pypi` to specify step-specific overrides.
1598
+
1599
+ Parameters
1600
+ ----------
1601
+ packages : Dict[str, str], default: {}
1602
+ Packages to use for this flow. The key is the name of the package
1603
+ and the value is the version to use.
1604
+ python : str, optional, default: None
1605
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1606
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1607
+ """
1608
+ ...
1609
+
1610
+ 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]]:
1611
+ """
1612
+ 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)
1613
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1614
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1615
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1616
+ starts only after all sensors finish.
1617
+
1618
+
1619
+ Parameters
1620
+ ----------
1621
+ timeout : int
1622
+ Time, in seconds before the task times out and fails. (Default: 3600)
1623
+ poke_interval : int
1624
+ Time in seconds that the job should wait in between each try. (Default: 60)
1625
+ mode : str
1626
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1627
+ exponential_backoff : bool
1628
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1629
+ pool : str
1630
+ the slot pool this task should run in,
1631
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1632
+ soft_fail : bool
1633
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1634
+ name : str
1635
+ Name of the sensor on Airflow
1636
+ description : str
1637
+ Description of sensor in the Airflow UI
1638
+ bucket_key : Union[str, List[str]]
1639
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1640
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1641
+ bucket_name : str
1642
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1643
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1644
+ wildcard_match : bool
1645
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1646
+ aws_conn_id : str
1647
+ a reference to the s3 connection on Airflow. (Default: None)
1648
+ verify : bool
1649
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1650
+ """
1651
+ ...
1652
+
1653
+ @typing.overload
1654
+ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1655
+ """
1656
+ Specifies the event(s) that this flow depends on.
1657
+
1658
+ ```
1659
+ @trigger(event='foo')
1660
+ ```
1661
+ or
1662
+ ```
1663
+ @trigger(events=['foo', 'bar'])
1664
+ ```
1665
+
1666
+ Additionally, you can specify the parameter mappings
1667
+ to map event payload to Metaflow parameters for the flow.
1668
+ ```
1669
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1670
+ ```
1671
+ or
1672
+ ```
1673
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1674
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1675
+ ```
1676
+
1677
+ 'parameters' can also be a list of strings and tuples like so:
1678
+ ```
1679
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1680
+ ```
1681
+ This is equivalent to:
1682
+ ```
1683
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1684
+ ```
1685
+
1686
+
1687
+ Parameters
1688
+ ----------
1689
+ event : Union[str, Dict[str, Any]], optional, default None
1690
+ Event dependency for this flow.
1691
+ events : List[Union[str, Dict[str, Any]]], default []
1692
+ Events dependency for this flow.
1693
+ options : Dict[str, Any], default {}
1694
+ Backend-specific configuration for tuning eventing behavior.
1695
+ """
1696
+ ...
1697
+
1698
+ @typing.overload
1699
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1700
+ ...
1701
+
1702
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1703
+ """
1704
+ Specifies the event(s) that this flow depends on.
1705
+
1706
+ ```
1707
+ @trigger(event='foo')
1708
+ ```
1709
+ or
1710
+ ```
1711
+ @trigger(events=['foo', 'bar'])
1712
+ ```
1713
+
1714
+ Additionally, you can specify the parameter mappings
1715
+ to map event payload to Metaflow parameters for the flow.
1716
+ ```
1717
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1718
+ ```
1719
+ or
1720
+ ```
1721
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1722
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1723
+ ```
1724
+
1725
+ 'parameters' can also be a list of strings and tuples like so:
1726
+ ```
1727
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1728
+ ```
1729
+ This is equivalent to:
1730
+ ```
1731
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1732
+ ```
1733
+
1734
+
1735
+ Parameters
1736
+ ----------
1737
+ event : Union[str, Dict[str, Any]], optional, default None
1738
+ Event dependency for this flow.
1739
+ events : List[Union[str, Dict[str, Any]]], default []
1740
+ Events dependency for this flow.
1741
+ options : Dict[str, Any], default {}
1742
+ Backend-specific configuration for tuning eventing behavior.
1743
+ """
1744
+ ...
1745
+
1746
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1747
+ """
1748
+ Specifies what flows belong to the same project.
1749
+
1750
+ A project-specific namespace is created for all flows that
1751
+ use the same `@project(name)`.
1752
+
1753
+
1754
+ Parameters
1755
+ ----------
1756
+ name : str
1757
+ Project name. Make sure that the name is unique amongst all
1758
+ projects that use the same production scheduler. The name may
1759
+ contain only lowercase alphanumeric characters and underscores.
1760
+
1761
+ branch : Optional[str], default None
1762
+ The branch to use. If not specified, the branch is set to
1763
+ `user.<username>` unless `production` is set to `True`. This can
1764
+ also be set on the command line using `--branch` as a top-level option.
1765
+ It is an error to specify `branch` in the decorator and on the command line.
1766
+
1767
+ production : bool, default False
1768
+ Whether or not the branch is the production branch. This can also be set on the
1769
+ command line using `--production` as a top-level option. It is an error to specify
1770
+ `production` in the decorator and on the command line.
1771
+ The project branch name will be:
1772
+ - if `branch` is specified:
1773
+ - if `production` is True: `prod.<branch>`
1774
+ - if `production` is False: `test.<branch>`
1775
+ - if `branch` is not specified:
1776
+ - if `production` is True: `prod`
1777
+ - if `production` is False: `user.<username>`
1778
+ """
1779
+ ...
1780
+
1781
+ 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]]:
1782
+ """
1783
+ 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.
1784
+ 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.
1785
+
1786
+
1787
+ Parameters
1788
+ ----------
1789
+ timeout : int
1790
+ Time, in seconds before the task times out and fails. (Default: 3600)
1791
+ poke_interval : int
1792
+ Time in seconds that the job should wait in between each try. (Default: 60)
1793
+ mode : str
1794
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1795
+ exponential_backoff : bool
1796
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1797
+ pool : str
1798
+ the slot pool this task should run in,
1799
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1800
+ soft_fail : bool
1801
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1802
+ name : str
1803
+ Name of the sensor on Airflow
1804
+ description : str
1805
+ Description of sensor in the Airflow UI
1806
+ external_dag_id : str
1807
+ The dag_id that contains the task you want to wait for.
1808
+ external_task_ids : List[str]
1809
+ The list of task_ids that you want to wait for.
1810
+ If None (default value) the sensor waits for the DAG. (Default: None)
1811
+ allowed_states : List[str]
1812
+ Iterable of allowed states, (Default: ['success'])
1813
+ failed_states : List[str]
1814
+ Iterable of failed or dis-allowed states. (Default: None)
1815
+ execution_delta : datetime.timedelta
1816
+ time difference with the previous execution to look at,
1817
+ the default is the same logical date as the current task or DAG. (Default: None)
1818
+ check_existence: bool
1819
+ Set to True to check if the external task exists or check if
1820
+ the DAG to wait for exists. (Default: True)
1821
+ """
1822
+ ...
1823
+
1808
1824
  pkg_name: str
1809
1825