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