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