ob-metaflow-stubs 6.0.3.188rc1__py2.py3-none-any.whl → 6.0.3.188rc3__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 (246) hide show
  1. metaflow-stubs/__init__.pyi +829 -824
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +39 -26
  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 +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +3 -3
  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 +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +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 +3 -3
  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 +2 -2
  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 +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  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 +12 -6
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +2 -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 +4 -44
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +15 -30
  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 +4 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +24 -0
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +53 -0
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +426 -0
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +64 -0
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +50 -0
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +241 -0
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +2 -2
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +46 -0
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +1 -11
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +4 -4
  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 +2 -2
  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 +4 -4
  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 +2 -1
  115. metaflow-stubs/multicore_utils.pyi +1 -1
  116. metaflow-stubs/ob_internal.pyi +2 -1
  117. metaflow-stubs/parameters.pyi +3 -3
  118. metaflow-stubs/plugins/__init__.pyi +8 -8
  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 +3 -3
  130. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  131. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  132. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  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 +2 -2
  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 +1 -1
  157. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  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 +2 -2
  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 +2 -2
  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 +3 -3
  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 +2 -2
  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 +1 -1
  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 +2 -2
  202. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  203. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  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 +2 -2
  211. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  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 +2 -2
  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 +5 -5
  224. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  225. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  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 +2 -2
  230. metaflow-stubs/system/__init__.pyi +1 -1
  231. metaflow-stubs/system/system_logger.pyi +1 -1
  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.188rc1.dist-info → ob_metaflow_stubs-6.0.3.188rc3.dist-info}/METADATA +1 -1
  240. ob_metaflow_stubs-6.0.3.188rc3.dist-info/RECORD +243 -0
  241. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/cli_to_config.pyi +0 -17
  242. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/secrets.pyi +0 -46
  243. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/validations.pyi +0 -24
  244. ob_metaflow_stubs-6.0.3.188rc1.dist-info/RECORD +0 -239
  245. {ob_metaflow_stubs-6.0.3.188rc1.dist-info → ob_metaflow_stubs-6.0.3.188rc3.dist-info}/WHEEL +0 -0
  246. {ob_metaflow_stubs-6.0.3.188rc1.dist-info → ob_metaflow_stubs-6.0.3.188rc3.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.18.1+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-07-07T22:26:05.548812 #
4
+ # Generated on 2025-07-10T08:45:58.339588 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -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 events as events
39
- from . import tuple_util as tuple_util
40
38
  from . import cards as cards
39
+ from . import tuple_util as tuple_util
41
40
  from . import metaflow_git as metaflow_git
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 conda_environment_yml_parser as conda_environment_yml_parser
47
48
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
49
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
- from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_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
@@ -73,6 +73,7 @@ from .mf_extensions.outerbounds.plugins.snowflake.snowflake import Snowflake as
73
73
  from .mf_extensions.outerbounds.plugins.checkpoint_datastores.nebius import nebius_checkpoints as nebius_checkpoints
74
74
  from .mf_extensions.outerbounds.plugins.checkpoint_datastores.coreweave import coreweave_checkpoints as coreweave_checkpoints
75
75
  from .mf_extensions.outerbounds.plugins.aws.assume_role_decorator import assume_role as assume_role
76
+ from .mf_extensions.outerbounds.plugins.apps.core.deployer import AppDeployer as AppDeployer
76
77
  from . import cli_components as cli_components
77
78
  from . import system as system
78
79
  from . import pylint_wrapper as pylint_wrapper
@@ -155,6 +156,243 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
156
  """
156
157
  ...
157
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]]]:
161
+ """
162
+ Specifies the number of times the task corresponding
163
+ to a step needs to be retried.
164
+
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
+
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.
172
+
173
+
174
+ Parameters
175
+ ----------
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.
180
+ """
181
+ ...
182
+
183
+ @typing.overload
184
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
185
+ ...
186
+
187
+ @typing.overload
188
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
189
+ ...
190
+
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):
192
+ """
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)`.
199
+
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.
203
+
204
+
205
+ Parameters
206
+ ----------
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.
211
+ """
212
+ ...
213
+
214
+ @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]]]:
216
+ """
217
+ Specifies secrets to be retrieved and injected as environment variables prior to
218
+ the execution of a step.
219
+
220
+
221
+ Parameters
222
+ ----------
223
+ sources : List[Union[str, Dict[str, Any]]], default: []
224
+ List of secret specs, defining how the secrets are to be retrieved
225
+ """
226
+ ...
227
+
228
+ @typing.overload
229
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
230
+ ...
231
+
232
+ @typing.overload
233
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
234
+ ...
235
+
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]]] = []):
237
+ """
238
+ Specifies secrets to be retrieved and injected as environment variables prior to
239
+ the execution of a step.
240
+
241
+
242
+ Parameters
243
+ ----------
244
+ sources : List[Union[str, Dict[str, Any]]], default: []
245
+ List of secret specs, defining how the secrets are to be retrieved
246
+ """
247
+ ...
248
+
249
+ @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]]]:
251
+ """
252
+ Enables checkpointing for a step.
253
+
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
+ ```
299
+
300
+
301
+ Parameters
302
+ ----------
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.
316
+ """
317
+ ...
318
+
319
+ @typing.overload
320
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
321
+ ...
322
+
323
+ @typing.overload
324
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
325
+ ...
326
+
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):
328
+ """
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
371
+
372
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
373
+ for i in range(self.epochs):
374
+ ...
375
+ ```
376
+
377
+
378
+ Parameters
379
+ ----------
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.
393
+ """
394
+ ...
395
+
158
396
  @typing.overload
159
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]]]:
160
398
  """
@@ -284,7 +522,7 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
284
522
  """
285
523
  ...
286
524
 
287
- 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]]]:
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]]]:
288
526
  """
289
527
  Specifies that this step should execute on DGX cloud.
290
528
 
@@ -295,22 +533,67 @@ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Cal
295
533
  Number of GPUs to use.
296
534
  gpu_type : str
297
535
  Type of Nvidia GPU to use.
536
+ queue_timeout : int
537
+ Time to keep the job in NVCF's queue.
298
538
  """
299
539
  ...
300
540
 
301
- 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]]]:
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]]]:
302
543
  """
303
- Specifies that this step should execute on DGX cloud.
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.
304
550
 
305
551
 
306
552
  Parameters
307
553
  ----------
308
- gpu : int
309
- Number of GPUs to use.
310
- gpu_type : str
311
- Type of Nvidia GPU to use.
312
- queue_timeout : int
313
- Time to keep the job in NVCF's queue.
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.
564
+ """
565
+ ...
566
+
567
+ @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]:
573
+ ...
574
+
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):
576
+ """
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.
314
597
  """
315
598
  ...
316
599
 
@@ -416,538 +699,317 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
416
699
  """
417
700
  ...
418
701
 
419
- @typing.overload
420
- 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]]]:
702
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
421
703
  """
422
- Specifies the Conda environment for the step.
423
-
424
- Information in this decorator will augment any
425
- attributes set in the `@conda_base` flow-level decorator. Hence,
426
- you can use `@conda_base` to set packages required by all
427
- steps and use `@conda` to specify step-specific overrides.
428
-
704
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
429
705
 
430
- Parameters
431
- ----------
432
- packages : Dict[str, str], default {}
433
- Packages to use for this step. The key is the name of the package
434
- and the value is the version to use.
435
- libraries : Dict[str, str], default {}
436
- Supported for backward compatibility. When used with packages, packages will take precedence.
437
- python : str, optional, default None
438
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
439
- that the version used will correspond to the version of the Python interpreter used to start the run.
440
- disabled : bool, default False
441
- If set to True, disables @conda.
442
- """
443
- ...
444
-
445
- @typing.overload
446
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
447
- ...
448
-
449
- @typing.overload
450
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
451
- ...
452
-
453
- 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):
454
- """
455
- Specifies the Conda environment for the step.
706
+ > Examples
456
707
 
457
- Information in this decorator will augment any
458
- attributes set in the `@conda_base` flow-level decorator. Hence,
459
- you can use `@conda_base` to set packages required by all
460
- steps and use `@conda` to specify step-specific overrides.
708
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
709
+ ```python
710
+ @huggingface_hub
711
+ @step
712
+ def pull_model_from_huggingface(self):
713
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
714
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
715
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
716
+ # value of the function is a reference to the model in the backend storage.
717
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
461
718
 
719
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
720
+ self.llama_model = current.huggingface_hub.snapshot_download(
721
+ repo_id=self.model_id,
722
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
723
+ )
724
+ self.next(self.train)
725
+ ```
462
726
 
463
- Parameters
464
- ----------
465
- packages : Dict[str, str], default {}
466
- Packages to use for this step. The key is the name of the package
467
- and the value is the version to use.
468
- libraries : Dict[str, str], default {}
469
- Supported for backward compatibility. When used with packages, packages will take precedence.
470
- python : str, optional, default None
471
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
472
- that the version used will correspond to the version of the Python interpreter used to start the run.
473
- disabled : bool, default False
474
- If set to True, disables @conda.
475
- """
476
- ...
477
-
478
- @typing.overload
479
- 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]]]:
480
- """
481
- Specifies the number of times the task corresponding
482
- to a step needs to be retried.
727
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
728
+ ```python
729
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
730
+ @step
731
+ def pull_model_from_huggingface(self):
732
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
733
+ ```
483
734
 
484
- This decorator is useful for handling transient errors, such as networking issues.
485
- If your task contains operations that can't be retried safely, e.g. database updates,
486
- it is advisable to annotate it with `@retry(times=0)`.
735
+ ```python
736
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
737
+ @step
738
+ def finetune_model(self):
739
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
740
+ # path_to_model will be /my-directory
741
+ ```
487
742
 
488
- This can be used in conjunction with the `@catch` decorator. The `@catch`
489
- decorator will execute a no-op task after all retries have been exhausted,
490
- ensuring that the flow execution can continue.
743
+ ```python
744
+ # Takes all the arguments passed to `snapshot_download`
745
+ # except for `local_dir`
746
+ @huggingface_hub(load=[
747
+ {
748
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
749
+ },
750
+ {
751
+ "repo_id": "myorg/mistral-lora",
752
+ "repo_type": "model",
753
+ },
754
+ ])
755
+ @step
756
+ def finetune_model(self):
757
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
758
+ # path_to_model will be /my-directory
759
+ ```
491
760
 
492
761
 
493
762
  Parameters
494
763
  ----------
495
- times : int, default 3
496
- Number of times to retry this task.
497
- minutes_between_retries : int, default 2
498
- Number of minutes between retries.
499
- """
500
- ...
501
-
502
- @typing.overload
503
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
504
- ...
505
-
506
- @typing.overload
507
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
508
- ...
509
-
510
- 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):
511
- """
512
- Specifies the number of times the task corresponding
513
- to a step needs to be retried.
764
+ temp_dir_root : str, optional
765
+ The root directory that will hold the temporary directory where objects will be downloaded.
514
766
 
515
- This decorator is useful for handling transient errors, such as networking issues.
516
- If your task contains operations that can't be retried safely, e.g. database updates,
517
- it is advisable to annotate it with `@retry(times=0)`.
767
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
768
+ The list of repos (models/datasets) to load.
518
769
 
519
- This can be used in conjunction with the `@catch` decorator. The `@catch`
520
- decorator will execute a no-op task after all retries have been exhausted,
521
- ensuring that the flow execution can continue.
770
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
522
771
 
772
+ - If repo (model/dataset) is not found in the datastore:
773
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
774
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
775
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
523
776
 
524
- Parameters
525
- ----------
526
- times : int, default 3
527
- Number of times to retry this task.
528
- minutes_between_retries : int, default 2
529
- Number of minutes between retries.
777
+ - If repo is found in the datastore:
778
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
530
779
  """
531
780
  ...
532
781
 
533
- 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]]]:
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]]]:
534
783
  """
535
- This decorator is used to run vllm APIs as Metaflow task sidecars.
784
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
536
785
 
537
786
  User code call
538
787
  --------------
539
- @vllm(
540
- model="...",
788
+ @ollama(
789
+ models=[...],
541
790
  ...
542
791
  )
543
792
 
544
793
  Valid backend options
545
794
  ---------------------
546
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.
547
798
 
548
799
  Valid model options
549
800
  -------------------
550
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
551
-
552
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
553
- If you need multiple models, you must create multiple @vllm decorators.
801
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
554
802
 
555
803
 
556
804
  Parameters
557
805
  ----------
558
- model: str
559
- HuggingFace model identifier to be served by vLLM.
806
+ models: list[str]
807
+ List of Ollama containers running models in sidecars.
560
808
  backend: str
561
- Determines where and how to run the vLLM process.
562
- openai_api_server: bool
563
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
564
- Default is False (uses native engine).
565
- Set to True for backward compatibility with existing code.
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.
566
816
  debug: bool
567
817
  Whether to turn on verbose debugging logs.
568
- card_refresh_interval: int
569
- Interval in seconds for refreshing the vLLM status card.
570
- Only used when openai_api_server=True.
571
- max_retries: int
572
- Maximum number of retries checking for vLLM server startup.
573
- Only used when openai_api_server=True.
574
- retry_alert_frequency: int
575
- Frequency of alert logs for vLLM server startup retries.
576
- Only used when openai_api_server=True.
577
- engine_args : dict
578
- Additional keyword arguments to pass to the vLLM engine.
579
- For example, `tensor_parallel_size=2`.
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.
580
822
  """
581
823
  ...
582
824
 
583
- @typing.overload
584
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
585
- """
586
- Decorator prototype for all step decorators. This function gets specialized
587
- and imported for all decorators types by _import_plugin_decorators().
588
- """
589
- ...
590
-
591
- @typing.overload
592
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
593
- ...
594
-
595
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
596
- """
597
- Decorator prototype for all step decorators. This function gets specialized
598
- and imported for all decorators types by _import_plugin_decorators().
599
- """
600
- ...
601
-
602
- @typing.overload
603
- 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]]]:
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]]]:
604
826
  """
605
- Creates a human-readable report, a Metaflow Card, after this step completes.
606
-
607
- Note that you may add multiple `@card` decorators in a step with different parameters.
827
+ Specifies that this step should execute on DGX cloud.
608
828
 
609
829
 
610
830
  Parameters
611
831
  ----------
612
- type : str, default 'default'
613
- Card type.
614
- id : str, optional, default None
615
- If multiple cards are present, use this id to identify this card.
616
- options : Dict[str, Any], default {}
617
- Options passed to the card. The contents depend on the card type.
618
- timeout : int, default 45
619
- Interrupt reporting if it takes more than this many seconds.
832
+ gpu : int
833
+ Number of GPUs to use.
834
+ gpu_type : str
835
+ Type of Nvidia GPU to use.
620
836
  """
621
837
  ...
622
838
 
623
839
  @typing.overload
624
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
625
- ...
626
-
627
- @typing.overload
628
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
629
- ...
630
-
631
- 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):
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]]]:
632
841
  """
633
- Creates a human-readable report, a Metaflow Card, after this step completes.
842
+ Specifies a timeout for your step.
634
843
 
635
- Note that you may add multiple `@card` decorators in a step with different parameters.
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.
636
852
 
637
853
 
638
854
  Parameters
639
855
  ----------
640
- type : str, default 'default'
641
- Card type.
642
- id : str, optional, default None
643
- If multiple cards are present, use this id to identify this card.
644
- options : Dict[str, Any], default {}
645
- Options passed to the card. The contents depend on the card type.
646
- timeout : int, default 45
647
- Interrupt reporting if it takes more than this many seconds.
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.
648
862
  """
649
863
  ...
650
864
 
651
865
  @typing.overload
652
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
653
- """
654
- Internal decorator to support Fast bakery
655
- """
866
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
656
867
  ...
657
868
 
658
869
  @typing.overload
659
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
660
- ...
661
-
662
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
663
- """
664
- Internal decorator to support Fast bakery
665
- """
870
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
666
871
  ...
667
872
 
668
- @typing.overload
669
- 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]]]:
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):
670
874
  """
671
- Enables checkpointing for a step.
672
-
673
- > Examples
674
-
675
- - Saving Checkpoints
676
-
677
- ```python
678
- @checkpoint
679
- @step
680
- def train(self):
681
- model = create_model(self.parameters, checkpoint_path = None)
682
- for i in range(self.epochs):
683
- # some training logic
684
- loss = model.train(self.dataset)
685
- if i % 10 == 0:
686
- model.save(
687
- current.checkpoint.directory,
688
- )
689
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
690
- # and returns a reference dictionary to the checkpoint saved in the datastore
691
- self.latest_checkpoint = current.checkpoint.save(
692
- name="epoch_checkpoint",
693
- metadata={
694
- "epoch": i,
695
- "loss": loss,
696
- }
697
- )
698
- ```
875
+ Specifies a timeout for your step.
699
876
 
700
- - Using Loaded Checkpoints
877
+ This decorator is useful if this step may hang indefinitely.
701
878
 
702
- ```python
703
- @retry(times=3)
704
- @checkpoint
705
- @step
706
- def train(self):
707
- # Assume that the task has restarted and the previous attempt of the task
708
- # saved a checkpoint
709
- checkpoint_path = None
710
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
711
- print("Loaded checkpoint from the previous attempt")
712
- checkpoint_path = current.checkpoint.directory
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.
713
882
 
714
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
715
- for i in range(self.epochs):
716
- ...
717
- ```
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.
718
885
 
719
886
 
720
887
  Parameters
721
888
  ----------
722
- load_policy : str, default: "fresh"
723
- The policy for loading the checkpoint. The following policies are supported:
724
- - "eager": Loads the the latest available checkpoint within the namespace.
725
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
726
- will be loaded at the start of the task.
727
- - "none": Do not load any checkpoint
728
- - "fresh": Loads the lastest checkpoint created within the running Task.
729
- This mode helps loading checkpoints across various retry attempts of the same task.
730
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
731
- created within the task will be loaded when the task is retries execution on failure.
732
-
733
- temp_dir_root : str, default: None
734
- The root directory under which `current.checkpoint.directory` will be created.
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.
735
895
  """
736
896
  ...
737
897
 
738
- @typing.overload
739
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
740
- ...
741
-
742
- @typing.overload
743
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
744
- ...
745
-
746
- 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):
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]]]:
747
899
  """
748
- Enables checkpointing for a step.
749
-
750
- > Examples
751
-
752
- - Saving Checkpoints
900
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
753
901
 
754
- ```python
755
- @checkpoint
756
- @step
757
- def train(self):
758
- model = create_model(self.parameters, checkpoint_path = None)
759
- for i in range(self.epochs):
760
- # some training logic
761
- loss = model.train(self.dataset)
762
- if i % 10 == 0:
763
- model.save(
764
- current.checkpoint.directory,
765
- )
766
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
767
- # and returns a reference dictionary to the checkpoint saved in the datastore
768
- self.latest_checkpoint = current.checkpoint.save(
769
- name="epoch_checkpoint",
770
- metadata={
771
- "epoch": i,
772
- "loss": loss,
773
- }
774
- )
775
- ```
902
+ User code call
903
+ --------------
904
+ @vllm(
905
+ model="...",
906
+ ...
907
+ )
776
908
 
777
- - Using Loaded Checkpoints
909
+ Valid backend options
910
+ ---------------------
911
+ - 'local': Run as a separate process on the local task machine.
778
912
 
779
- ```python
780
- @retry(times=3)
781
- @checkpoint
782
- @step
783
- def train(self):
784
- # Assume that the task has restarted and the previous attempt of the task
785
- # saved a checkpoint
786
- checkpoint_path = None
787
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
788
- print("Loaded checkpoint from the previous attempt")
789
- checkpoint_path = current.checkpoint.directory
913
+ Valid model options
914
+ -------------------
915
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
790
916
 
791
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
792
- for i in range(self.epochs):
793
- ...
794
- ```
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.
795
919
 
796
920
 
797
921
  Parameters
798
922
  ----------
799
- load_policy : str, default: "fresh"
800
- The policy for loading the checkpoint. The following policies are supported:
801
- - "eager": Loads the the latest available checkpoint within the namespace.
802
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
803
- will be loaded at the start of the task.
804
- - "none": Do not load any checkpoint
805
- - "fresh": Loads the lastest checkpoint created within the running Task.
806
- This mode helps loading checkpoints across various retry attempts of the same task.
807
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
808
- created within the task will be loaded when the task is retries execution on failure.
809
-
810
- temp_dir_root : str, default: None
811
- The root directory under which `current.checkpoint.directory` will be created.
923
+ model: str
924
+ HuggingFace model identifier to be served by vLLM.
925
+ 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.
931
+ debug: bool
932
+ 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`.
812
945
  """
813
946
  ...
814
947
 
815
- 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]]]:
948
+ @typing.overload
949
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
816
950
  """
817
- Specifies that this step should execute on Kubernetes.
818
-
819
-
820
- Parameters
821
- ----------
822
- cpu : int, default 1
823
- Number of CPUs required for this step. If `@resources` is
824
- also present, the maximum value from all decorators is used.
825
- memory : int, default 4096
826
- Memory size (in MB) required for this step. If
827
- `@resources` is also present, the maximum value from all decorators is
828
- used.
829
- disk : int, default 10240
830
- Disk size (in MB) required for this step. If
831
- `@resources` is also present, the maximum value from all decorators is
832
- used.
833
- image : str, optional, default None
834
- Docker image to use when launching on Kubernetes. If not specified, and
835
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
836
- not, a default Docker image mapping to the current version of Python is used.
837
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
838
- If given, the imagePullPolicy to be applied to the Docker image of the step.
839
- image_pull_secrets: List[str], default []
840
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
841
- Kubernetes image pull secrets to use when pulling container images
842
- in Kubernetes.
843
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
844
- Kubernetes service account to use when launching pod in Kubernetes.
845
- secrets : List[str], optional, default None
846
- Kubernetes secrets to use when launching pod in Kubernetes. These
847
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
848
- in Metaflow configuration.
849
- node_selector: Union[Dict[str,str], str], optional, default None
850
- Kubernetes node selector(s) to apply to the pod running the task.
851
- Can be passed in as a comma separated string of values e.g.
852
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
853
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
854
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
855
- Kubernetes namespace to use when launching pod in Kubernetes.
856
- gpu : int, optional, default None
857
- Number of GPUs required for this step. A value of zero implies that
858
- the scheduled node should not have GPUs.
859
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
860
- The vendor of the GPUs to be used for this step.
861
- tolerations : List[str], default []
862
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
863
- Kubernetes tolerations to use when launching pod in Kubernetes.
864
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
865
- Kubernetes labels to use when launching pod in Kubernetes.
866
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
867
- Kubernetes annotations to use when launching pod in Kubernetes.
868
- use_tmpfs : bool, default False
869
- This enables an explicit tmpfs mount for this step.
870
- tmpfs_tempdir : bool, default True
871
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
872
- tmpfs_size : int, optional, default: None
873
- The value for the size (in MiB) of the tmpfs mount for this step.
874
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
875
- memory allocated for this step.
876
- tmpfs_path : str, optional, default /metaflow_temp
877
- Path to tmpfs mount for this step.
878
- persistent_volume_claims : Dict[str, str], optional, default None
879
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
880
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
881
- shared_memory: int, optional
882
- Shared memory size (in MiB) required for this step
883
- port: int, optional
884
- Port number to specify in the Kubernetes job object
885
- compute_pool : str, optional, default None
886
- Compute pool to be used for for this step.
887
- If not specified, any accessible compute pool within the perimeter is used.
888
- hostname_resolution_timeout: int, default 10 * 60
889
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
890
- Only applicable when @parallel is used.
891
- qos: str, default: Burstable
892
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
893
-
894
- security_context: Dict[str, Any], optional, default None
895
- Container security context. Applies to the task container. Allows the following keys:
896
- - privileged: bool, optional, default None
897
- - allow_privilege_escalation: bool, optional, default None
898
- - run_as_user: int, optional, default None
899
- - run_as_group: int, optional, default None
900
- - run_as_non_root: bool, optional, default None
951
+ Decorator prototype for all step decorators. This function gets specialized
952
+ and imported for all decorators types by _import_plugin_decorators().
901
953
  """
902
954
  ...
903
955
 
904
956
  @typing.overload
905
- 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]]]:
957
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
958
+ ...
959
+
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
+ """
965
+ ...
966
+
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]]]:
906
969
  """
907
- Specifies secrets to be retrieved and injected as environment variables prior to
908
- the execution of a step.
970
+ Creates a human-readable report, a Metaflow Card, after this step completes.
971
+
972
+ Note that you may add multiple `@card` decorators in a step with different parameters.
909
973
 
910
974
 
911
975
  Parameters
912
976
  ----------
913
- sources : List[Union[str, Dict[str, Any]]], default: []
914
- List of secret specs, defining how the secrets are to be retrieved
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.
915
985
  """
916
986
  ...
917
987
 
918
988
  @typing.overload
919
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
989
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
920
990
  ...
921
991
 
922
992
  @typing.overload
923
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
993
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
924
994
  ...
925
995
 
926
- 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]]] = []):
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):
927
997
  """
928
- Specifies secrets to be retrieved and injected as environment variables prior to
929
- the execution of a step.
930
-
998
+ Creates a human-readable report, a Metaflow Card, after this step completes.
931
999
 
932
- Parameters
933
- ----------
934
- sources : List[Union[str, Dict[str, Any]]], default: []
935
- List of secret specs, defining how the secrets are to be retrieved
936
- """
937
- ...
938
-
939
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
940
- """
941
- Specifies that this step is used to deploy an instance of the app.
942
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
1000
+ Note that you may add multiple `@card` decorators in a step with different parameters.
943
1001
 
944
1002
 
945
1003
  Parameters
946
1004
  ----------
947
- app_port : int
948
- Number of GPUs to use.
949
- app_name : str
950
- Name of the app to deploy.
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.
951
1013
  """
952
1014
  ...
953
1015
 
@@ -1030,185 +1092,128 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
1030
1092
  """
1031
1093
  ...
1032
1094
 
1033
- 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]]]:
1034
- """
1035
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
1036
-
1037
- User code call
1038
- --------------
1039
- @ollama(
1040
- models=[...],
1041
- ...
1042
- )
1043
-
1044
- Valid backend options
1045
- ---------------------
1046
- - 'local': Run as a separate process on the local task machine.
1047
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
1048
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
1049
-
1050
- Valid model options
1051
- -------------------
1052
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
1053
-
1054
-
1055
- Parameters
1056
- ----------
1057
- models: list[str]
1058
- List of Ollama containers running models in sidecars.
1059
- backend: str
1060
- Determines where and how to run the Ollama process.
1061
- force_pull: bool
1062
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
1063
- cache_update_policy: str
1064
- Cache update policy: "auto", "force", or "never".
1065
- force_cache_update: bool
1066
- Simple override for "force" cache update policy.
1067
- debug: bool
1068
- Whether to turn on verbose debugging logs.
1069
- circuit_breaker_config: dict
1070
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
1071
- timeout_config: dict
1072
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
1073
- """
1074
- ...
1075
-
1076
1095
  @typing.overload
1077
- 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]]]:
1096
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1078
1097
  """
1079
- Specifies a timeout for your step.
1080
-
1081
- This decorator is useful if this step may hang indefinitely.
1082
-
1083
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1084
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1085
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1086
-
1087
- Note that all the values specified in parameters are added together so if you specify
1088
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1089
-
1090
-
1091
- Parameters
1092
- ----------
1093
- seconds : int, default 0
1094
- Number of seconds to wait prior to timing out.
1095
- minutes : int, default 0
1096
- Number of minutes to wait prior to timing out.
1097
- hours : int, default 0
1098
- Number of hours to wait prior to timing out.
1098
+ Internal decorator to support Fast bakery
1099
1099
  """
1100
1100
  ...
1101
1101
 
1102
1102
  @typing.overload
1103
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1104
- ...
1105
-
1106
- @typing.overload
1107
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1103
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1108
1104
  ...
1109
1105
 
1110
- 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):
1106
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1111
1107
  """
1112
- Specifies a timeout for your step.
1113
-
1114
- This decorator is useful if this step may hang indefinitely.
1115
-
1116
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1117
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1118
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1119
-
1120
- Note that all the values specified in parameters are added together so if you specify
1121
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1122
-
1123
-
1124
- Parameters
1125
- ----------
1126
- seconds : int, default 0
1127
- Number of seconds to wait prior to timing out.
1128
- minutes : int, default 0
1129
- Number of minutes to wait prior to timing out.
1130
- hours : int, default 0
1131
- Number of hours to wait prior to timing out.
1108
+ Internal decorator to support Fast bakery
1132
1109
  """
1133
1110
  ...
1134
1111
 
1135
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
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]]]:
1136
1113
  """
1137
- Decorator that helps cache, version and store models/datasets from huggingface hub.
1138
-
1139
- > Examples
1140
-
1141
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
1142
- ```python
1143
- @huggingface_hub
1144
- @step
1145
- def pull_model_from_huggingface(self):
1146
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
1147
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
1148
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
1149
- # value of the function is a reference to the model in the backend storage.
1150
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
1151
-
1152
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
1153
- self.llama_model = current.huggingface_hub.snapshot_download(
1154
- repo_id=self.model_id,
1155
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1156
- )
1157
- self.next(self.train)
1158
- ```
1159
-
1160
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
1161
- ```python
1162
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1163
- @step
1164
- def pull_model_from_huggingface(self):
1165
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1166
- ```
1167
-
1168
- ```python
1169
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1170
- @step
1171
- def finetune_model(self):
1172
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1173
- # path_to_model will be /my-directory
1174
- ```
1175
-
1176
- ```python
1177
- # Takes all the arguments passed to `snapshot_download`
1178
- # except for `local_dir`
1179
- @huggingface_hub(load=[
1180
- {
1181
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1182
- },
1183
- {
1184
- "repo_id": "myorg/mistral-lora",
1185
- "repo_type": "model",
1186
- },
1187
- ])
1188
- @step
1189
- def finetune_model(self):
1190
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1191
- # path_to_model will be /my-directory
1192
- ```
1114
+ Specifies that this step should execute on Kubernetes.
1193
1115
 
1194
1116
 
1195
1117
  Parameters
1196
1118
  ----------
1197
- temp_dir_root : str, optional
1198
- The root directory that will hold the temporary directory where objects will be downloaded.
1199
-
1200
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1201
- The list of repos (models/datasets) to load.
1202
-
1203
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1204
-
1205
- - If repo (model/dataset) is not found in the datastore:
1206
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1207
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1208
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
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
1209
1190
 
1210
- - If repo is found in the datastore:
1211
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
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().
1212
1217
  """
1213
1218
  ...
1214
1219
 
@@ -1246,53 +1251,136 @@ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], Non
1246
1251
  ...
1247
1252
 
1248
1253
  @typing.overload
1249
- 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]]:
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]]:
1250
1255
  """
1251
- Specifies the times when the flow should be run when running on a
1252
- production scheduler.
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
+ ```
1253
1285
 
1254
1286
 
1255
1287
  Parameters
1256
1288
  ----------
1257
- hourly : bool, default False
1258
- Run the workflow hourly.
1259
- daily : bool, default True
1260
- Run the workflow daily.
1261
- weekly : bool, default False
1262
- Run the workflow weekly.
1263
- cron : str, optional, default None
1264
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1265
- specified by this expression.
1266
- timezone : str, optional, default None
1267
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1268
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
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.
1269
1295
  """
1270
1296
  ...
1271
1297
 
1272
1298
  @typing.overload
1273
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1299
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1274
1300
  ...
1275
1301
 
1276
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
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] = {}):
1277
1303
  """
1278
- Specifies the times when the flow should be run when running on a
1279
- production scheduler.
1304
+ Specifies the event(s) that this flow depends on.
1305
+
1306
+ ```
1307
+ @trigger(event='foo')
1308
+ ```
1309
+ or
1310
+ ```
1311
+ @trigger(events=['foo', 'bar'])
1312
+ ```
1313
+
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
+ ```
1324
+
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
+ ```
1280
1333
 
1281
1334
 
1282
1335
  Parameters
1283
1336
  ----------
1284
- hourly : bool, default False
1285
- Run the workflow hourly.
1286
- daily : bool, default True
1287
- Run the workflow daily.
1288
- weekly : bool, default False
1289
- Run the workflow weekly.
1290
- cron : str, optional, default None
1291
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1292
- specified by this expression.
1293
- timezone : str, optional, default None
1294
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1295
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
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.
1343
+ """
1344
+ ...
1345
+
1346
+ @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
+ """
1363
+ ...
1364
+
1365
+ @typing.overload
1366
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1367
+ ...
1368
+
1369
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1370
+ """
1371
+ Specifies the PyPI packages for all steps of the flow.
1372
+
1373
+ Use `@pypi_base` to set common packages required by all
1374
+ steps and use `@pypi` to specify step-specific overrides.
1375
+
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.
1296
1384
  """
1297
1385
  ...
1298
1386
 
@@ -1339,54 +1427,38 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1339
1427
  """
1340
1428
  ...
1341
1429
 
1342
- @typing.overload
1343
- 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]]:
1430
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1344
1431
  """
1345
- Specifies the Conda environment for all steps of the flow.
1432
+ Specifies what flows belong to the same project.
1346
1433
 
1347
- Use `@conda_base` to set common libraries required by all
1348
- steps and use `@conda` to specify step-specific additions.
1434
+ A project-specific namespace is created for all flows that
1435
+ use the same `@project(name)`.
1349
1436
 
1350
1437
 
1351
1438
  Parameters
1352
1439
  ----------
1353
- packages : Dict[str, str], default {}
1354
- Packages to use for this flow. The key is the name of the package
1355
- and the value is the version to use.
1356
- libraries : Dict[str, str], default {}
1357
- Supported for backward compatibility. When used with packages, packages will take precedence.
1358
- python : str, optional, default None
1359
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1360
- that the version used will correspond to the version of the Python interpreter used to start the run.
1361
- disabled : bool, default False
1362
- If set to True, disables Conda.
1363
- """
1364
- ...
1365
-
1366
- @typing.overload
1367
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1368
- ...
1369
-
1370
- 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):
1371
- """
1372
- Specifies the Conda environment for all steps of the flow.
1373
-
1374
- Use `@conda_base` to set common libraries required by all
1375
- steps and use `@conda` to specify step-specific additions.
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.
1376
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.
1377
1450
 
1378
- Parameters
1379
- ----------
1380
- packages : Dict[str, str], default {}
1381
- Packages to use for this flow. The key is the name of the package
1382
- and the value is the version to use.
1383
- libraries : Dict[str, str], default {}
1384
- Supported for backward compatibility. When used with packages, packages will take precedence.
1385
- python : str, optional, default None
1386
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1387
- that the version used will correspond to the version of the Python interpreter used to start the run.
1388
- disabled : bool, default False
1389
- If set to True, disables Conda.
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>`
1390
1462
  """
1391
1463
  ...
1392
1464
 
@@ -1492,136 +1564,96 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1492
1564
  ...
1493
1565
 
1494
1566
  @typing.overload
1495
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
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]]:
1496
1568
  """
1497
- Specifies the PyPI packages for all steps of the flow.
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.
1498
1573
 
1499
- Use `@pypi_base` to set common packages required by all
1500
- steps and use `@pypi` to specify step-specific overrides.
1501
1574
 
1502
1575
  Parameters
1503
1576
  ----------
1504
- packages : Dict[str, str], default: {}
1577
+ packages : Dict[str, str], default {}
1505
1578
  Packages to use for this flow. The key is the name of the package
1506
1579
  and the value is the version to use.
1507
- python : str, optional, default: None
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
1508
1583
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1509
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.
1510
1587
  """
1511
1588
  ...
1512
1589
 
1513
1590
  @typing.overload
1514
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1591
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1515
1592
  ...
1516
1593
 
1517
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
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):
1518
1595
  """
1519
- Specifies the PyPI packages for all steps of the flow.
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.
1520
1600
 
1521
- Use `@pypi_base` to set common packages required by all
1522
- steps and use `@pypi` to specify step-specific overrides.
1523
1601
 
1524
1602
  Parameters
1525
1603
  ----------
1526
- packages : Dict[str, str], default: {}
1604
+ packages : Dict[str, str], default {}
1527
1605
  Packages to use for this flow. The key is the name of the package
1528
1606
  and the value is the version to use.
1529
- python : str, optional, default: None
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
1530
1610
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1531
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.
1532
1614
  """
1533
1615
  ...
1534
1616
 
1535
- @typing.overload
1536
- 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]]:
1537
- """
1538
- Specifies the event(s) that this flow depends on.
1539
-
1540
- ```
1541
- @trigger(event='foo')
1542
- ```
1543
- or
1544
- ```
1545
- @trigger(events=['foo', 'bar'])
1546
- ```
1547
-
1548
- Additionally, you can specify the parameter mappings
1549
- to map event payload to Metaflow parameters for the flow.
1550
- ```
1551
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1552
- ```
1553
- or
1554
- ```
1555
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1556
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1557
- ```
1558
-
1559
- 'parameters' can also be a list of strings and tuples like so:
1560
- ```
1561
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1562
- ```
1563
- This is equivalent to:
1564
- ```
1565
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1566
- ```
1567
-
1568
-
1569
- Parameters
1570
- ----------
1571
- event : Union[str, Dict[str, Any]], optional, default None
1572
- Event dependency for this flow.
1573
- events : List[Union[str, Dict[str, Any]]], default []
1574
- Events dependency for this flow.
1575
- options : Dict[str, Any], default {}
1576
- Backend-specific configuration for tuning eventing behavior.
1577
- """
1578
- ...
1579
-
1580
- @typing.overload
1581
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1582
- ...
1583
-
1584
- 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] = {}):
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]]:
1585
1618
  """
1586
- Specifies the event(s) that this flow depends on.
1587
-
1588
- ```
1589
- @trigger(event='foo')
1590
- ```
1591
- or
1592
- ```
1593
- @trigger(events=['foo', 'bar'])
1594
- ```
1595
-
1596
- Additionally, you can specify the parameter mappings
1597
- to map event payload to Metaflow parameters for the flow.
1598
- ```
1599
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1600
- ```
1601
- or
1602
- ```
1603
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1604
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1605
- ```
1606
-
1607
- 'parameters' can also be a list of strings and tuples like so:
1608
- ```
1609
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1610
- ```
1611
- This is equivalent to:
1612
- ```
1613
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1614
- ```
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.
1615
1621
 
1616
1622
 
1617
1623
  Parameters
1618
1624
  ----------
1619
- event : Union[str, Dict[str, Any]], optional, default None
1620
- Event dependency for this flow.
1621
- events : List[Union[str, Dict[str, Any]]], default []
1622
- Events dependency for this flow.
1623
- options : Dict[str, Any], default {}
1624
- Backend-specific configuration for tuning eventing behavior.
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)
1625
1657
  """
1626
1658
  ...
1627
1659
 
@@ -1739,81 +1771,54 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1739
1771
  """
1740
1772
  ...
1741
1773
 
1742
- 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]]:
1774
+ @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]]:
1743
1776
  """
1744
- 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.
1745
- 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.
1777
+ Specifies the times when the flow should be run when running on a
1778
+ production scheduler.
1746
1779
 
1747
1780
 
1748
1781
  Parameters
1749
1782
  ----------
1750
- timeout : int
1751
- Time, in seconds before the task times out and fails. (Default: 3600)
1752
- poke_interval : int
1753
- Time in seconds that the job should wait in between each try. (Default: 60)
1754
- mode : str
1755
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1756
- exponential_backoff : bool
1757
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1758
- pool : str
1759
- the slot pool this task should run in,
1760
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1761
- soft_fail : bool
1762
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1763
- name : str
1764
- Name of the sensor on Airflow
1765
- description : str
1766
- Description of sensor in the Airflow UI
1767
- external_dag_id : str
1768
- The dag_id that contains the task you want to wait for.
1769
- external_task_ids : List[str]
1770
- The list of task_ids that you want to wait for.
1771
- If None (default value) the sensor waits for the DAG. (Default: None)
1772
- allowed_states : List[str]
1773
- Iterable of allowed states, (Default: ['success'])
1774
- failed_states : List[str]
1775
- Iterable of failed or dis-allowed states. (Default: None)
1776
- execution_delta : datetime.timedelta
1777
- time difference with the previous execution to look at,
1778
- the default is the same logical date as the current task or DAG. (Default: None)
1779
- check_existence: bool
1780
- Set to True to check if the external task exists or check if
1781
- the DAG to wait for exists. (Default: True)
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).
1782
1795
  """
1783
1796
  ...
1784
1797
 
1785
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1798
+ @typing.overload
1799
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1800
+ ...
1801
+
1802
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1786
1803
  """
1787
- Specifies what flows belong to the same project.
1788
-
1789
- A project-specific namespace is created for all flows that
1790
- use the same `@project(name)`.
1804
+ Specifies the times when the flow should be run when running on a
1805
+ production scheduler.
1791
1806
 
1792
1807
 
1793
1808
  Parameters
1794
1809
  ----------
1795
- name : str
1796
- Project name. Make sure that the name is unique amongst all
1797
- projects that use the same production scheduler. The name may
1798
- contain only lowercase alphanumeric characters and underscores.
1799
-
1800
- branch : Optional[str], default None
1801
- The branch to use. If not specified, the branch is set to
1802
- `user.<username>` unless `production` is set to `True`. This can
1803
- also be set on the command line using `--branch` as a top-level option.
1804
- It is an error to specify `branch` in the decorator and on the command line.
1805
-
1806
- production : bool, default False
1807
- Whether or not the branch is the production branch. This can also be set on the
1808
- command line using `--production` as a top-level option. It is an error to specify
1809
- `production` in the decorator and on the command line.
1810
- The project branch name will be:
1811
- - if `branch` is specified:
1812
- - if `production` is True: `prod.<branch>`
1813
- - if `production` is False: `test.<branch>`
1814
- - if `branch` is not specified:
1815
- - if `production` is True: `prod`
1816
- - if `production` is False: `user.<username>`
1810
+ hourly : bool, default False
1811
+ Run the workflow hourly.
1812
+ daily : bool, default True
1813
+ Run the workflow daily.
1814
+ weekly : bool, default False
1815
+ Run the workflow weekly.
1816
+ cron : str, optional, default None
1817
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1818
+ specified by this expression.
1819
+ timezone : str, optional, default None
1820
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1821
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1817
1822
  """
1818
1823
  ...
1819
1824