ob-metaflow-stubs 6.0.3.188rc4__py2.py3-none-any.whl → 6.0.4.1rc0__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 +989 -979
  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 +8 -9
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +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 +61 -61
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +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 +4 -4
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +3 -3
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +4 -4
  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 +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +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 +3 -3
  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 +3 -3
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +4 -4
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +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 +30 -28
  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 +4 -4
  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 +3 -3
  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 +2 -2
  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 +3 -3
  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 +15 -13
  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 +8 -35
  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 +4 -4
  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 +5 -5
  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 +4 -4
  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 +5 -5
  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 +2 -2
  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 +3 -3
  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 +5 -5
  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 +5 -4
  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 +6 -2
  220. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +4 -4
  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 +3 -3
  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 +3 -2
  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 +3 -3
  237. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  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 +6 -8
  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.1rc0.dist-info}/METADATA +1 -1
  257. ob_metaflow_stubs-6.0.4.1rc0.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.1rc0.dist-info}/WHEEL +0 -0
  261. {ob_metaflow_stubs-6.0.3.188rc4.dist-info → ob_metaflow_stubs-6.0.4.1rc0.dist-info}/top_level.txt +0 -0
@@ -1,19 +1,20 @@
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:03:25.730478 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
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,20 +35,22 @@ 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
40
43
  from . import tuple_util as tuple_util
41
44
  from . import events as events
45
+ from . import metaflow_git as metaflow_git
42
46
  from . import runner as runner
43
47
  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
47
51
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
- from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
52
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
53
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_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,452 +162,412 @@ 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
+ @typing.overload
166
+ 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]]]:
160
167
  """
161
- Specifies that this step should execute on Kubernetes.
168
+ Specifies that the step will success under all circumstances.
169
+
170
+ The decorator will create an optional artifact, specified by `var`, which
171
+ contains the exception raised. You can use it to detect the presence
172
+ of errors, indicating that all happy-path artifacts produced by the step
173
+ are missing.
162
174
 
163
175
 
164
176
  Parameters
165
177
  ----------
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
178
+ var : str, optional, default None
179
+ Name of the artifact in which to store the caught exception.
180
+ If not specified, the exception is not stored.
181
+ print_exception : bool, default True
182
+ Determines whether or not the exception is printed to
183
+ stdout when caught.
245
184
  """
246
185
  ...
247
186
 
248
187
  @typing.overload
249
- def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
188
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
189
+ ...
190
+
191
+ @typing.overload
192
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
193
+ ...
194
+
195
+ 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):
250
196
  """
251
- Specifies the Conda environment for the step.
197
+ Specifies that the step will success under all circumstances.
252
198
 
253
- Information in this decorator will augment any
254
- attributes set in the `@conda_base` flow-level decorator. Hence,
255
- you can use `@conda_base` to set packages required by all
256
- steps and use `@conda` to specify step-specific overrides.
199
+ The decorator will create an optional artifact, specified by `var`, which
200
+ contains the exception raised. You can use it to detect the presence
201
+ of errors, indicating that all happy-path artifacts produced by the step
202
+ are missing.
257
203
 
258
204
 
259
205
  Parameters
260
206
  ----------
261
- packages : Dict[str, str], default {}
262
- Packages to use for this step. The key is the name of the package
263
- and the value is the version to use.
264
- libraries : Dict[str, str], default {}
265
- Supported for backward compatibility. When used with packages, packages will take precedence.
266
- python : str, optional, default None
267
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
268
- that the version used will correspond to the version of the Python interpreter used to start the run.
269
- disabled : bool, default False
270
- If set to True, disables @conda.
207
+ var : str, optional, default None
208
+ Name of the artifact in which to store the caught exception.
209
+ If not specified, the exception is not stored.
210
+ print_exception : bool, default True
211
+ Determines whether or not the exception is printed to
212
+ stdout when caught.
271
213
  """
272
214
  ...
273
215
 
274
216
  @typing.overload
275
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
217
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
218
+ """
219
+ Decorator prototype for all step decorators. This function gets specialized
220
+ and imported for all decorators types by _import_plugin_decorators().
221
+ """
276
222
  ...
277
223
 
278
224
  @typing.overload
279
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
225
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
280
226
  ...
281
227
 
282
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
228
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
283
229
  """
284
- Specifies the Conda environment for the step.
285
-
286
- Information in this decorator will augment any
287
- attributes set in the `@conda_base` flow-level decorator. Hence,
288
- you can use `@conda_base` to set packages required by all
289
- steps and use `@conda` to specify step-specific overrides.
290
-
291
-
292
- Parameters
293
- ----------
294
- packages : Dict[str, str], default {}
295
- Packages to use for this step. The key is the name of the package
296
- and the value is the version to use.
297
- libraries : Dict[str, str], default {}
298
- Supported for backward compatibility. When used with packages, packages will take precedence.
299
- python : str, optional, default None
300
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
301
- that the version used will correspond to the version of the Python interpreter used to start the run.
302
- disabled : bool, default False
303
- If set to True, disables @conda.
230
+ Decorator prototype for all step decorators. This function gets specialized
231
+ and imported for all decorators types by _import_plugin_decorators().
304
232
  """
305
233
  ...
306
234
 
307
235
  @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]]]:
236
+ 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]]]:
309
237
  """
310
- Specifies the resources needed when executing this step.
238
+ Specifies the number of times the task corresponding
239
+ to a step needs to be retried.
311
240
 
312
- Use `@resources` to specify the resource requirements
313
- independently of the specific compute layer (`@batch`, `@kubernetes`).
241
+ This decorator is useful for handling transient errors, such as networking issues.
242
+ If your task contains operations that can't be retried safely, e.g. database updates,
243
+ it is advisable to annotate it with `@retry(times=0)`.
314
244
 
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`.
245
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
246
+ decorator will execute a no-op task after all retries have been exhausted,
247
+ ensuring that the flow execution can continue.
325
248
 
326
249
 
327
250
  Parameters
328
251
  ----------
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.
252
+ times : int, default 3
253
+ Number of times to retry this task.
254
+ minutes_between_retries : int, default 2
255
+ Number of minutes between retries.
340
256
  """
341
257
  ...
342
258
 
343
259
  @typing.overload
344
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
260
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
345
261
  ...
346
262
 
347
263
  @typing.overload
348
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
264
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
349
265
  ...
350
266
 
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):
267
+ 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):
352
268
  """
353
- Specifies the resources needed when executing this step.
269
+ Specifies the number of times the task corresponding
270
+ to a step needs to be retried.
354
271
 
355
- Use `@resources` to specify the resource requirements
356
- independently of the specific compute layer (`@batch`, `@kubernetes`).
272
+ This decorator is useful for handling transient errors, such as networking issues.
273
+ If your task contains operations that can't be retried safely, e.g. database updates,
274
+ it is advisable to annotate it with `@retry(times=0)`.
357
275
 
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`.
276
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
277
+ decorator will execute a no-op task after all retries have been exhausted,
278
+ ensuring that the flow execution can continue.
368
279
 
369
280
 
370
281
  Parameters
371
282
  ----------
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.
283
+ times : int, default 3
284
+ Number of times to retry this task.
285
+ minutes_between_retries : int, default 2
286
+ Number of minutes between retries.
383
287
  """
384
288
  ...
385
289
 
386
- @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]]]:
290
+ 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]]]:
388
291
  """
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.
292
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
392
293
 
294
+ User code call
295
+ --------------
296
+ @ollama(
297
+ models=[...],
298
+ ...
299
+ )
393
300
 
394
- Parameters
395
- ----------
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.
404
- """
405
- ...
406
-
407
- @typing.overload
408
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
409
- ...
410
-
411
- @typing.overload
412
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
413
- ...
414
-
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):
416
- """
417
- Creates a human-readable report, a Metaflow Card, after this step completes.
301
+ Valid backend options
302
+ ---------------------
303
+ - 'local': Run as a separate process on the local task machine.
304
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
305
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
418
306
 
419
- Note that you may add multiple `@card` decorators in a step with different parameters.
307
+ Valid model options
308
+ -------------------
309
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
420
310
 
421
311
 
422
312
  Parameters
423
313
  ----------
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.
314
+ models: list[str]
315
+ List of Ollama containers running models in sidecars.
316
+ backend: str
317
+ Determines where and how to run the Ollama process.
318
+ force_pull: bool
319
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
320
+ cache_update_policy: str
321
+ Cache update policy: "auto", "force", or "never".
322
+ force_cache_update: bool
323
+ Simple override for "force" cache update policy.
324
+ debug: bool
325
+ Whether to turn on verbose debugging logs.
326
+ circuit_breaker_config: dict
327
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
328
+ timeout_config: dict
329
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
432
330
  """
433
331
  ...
434
332
 
435
333
  @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]]]:
334
+ 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]]]:
437
335
  """
438
- Specifies that the step will success under all circumstances.
439
-
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.
336
+ Specifies secrets to be retrieved and injected as environment variables prior to
337
+ the execution of a step.
444
338
 
445
339
 
446
340
  Parameters
447
341
  ----------
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.
342
+ sources : List[Union[str, Dict[str, Any]]], default: []
343
+ List of secret specs, defining how the secrets are to be retrieved
344
+ role : str, optional, default: None
345
+ Role to use for fetching secrets
454
346
  """
455
347
  ...
456
348
 
457
349
  @typing.overload
458
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
350
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
459
351
  ...
460
352
 
461
353
  @typing.overload
462
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
354
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
463
355
  ...
464
356
 
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):
357
+ 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):
466
358
  """
467
- Specifies that the step will success under all circumstances.
468
-
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.
359
+ Specifies secrets to be retrieved and injected as environment variables prior to
360
+ the execution of a step.
473
361
 
474
362
 
475
363
  Parameters
476
364
  ----------
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.
483
- """
484
- ...
485
-
486
- @typing.overload
487
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
488
- """
489
- Internal decorator to support Fast bakery
365
+ sources : List[Union[str, Dict[str, Any]]], default: []
366
+ List of secret specs, defining how the secrets are to be retrieved
367
+ role : str, optional, default: None
368
+ Role to use for fetching secrets
490
369
  """
491
370
  ...
492
371
 
493
- @typing.overload
494
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
495
- ...
496
-
497
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
372
+ 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]]]:
498
373
  """
499
- Internal decorator to support Fast bakery
374
+ Specifies that this step should execute on DGX cloud.
375
+
376
+
377
+ Parameters
378
+ ----------
379
+ gpu : int
380
+ Number of GPUs to use.
381
+ gpu_type : str
382
+ Type of Nvidia GPU to use.
500
383
  """
501
384
  ...
502
385
 
503
386
  @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]]]:
387
+ 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]]]:
505
388
  """
506
- Specifies the number of times the task corresponding
507
- to a step needs to be retried.
389
+ Enables loading / saving of models within a step.
508
390
 
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)`.
391
+ > Examples
392
+ - Saving Models
393
+ ```python
394
+ @model
395
+ @step
396
+ def train(self):
397
+ # current.model.save returns a dictionary reference to the model saved
398
+ self.my_model = current.model.save(
399
+ path_to_my_model,
400
+ label="my_model",
401
+ metadata={
402
+ "epochs": 10,
403
+ "batch-size": 32,
404
+ "learning-rate": 0.001,
405
+ }
406
+ )
407
+ self.next(self.test)
512
408
 
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.
409
+ @model(load="my_model")
410
+ @step
411
+ def test(self):
412
+ # `current.model.loaded` returns a dictionary of the loaded models
413
+ # where the key is the name of the artifact and the value is the path to the model
414
+ print(os.listdir(current.model.loaded["my_model"]))
415
+ self.next(self.end)
416
+ ```
417
+
418
+ - Loading models
419
+ ```python
420
+ @step
421
+ def train(self):
422
+ # current.model.load returns the path to the model loaded
423
+ checkpoint_path = current.model.load(
424
+ self.checkpoint_key,
425
+ )
426
+ model_path = current.model.load(
427
+ self.model,
428
+ )
429
+ self.next(self.test)
430
+ ```
516
431
 
517
432
 
518
433
  Parameters
519
434
  ----------
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.
435
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
436
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
437
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
438
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
439
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
440
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
441
+
442
+ temp_dir_root : str, default: None
443
+ The root directory under which `current.model.loaded` will store loaded models
524
444
  """
525
445
  ...
526
446
 
527
447
  @typing.overload
528
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
448
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
529
449
  ...
530
450
 
531
451
  @typing.overload
532
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
452
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
533
453
  ...
534
454
 
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):
455
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
536
456
  """
537
- Specifies the number of times the task corresponding
538
- to a step needs to be retried.
457
+ Enables loading / saving of models within a step.
539
458
 
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)`.
459
+ > Examples
460
+ - Saving Models
461
+ ```python
462
+ @model
463
+ @step
464
+ def train(self):
465
+ # current.model.save returns a dictionary reference to the model saved
466
+ self.my_model = current.model.save(
467
+ path_to_my_model,
468
+ label="my_model",
469
+ metadata={
470
+ "epochs": 10,
471
+ "batch-size": 32,
472
+ "learning-rate": 0.001,
473
+ }
474
+ )
475
+ self.next(self.test)
543
476
 
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.
477
+ @model(load="my_model")
478
+ @step
479
+ def test(self):
480
+ # `current.model.loaded` returns a dictionary of the loaded models
481
+ # where the key is the name of the artifact and the value is the path to the model
482
+ print(os.listdir(current.model.loaded["my_model"]))
483
+ self.next(self.end)
484
+ ```
485
+
486
+ - Loading models
487
+ ```python
488
+ @step
489
+ def train(self):
490
+ # current.model.load returns the path to the model loaded
491
+ checkpoint_path = current.model.load(
492
+ self.checkpoint_key,
493
+ )
494
+ model_path = current.model.load(
495
+ self.model,
496
+ )
497
+ self.next(self.test)
498
+ ```
547
499
 
548
500
 
549
501
  Parameters
550
502
  ----------
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.
503
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
504
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
505
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
506
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
507
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
508
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
509
+
510
+ temp_dir_root : str, default: None
511
+ The root directory under which `current.model.loaded` will store loaded models
555
512
  """
556
513
  ...
557
514
 
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]]]:
559
- """
560
- This decorator is used to run vllm APIs as Metaflow task sidecars.
515
+ @typing.overload
516
+ 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]]]:
517
+ """
518
+ Specifies a timeout for your step.
561
519
 
562
- User code call
563
- --------------
564
- @vllm(
565
- model="...",
566
- ...
567
- )
520
+ This decorator is useful if this step may hang indefinitely.
568
521
 
569
- Valid backend options
570
- ---------------------
571
- - 'local': Run as a separate process on the local task machine.
522
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
523
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
524
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
572
525
 
573
- Valid model options
574
- -------------------
575
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
526
+ Note that all the values specified in parameters are added together so if you specify
527
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
576
528
 
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.
529
+
530
+ Parameters
531
+ ----------
532
+ seconds : int, default 0
533
+ Number of seconds to wait prior to timing out.
534
+ minutes : int, default 0
535
+ Number of minutes to wait prior to timing out.
536
+ hours : int, default 0
537
+ Number of hours to wait prior to timing out.
538
+ """
539
+ ...
540
+
541
+ @typing.overload
542
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
543
+ ...
544
+
545
+ @typing.overload
546
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
+ ...
548
+
549
+ 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):
550
+ """
551
+ Specifies a timeout for your step.
552
+
553
+ This decorator is useful if this step may hang indefinitely.
554
+
555
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
556
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
557
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
558
+
559
+ Note that all the values specified in parameters are added together so if you specify
560
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
579
561
 
580
562
 
581
563
  Parameters
582
564
  ----------
583
- model: str
584
- HuggingFace model identifier to be served by vLLM.
585
- backend: str
586
- Determines where and how to run the vLLM process.
587
- openai_api_server: bool
588
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
589
- Default is False (uses native engine).
590
- Set to True for backward compatibility with existing code.
591
- debug: bool
592
- Whether to turn on verbose debugging logs.
593
- card_refresh_interval: int
594
- Interval in seconds for refreshing the vLLM status card.
595
- Only used when openai_api_server=True.
596
- max_retries: int
597
- Maximum number of retries checking for vLLM server startup.
598
- Only used when openai_api_server=True.
599
- retry_alert_frequency: int
600
- Frequency of alert logs for vLLM server startup retries.
601
- Only used when openai_api_server=True.
602
- engine_args : dict
603
- Additional keyword arguments to pass to the vLLM engine.
604
- For example, `tensor_parallel_size=2`.
565
+ seconds : int, default 0
566
+ Number of seconds to wait prior to timing out.
567
+ minutes : int, default 0
568
+ Number of minutes to wait prior to timing out.
569
+ hours : int, default 0
570
+ Number of hours to wait prior to timing out.
605
571
  """
606
572
  ...
607
573
 
@@ -686,192 +652,140 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
686
652
  ...
687
653
 
688
654
  @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]]]:
655
+ def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
690
656
  """
691
- Enables checkpointing for a step.
692
-
693
- > Examples
694
-
695
- - Saving Checkpoints
696
-
697
- ```python
698
- @checkpoint
699
- @step
700
- def train(self):
701
- model = create_model(self.parameters, checkpoint_path = None)
702
- for i in range(self.epochs):
703
- # some training logic
704
- loss = model.train(self.dataset)
705
- if i % 10 == 0:
706
- model.save(
707
- current.checkpoint.directory,
708
- )
709
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
710
- # and returns a reference dictionary to the checkpoint saved in the datastore
711
- self.latest_checkpoint = current.checkpoint.save(
712
- name="epoch_checkpoint",
713
- metadata={
714
- "epoch": i,
715
- "loss": loss,
716
- }
717
- )
718
- ```
719
-
720
- - Using Loaded Checkpoints
721
-
722
- ```python
723
- @retry(times=3)
724
- @checkpoint
725
- @step
726
- def train(self):
727
- # Assume that the task has restarted and the previous attempt of the task
728
- # saved a checkpoint
729
- checkpoint_path = None
730
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
731
- print("Loaded checkpoint from the previous attempt")
732
- checkpoint_path = current.checkpoint.directory
657
+ Specifies the Conda environment for the step.
733
658
 
734
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
735
- for i in range(self.epochs):
736
- ...
737
- ```
659
+ Information in this decorator will augment any
660
+ attributes set in the `@conda_base` flow-level decorator. Hence,
661
+ you can use `@conda_base` to set packages required by all
662
+ steps and use `@conda` to specify step-specific overrides.
738
663
 
739
664
 
740
665
  Parameters
741
666
  ----------
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.
752
-
753
- temp_dir_root : str, default: None
754
- The root directory under which `current.checkpoint.directory` will be created.
667
+ packages : Dict[str, str], default {}
668
+ Packages to use for this step. The key is the name of the package
669
+ and the value is the version to use.
670
+ libraries : Dict[str, str], default {}
671
+ Supported for backward compatibility. When used with packages, packages will take precedence.
672
+ python : str, optional, default None
673
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
674
+ that the version used will correspond to the version of the Python interpreter used to start the run.
675
+ disabled : bool, default False
676
+ If set to True, disables @conda.
755
677
  """
756
678
  ...
757
679
 
758
680
  @typing.overload
759
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
681
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
760
682
  ...
761
683
 
762
684
  @typing.overload
763
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
685
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
764
686
  ...
765
687
 
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):
688
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
767
689
  """
768
- Enables checkpointing for a step.
769
-
770
- > Examples
771
-
772
- - Saving Checkpoints
773
-
774
- ```python
775
- @checkpoint
776
- @step
777
- def train(self):
778
- model = create_model(self.parameters, checkpoint_path = None)
779
- for i in range(self.epochs):
780
- # some training logic
781
- loss = model.train(self.dataset)
782
- if i % 10 == 0:
783
- model.save(
784
- current.checkpoint.directory,
785
- )
786
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
787
- # and returns a reference dictionary to the checkpoint saved in the datastore
788
- self.latest_checkpoint = current.checkpoint.save(
789
- name="epoch_checkpoint",
790
- metadata={
791
- "epoch": i,
792
- "loss": loss,
793
- }
794
- )
795
- ```
796
-
797
- - Using Loaded Checkpoints
798
-
799
- ```python
800
- @retry(times=3)
801
- @checkpoint
802
- @step
803
- def train(self):
804
- # Assume that the task has restarted and the previous attempt of the task
805
- # saved a checkpoint
806
- checkpoint_path = None
807
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
808
- print("Loaded checkpoint from the previous attempt")
809
- checkpoint_path = current.checkpoint.directory
690
+ Specifies the Conda environment for the step.
810
691
 
811
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
812
- for i in range(self.epochs):
813
- ...
814
- ```
692
+ Information in this decorator will augment any
693
+ attributes set in the `@conda_base` flow-level decorator. Hence,
694
+ you can use `@conda_base` to set packages required by all
695
+ steps and use `@conda` to specify step-specific overrides.
815
696
 
816
697
 
817
698
  Parameters
818
699
  ----------
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.
700
+ packages : Dict[str, str], default {}
701
+ Packages to use for this step. The key is the name of the package
702
+ and the value is the version to use.
703
+ libraries : Dict[str, str], default {}
704
+ Supported for backward compatibility. When used with packages, packages will take precedence.
705
+ python : str, optional, default None
706
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
707
+ that the version used will correspond to the version of the Python interpreter used to start the run.
708
+ disabled : bool, default False
709
+ If set to True, disables @conda.
832
710
  """
833
711
  ...
834
712
 
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]]]:
713
+ @typing.overload
714
+ 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]]]:
836
715
  """
837
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
716
+ Specifies the resources needed when executing this step.
838
717
 
839
- User code call
840
- --------------
841
- @ollama(
842
- models=[...],
843
- ...
844
- )
718
+ Use `@resources` to specify the resource requirements
719
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
845
720
 
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.
721
+ You can choose the compute layer on the command line by executing e.g.
722
+ ```
723
+ python myflow.py run --with batch
724
+ ```
725
+ or
726
+ ```
727
+ python myflow.py run --with kubernetes
728
+ ```
729
+ which executes the flow on the desired system using the
730
+ requirements specified in `@resources`.
851
731
 
852
- Valid model options
853
- -------------------
854
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
732
+
733
+ Parameters
734
+ ----------
735
+ cpu : int, default 1
736
+ Number of CPUs required for this step.
737
+ gpu : int, optional, default None
738
+ Number of GPUs required for this step.
739
+ disk : int, optional, default None
740
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
741
+ memory : int, default 4096
742
+ Memory size (in MB) required for this step.
743
+ shared_memory : int, optional, default None
744
+ The value for the size (in MiB) of the /dev/shm volume for this step.
745
+ This parameter maps to the `--shm-size` option in Docker.
746
+ """
747
+ ...
748
+
749
+ @typing.overload
750
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
751
+ ...
752
+
753
+ @typing.overload
754
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
755
+ ...
756
+
757
+ 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):
758
+ """
759
+ Specifies the resources needed when executing this step.
760
+
761
+ Use `@resources` to specify the resource requirements
762
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
763
+
764
+ You can choose the compute layer on the command line by executing e.g.
765
+ ```
766
+ python myflow.py run --with batch
767
+ ```
768
+ or
769
+ ```
770
+ python myflow.py run --with kubernetes
771
+ ```
772
+ which executes the flow on the desired system using the
773
+ requirements specified in `@resources`.
855
774
 
856
775
 
857
776
  Parameters
858
777
  ----------
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.
778
+ cpu : int, default 1
779
+ Number of CPUs required for this step.
780
+ gpu : int, optional, default None
781
+ Number of GPUs required for this step.
782
+ disk : int, optional, default None
783
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
784
+ memory : int, default 4096
785
+ Memory size (in MB) required for this step.
786
+ shared_memory : int, optional, default None
787
+ The value for the size (in MiB) of the /dev/shm volume for this step.
788
+ This parameter maps to the `--shm-size` option in Docker.
875
789
  """
876
790
  ...
877
791
 
@@ -927,199 +841,209 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
927
841
  ...
928
842
 
929
843
  @typing.overload
930
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
931
- """
932
- Decorator prototype for all step decorators. This function gets specialized
933
- and imported for all decorators types by _import_plugin_decorators().
934
- """
935
- ...
936
-
937
- @typing.overload
938
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
939
- ...
940
-
941
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
942
- """
943
- Decorator prototype for all step decorators. This function gets specialized
944
- and imported for all decorators types by _import_plugin_decorators().
945
- """
946
- ...
947
-
948
- @typing.overload
949
- def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
844
+ 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]]]:
950
845
  """
951
- Enables loading / saving of models within a step.
952
-
953
- > Examples
954
- - Saving Models
955
- ```python
956
- @model
957
- @step
958
- def train(self):
959
- # current.model.save returns a dictionary reference to the model saved
960
- self.my_model = current.model.save(
961
- path_to_my_model,
962
- label="my_model",
963
- metadata={
964
- "epochs": 10,
965
- "batch-size": 32,
966
- "learning-rate": 0.001,
967
- }
968
- )
969
- self.next(self.test)
970
-
971
- @model(load="my_model")
972
- @step
973
- def test(self):
974
- # `current.model.loaded` returns a dictionary of the loaded models
975
- # where the key is the name of the artifact and the value is the path to the model
976
- print(os.listdir(current.model.loaded["my_model"]))
977
- self.next(self.end)
978
- ```
846
+ Creates a human-readable report, a Metaflow Card, after this step completes.
979
847
 
980
- - Loading models
981
- ```python
982
- @step
983
- def train(self):
984
- # current.model.load returns the path to the model loaded
985
- checkpoint_path = current.model.load(
986
- self.checkpoint_key,
987
- )
988
- model_path = current.model.load(
989
- self.model,
990
- )
991
- self.next(self.test)
992
- ```
848
+ Note that you may add multiple `@card` decorators in a step with different parameters.
993
849
 
994
850
 
995
851
  Parameters
996
852
  ----------
997
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
998
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
999
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1000
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1001
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1002
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1003
-
1004
- temp_dir_root : str, default: None
1005
- The root directory under which `current.model.loaded` will store loaded models
853
+ type : str, default 'default'
854
+ Card type.
855
+ id : str, optional, default None
856
+ If multiple cards are present, use this id to identify this card.
857
+ options : Dict[str, Any], default {}
858
+ Options passed to the card. The contents depend on the card type.
859
+ timeout : int, default 45
860
+ Interrupt reporting if it takes more than this many seconds.
1006
861
  """
1007
862
  ...
1008
863
 
1009
864
  @typing.overload
1010
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
865
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1011
866
  ...
1012
867
 
1013
868
  @typing.overload
1014
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
869
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1015
870
  ...
1016
871
 
1017
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
872
+ 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):
1018
873
  """
1019
- Enables loading / saving of models within a step.
1020
-
1021
- > Examples
1022
- - Saving Models
1023
- ```python
1024
- @model
1025
- @step
1026
- def train(self):
1027
- # current.model.save returns a dictionary reference to the model saved
1028
- self.my_model = current.model.save(
1029
- path_to_my_model,
1030
- label="my_model",
1031
- metadata={
1032
- "epochs": 10,
1033
- "batch-size": 32,
1034
- "learning-rate": 0.001,
1035
- }
1036
- )
1037
- self.next(self.test)
1038
-
1039
- @model(load="my_model")
1040
- @step
1041
- def test(self):
1042
- # `current.model.loaded` returns a dictionary of the loaded models
1043
- # where the key is the name of the artifact and the value is the path to the model
1044
- print(os.listdir(current.model.loaded["my_model"]))
1045
- self.next(self.end)
1046
- ```
1047
-
1048
- - Loading models
1049
- ```python
1050
- @step
1051
- def train(self):
1052
- # current.model.load returns the path to the model loaded
1053
- checkpoint_path = current.model.load(
1054
- self.checkpoint_key,
1055
- )
1056
- model_path = current.model.load(
1057
- self.model,
1058
- )
1059
- self.next(self.test)
1060
- ```
1061
-
1062
-
1063
- Parameters
1064
- ----------
1065
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1066
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1067
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1068
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1069
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1070
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
874
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1071
875
 
1072
- temp_dir_root : str, default: None
1073
- The root directory under which `current.model.loaded` will store loaded models
1074
- """
1075
- ...
1076
-
1077
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1078
- """
1079
- Specifies that this step should execute on DGX cloud.
876
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1080
877
 
1081
878
 
1082
879
  Parameters
1083
880
  ----------
1084
- gpu : int
1085
- Number of GPUs to use.
1086
- gpu_type : str
1087
- Type of Nvidia GPU to use.
881
+ type : str, default 'default'
882
+ Card type.
883
+ id : str, optional, default None
884
+ If multiple cards are present, use this id to identify this card.
885
+ options : Dict[str, Any], default {}
886
+ Options passed to the card. The contents depend on the card type.
887
+ timeout : int, default 45
888
+ Interrupt reporting if it takes more than this many seconds.
1088
889
  """
1089
890
  ...
1090
891
 
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]]]:
892
+ 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]]]:
1093
893
  """
1094
- Specifies secrets to be retrieved and injected as environment variables prior to
1095
- the execution of a step.
894
+ Specifies that this step should execute on Kubernetes.
1096
895
 
1097
896
 
1098
897
  Parameters
1099
898
  ----------
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
-
899
+ cpu : int, default 1
900
+ Number of CPUs required for this step. If `@resources` is
901
+ also present, the maximum value from all decorators is used.
902
+ memory : int, default 4096
903
+ Memory size (in MB) required for this step. If
904
+ `@resources` is also present, the maximum value from all decorators is
905
+ used.
906
+ disk : int, default 10240
907
+ Disk size (in MB) required for this step. If
908
+ `@resources` is also present, the maximum value from all decorators is
909
+ used.
910
+ image : str, optional, default None
911
+ Docker image to use when launching on Kubernetes. If not specified, and
912
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
913
+ not, a default Docker image mapping to the current version of Python is used.
914
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
915
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
916
+ image_pull_secrets: List[str], default []
917
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
918
+ Kubernetes image pull secrets to use when pulling container images
919
+ in Kubernetes.
920
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
921
+ Kubernetes service account to use when launching pod in Kubernetes.
922
+ secrets : List[str], optional, default None
923
+ Kubernetes secrets to use when launching pod in Kubernetes. These
924
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
925
+ in Metaflow configuration.
926
+ node_selector: Union[Dict[str,str], str], optional, default None
927
+ Kubernetes node selector(s) to apply to the pod running the task.
928
+ Can be passed in as a comma separated string of values e.g.
929
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
930
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
931
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
932
+ Kubernetes namespace to use when launching pod in Kubernetes.
933
+ gpu : int, optional, default None
934
+ Number of GPUs required for this step. A value of zero implies that
935
+ the scheduled node should not have GPUs.
936
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
937
+ The vendor of the GPUs to be used for this step.
938
+ tolerations : List[str], default []
939
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
940
+ Kubernetes tolerations to use when launching pod in Kubernetes.
941
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
942
+ Kubernetes labels to use when launching pod in Kubernetes.
943
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
944
+ Kubernetes annotations to use when launching pod in Kubernetes.
945
+ use_tmpfs : bool, default False
946
+ This enables an explicit tmpfs mount for this step.
947
+ tmpfs_tempdir : bool, default True
948
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
949
+ tmpfs_size : int, optional, default: None
950
+ The value for the size (in MiB) of the tmpfs mount for this step.
951
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
952
+ memory allocated for this step.
953
+ tmpfs_path : str, optional, default /metaflow_temp
954
+ Path to tmpfs mount for this step.
955
+ persistent_volume_claims : Dict[str, str], optional, default None
956
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
957
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
958
+ shared_memory: int, optional
959
+ Shared memory size (in MiB) required for this step
960
+ port: int, optional
961
+ Port number to specify in the Kubernetes job object
962
+ compute_pool : str, optional, default None
963
+ Compute pool to be used for for this step.
964
+ If not specified, any accessible compute pool within the perimeter is used.
965
+ hostname_resolution_timeout: int, default 10 * 60
966
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
967
+ Only applicable when @parallel is used.
968
+ qos: str, default: Burstable
969
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
970
+
971
+ security_context: Dict[str, Any], optional, default None
972
+ Container security context. Applies to the task container. Allows the following keys:
973
+ - privileged: bool, optional, default None
974
+ - allow_privilege_escalation: bool, optional, default None
975
+ - run_as_user: int, optional, default None
976
+ - run_as_group: int, optional, default None
977
+ - run_as_non_root: bool, optional, default None
978
+ """
979
+ ...
980
+
1105
981
  @typing.overload
1106
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
982
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
983
+ """
984
+ Decorator prototype for all step decorators. This function gets specialized
985
+ and imported for all decorators types by _import_plugin_decorators().
986
+ """
1107
987
  ...
1108
988
 
1109
989
  @typing.overload
1110
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
990
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1111
991
  ...
1112
992
 
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]]] = []):
993
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1114
994
  """
1115
- Specifies secrets to be retrieved and injected as environment variables prior to
1116
- the execution of a step.
995
+ Decorator prototype for all step decorators. This function gets specialized
996
+ and imported for all decorators types by _import_plugin_decorators().
997
+ """
998
+ ...
999
+
1000
+ @typing.overload
1001
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1002
+ """
1003
+ Internal decorator to support Fast bakery
1004
+ """
1005
+ ...
1006
+
1007
+ @typing.overload
1008
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1009
+ ...
1010
+
1011
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1012
+ """
1013
+ Internal decorator to support Fast bakery
1014
+ """
1015
+ ...
1016
+
1017
+ @typing.overload
1018
+ 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]]]:
1019
+ """
1020
+ Specifies environment variables to be set prior to the execution of a step.
1117
1021
 
1118
1022
 
1119
1023
  Parameters
1120
1024
  ----------
1121
- sources : List[Union[str, Dict[str, Any]]], default: []
1122
- List of secret specs, defining how the secrets are to be retrieved
1025
+ vars : Dict[str, str], default {}
1026
+ Dictionary of environment variables to set.
1027
+ """
1028
+ ...
1029
+
1030
+ @typing.overload
1031
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1032
+ ...
1033
+
1034
+ @typing.overload
1035
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1036
+ ...
1037
+
1038
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1039
+ """
1040
+ Specifies environment variables to be set prior to the execution of a step.
1041
+
1042
+
1043
+ Parameters
1044
+ ----------
1045
+ vars : Dict[str, str], default {}
1046
+ Dictionary of environment variables to set.
1123
1047
  """
1124
1048
  ...
1125
1049
 
@@ -1140,113 +1064,386 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
1140
1064
  ...
1141
1065
 
1142
1066
  @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]]]:
1067
+ 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]]]:
1144
1068
  """
1145
- Specifies environment variables to be set prior to the execution of a step.
1069
+ Enables checkpointing for a step.
1070
+
1071
+ > Examples
1072
+
1073
+ - Saving Checkpoints
1074
+
1075
+ ```python
1076
+ @checkpoint
1077
+ @step
1078
+ def train(self):
1079
+ model = create_model(self.parameters, checkpoint_path = None)
1080
+ for i in range(self.epochs):
1081
+ # some training logic
1082
+ loss = model.train(self.dataset)
1083
+ if i % 10 == 0:
1084
+ model.save(
1085
+ current.checkpoint.directory,
1086
+ )
1087
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
1088
+ # and returns a reference dictionary to the checkpoint saved in the datastore
1089
+ self.latest_checkpoint = current.checkpoint.save(
1090
+ name="epoch_checkpoint",
1091
+ metadata={
1092
+ "epoch": i,
1093
+ "loss": loss,
1094
+ }
1095
+ )
1096
+ ```
1097
+
1098
+ - Using Loaded Checkpoints
1099
+
1100
+ ```python
1101
+ @retry(times=3)
1102
+ @checkpoint
1103
+ @step
1104
+ def train(self):
1105
+ # Assume that the task has restarted and the previous attempt of the task
1106
+ # saved a checkpoint
1107
+ checkpoint_path = None
1108
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1109
+ print("Loaded checkpoint from the previous attempt")
1110
+ checkpoint_path = current.checkpoint.directory
1111
+
1112
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1113
+ for i in range(self.epochs):
1114
+ ...
1115
+ ```
1146
1116
 
1147
1117
 
1148
1118
  Parameters
1149
1119
  ----------
1150
- vars : Dict[str, str], default {}
1151
- Dictionary of environment variables to set.
1120
+ load_policy : str, default: "fresh"
1121
+ The policy for loading the checkpoint. The following policies are supported:
1122
+ - "eager": Loads the the latest available checkpoint within the namespace.
1123
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1124
+ will be loaded at the start of the task.
1125
+ - "none": Do not load any checkpoint
1126
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1127
+ This mode helps loading checkpoints across various retry attempts of the same task.
1128
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1129
+ created within the task will be loaded when the task is retries execution on failure.
1130
+
1131
+ temp_dir_root : str, default: None
1132
+ The root directory under which `current.checkpoint.directory` will be created.
1152
1133
  """
1153
1134
  ...
1154
1135
 
1155
1136
  @typing.overload
1156
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1137
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1157
1138
  ...
1158
1139
 
1159
1140
  @typing.overload
1160
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1141
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1161
1142
  ...
1162
1143
 
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] = {}):
1144
+ 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):
1164
1145
  """
1165
- Specifies environment variables to be set prior to the execution of a step.
1146
+ Enables checkpointing for a step.
1147
+
1148
+ > Examples
1149
+
1150
+ - Saving Checkpoints
1151
+
1152
+ ```python
1153
+ @checkpoint
1154
+ @step
1155
+ def train(self):
1156
+ model = create_model(self.parameters, checkpoint_path = None)
1157
+ for i in range(self.epochs):
1158
+ # some training logic
1159
+ loss = model.train(self.dataset)
1160
+ if i % 10 == 0:
1161
+ model.save(
1162
+ current.checkpoint.directory,
1163
+ )
1164
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
1165
+ # and returns a reference dictionary to the checkpoint saved in the datastore
1166
+ self.latest_checkpoint = current.checkpoint.save(
1167
+ name="epoch_checkpoint",
1168
+ metadata={
1169
+ "epoch": i,
1170
+ "loss": loss,
1171
+ }
1172
+ )
1173
+ ```
1174
+
1175
+ - Using Loaded Checkpoints
1176
+
1177
+ ```python
1178
+ @retry(times=3)
1179
+ @checkpoint
1180
+ @step
1181
+ def train(self):
1182
+ # Assume that the task has restarted and the previous attempt of the task
1183
+ # saved a checkpoint
1184
+ checkpoint_path = None
1185
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1186
+ print("Loaded checkpoint from the previous attempt")
1187
+ checkpoint_path = current.checkpoint.directory
1188
+
1189
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1190
+ for i in range(self.epochs):
1191
+ ...
1192
+ ```
1166
1193
 
1167
1194
 
1168
1195
  Parameters
1169
1196
  ----------
1170
- vars : Dict[str, str], default {}
1171
- Dictionary of environment variables to set.
1197
+ load_policy : str, default: "fresh"
1198
+ The policy for loading the checkpoint. The following policies are supported:
1199
+ - "eager": Loads the the latest available checkpoint within the namespace.
1200
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1201
+ will be loaded at the start of the task.
1202
+ - "none": Do not load any checkpoint
1203
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1204
+ This mode helps loading checkpoints across various retry attempts of the same task.
1205
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1206
+ created within the task will be loaded when the task is retries execution on failure.
1207
+
1208
+ temp_dir_root : str, default: None
1209
+ The root directory under which `current.checkpoint.directory` will be created.
1172
1210
  """
1173
1211
  ...
1174
1212
 
1175
- @typing.overload
1176
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1213
+ 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]]]:
1177
1214
  """
1178
- Decorator prototype for all step decorators. This function gets specialized
1179
- and imported for all decorators types by _import_plugin_decorators().
1215
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
1216
+
1217
+ User code call
1218
+ --------------
1219
+ @vllm(
1220
+ model="...",
1221
+ ...
1222
+ )
1223
+
1224
+ Valid backend options
1225
+ ---------------------
1226
+ - 'local': Run as a separate process on the local task machine.
1227
+
1228
+ Valid model options
1229
+ -------------------
1230
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1231
+
1232
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
1233
+ If you need multiple models, you must create multiple @vllm decorators.
1234
+
1235
+
1236
+ Parameters
1237
+ ----------
1238
+ model: str
1239
+ HuggingFace model identifier to be served by vLLM.
1240
+ backend: str
1241
+ Determines where and how to run the vLLM process.
1242
+ openai_api_server: bool
1243
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
1244
+ Default is False (uses native engine).
1245
+ Set to True for backward compatibility with existing code.
1246
+ debug: bool
1247
+ Whether to turn on verbose debugging logs.
1248
+ card_refresh_interval: int
1249
+ Interval in seconds for refreshing the vLLM status card.
1250
+ Only used when openai_api_server=True.
1251
+ max_retries: int
1252
+ Maximum number of retries checking for vLLM server startup.
1253
+ Only used when openai_api_server=True.
1254
+ retry_alert_frequency: int
1255
+ Frequency of alert logs for vLLM server startup retries.
1256
+ Only used when openai_api_server=True.
1257
+ engine_args : dict
1258
+ Additional keyword arguments to pass to the vLLM engine.
1259
+ For example, `tensor_parallel_size=2`.
1180
1260
  """
1181
1261
  ...
1182
1262
 
1183
- @typing.overload
1184
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1263
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1264
+ """
1265
+ Specifies what flows belong to the same project.
1266
+
1267
+ A project-specific namespace is created for all flows that
1268
+ use the same `@project(name)`.
1269
+
1270
+
1271
+ Parameters
1272
+ ----------
1273
+ name : str
1274
+ Project name. Make sure that the name is unique amongst all
1275
+ projects that use the same production scheduler. The name may
1276
+ contain only lowercase alphanumeric characters and underscores.
1277
+
1278
+ branch : Optional[str], default None
1279
+ The branch to use. If not specified, the branch is set to
1280
+ `user.<username>` unless `production` is set to `True`. This can
1281
+ also be set on the command line using `--branch` as a top-level option.
1282
+ It is an error to specify `branch` in the decorator and on the command line.
1283
+
1284
+ production : bool, default False
1285
+ Whether or not the branch is the production branch. This can also be set on the
1286
+ command line using `--production` as a top-level option. It is an error to specify
1287
+ `production` in the decorator and on the command line.
1288
+ The project branch name will be:
1289
+ - if `branch` is specified:
1290
+ - if `production` is True: `prod.<branch>`
1291
+ - if `production` is False: `test.<branch>`
1292
+ - if `branch` is not specified:
1293
+ - if `production` is True: `prod`
1294
+ - if `production` is False: `user.<username>`
1295
+ """
1185
1296
  ...
1186
1297
 
1187
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1298
+ @typing.overload
1299
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1188
1300
  """
1189
- Decorator prototype for all step decorators. This function gets specialized
1190
- and imported for all decorators types by _import_plugin_decorators().
1301
+ Specifies the flow(s) that this flow depends on.
1302
+
1303
+ ```
1304
+ @trigger_on_finish(flow='FooFlow')
1305
+ ```
1306
+ or
1307
+ ```
1308
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1309
+ ```
1310
+ This decorator respects the @project decorator and triggers the flow
1311
+ when upstream runs within the same namespace complete successfully
1312
+
1313
+ Additionally, you can specify project aware upstream flow dependencies
1314
+ by specifying the fully qualified project_flow_name.
1315
+ ```
1316
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1317
+ ```
1318
+ or
1319
+ ```
1320
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1321
+ ```
1322
+
1323
+ You can also specify just the project or project branch (other values will be
1324
+ inferred from the current project or project branch):
1325
+ ```
1326
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1327
+ ```
1328
+
1329
+ Note that `branch` is typically one of:
1330
+ - `prod`
1331
+ - `user.bob`
1332
+ - `test.my_experiment`
1333
+ - `prod.staging`
1334
+
1335
+
1336
+ Parameters
1337
+ ----------
1338
+ flow : Union[str, Dict[str, str]], optional, default None
1339
+ Upstream flow dependency for this flow.
1340
+ flows : List[Union[str, Dict[str, str]]], default []
1341
+ Upstream flow dependencies for this flow.
1342
+ options : Dict[str, Any], default {}
1343
+ Backend-specific configuration for tuning eventing behavior.
1191
1344
  """
1192
1345
  ...
1193
1346
 
1194
1347
  @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]]]:
1348
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1349
+ ...
1350
+
1351
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1196
1352
  """
1197
- Specifies a timeout for your step.
1353
+ Specifies the flow(s) that this flow depends on.
1198
1354
 
1199
- This decorator is useful if this step may hang indefinitely.
1355
+ ```
1356
+ @trigger_on_finish(flow='FooFlow')
1357
+ ```
1358
+ or
1359
+ ```
1360
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1361
+ ```
1362
+ This decorator respects the @project decorator and triggers the flow
1363
+ when upstream runs within the same namespace complete successfully
1200
1364
 
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.
1365
+ Additionally, you can specify project aware upstream flow dependencies
1366
+ by specifying the fully qualified project_flow_name.
1367
+ ```
1368
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1369
+ ```
1370
+ or
1371
+ ```
1372
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1373
+ ```
1204
1374
 
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.
1375
+ You can also specify just the project or project branch (other values will be
1376
+ inferred from the current project or project branch):
1377
+ ```
1378
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1379
+ ```
1380
+
1381
+ Note that `branch` is typically one of:
1382
+ - `prod`
1383
+ - `user.bob`
1384
+ - `test.my_experiment`
1385
+ - `prod.staging`
1207
1386
 
1208
1387
 
1209
1388
  Parameters
1210
1389
  ----------
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.
1390
+ flow : Union[str, Dict[str, str]], optional, default None
1391
+ Upstream flow dependency for this flow.
1392
+ flows : List[Union[str, Dict[str, str]]], default []
1393
+ Upstream flow dependencies for this flow.
1394
+ options : Dict[str, Any], default {}
1395
+ Backend-specific configuration for tuning eventing behavior.
1217
1396
  """
1218
1397
  ...
1219
1398
 
1220
1399
  @typing.overload
1221
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1400
+ 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]]:
1401
+ """
1402
+ Specifies the times when the flow should be run when running on a
1403
+ production scheduler.
1404
+
1405
+
1406
+ Parameters
1407
+ ----------
1408
+ hourly : bool, default False
1409
+ Run the workflow hourly.
1410
+ daily : bool, default True
1411
+ Run the workflow daily.
1412
+ weekly : bool, default False
1413
+ Run the workflow weekly.
1414
+ cron : str, optional, default None
1415
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1416
+ specified by this expression.
1417
+ timezone : str, optional, default None
1418
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1419
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1420
+ """
1222
1421
  ...
1223
1422
 
1224
1423
  @typing.overload
1225
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1424
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1226
1425
  ...
1227
1426
 
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):
1427
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1229
1428
  """
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.
1237
-
1238
- Note that all the values specified in parameters are added together so if you specify
1239
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1429
+ Specifies the times when the flow should be run when running on a
1430
+ production scheduler.
1240
1431
 
1241
1432
 
1242
1433
  Parameters
1243
1434
  ----------
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.
1435
+ hourly : bool, default False
1436
+ Run the workflow hourly.
1437
+ daily : bool, default True
1438
+ Run the workflow daily.
1439
+ weekly : bool, default False
1440
+ Run the workflow weekly.
1441
+ cron : str, optional, default None
1442
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1443
+ specified by this expression.
1444
+ timezone : str, optional, default None
1445
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1446
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1250
1447
  """
1251
1448
  ...
1252
1449
 
@@ -1308,194 +1505,59 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1308
1505
  @step
1309
1506
  def start(self):
1310
1507
  with open("my_file.txt", "w") as f:
1311
- f.write("Hello, World!")
1312
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1313
- self.next(self.end)
1314
-
1315
- ```
1316
-
1317
- - Accessing objects stored in external datastores after task execution.
1318
-
1319
- ```python
1320
- run = Run("CheckpointsTestsFlow/8992")
1321
- with artifact_store_from(run=run, config={
1322
- "client_params": {
1323
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1324
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1325
- },
1326
- }):
1327
- with Checkpoint() as cp:
1328
- latest = cp.list(
1329
- task=run["start"].task
1330
- )[0]
1331
- print(latest)
1332
- cp.load(
1333
- latest,
1334
- "test-checkpoints"
1335
- )
1336
-
1337
- task = Task("TorchTuneFlow/8484/train/53673")
1338
- with artifact_store_from(run=run, config={
1339
- "client_params": {
1340
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1341
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1342
- },
1343
- }):
1344
- load_model(
1345
- task.data.model_ref,
1346
- "test-models"
1347
- )
1348
- ```
1349
- Parameters:
1350
- ----------
1351
-
1352
- type: str
1353
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1354
-
1355
- config: dict or Callable
1356
- Dictionary of configuration options for the datastore. The following keys are required:
1357
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1358
- - example: 's3://bucket-name/path/to/root'
1359
- - example: 'gs://bucket-name/path/to/root'
1360
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1361
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1362
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1363
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1364
- """
1365
- ...
1366
-
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
- @typing.overload
1452
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1453
- """
1454
- Specifies the times when the flow should be run when running on a
1455
- production scheduler.
1508
+ f.write("Hello, World!")
1509
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1510
+ self.next(self.end)
1456
1511
 
1512
+ ```
1457
1513
 
1458
- Parameters
1459
- ----------
1460
- hourly : bool, default False
1461
- Run the workflow hourly.
1462
- daily : bool, default True
1463
- Run the workflow daily.
1464
- weekly : bool, default False
1465
- Run the workflow weekly.
1466
- cron : str, optional, default None
1467
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1468
- specified by this expression.
1469
- timezone : str, optional, default None
1470
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1471
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1472
- """
1473
- ...
1474
-
1475
- @typing.overload
1476
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1477
- ...
1478
-
1479
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1480
- """
1481
- Specifies the times when the flow should be run when running on a
1482
- production scheduler.
1514
+ - Accessing objects stored in external datastores after task execution.
1483
1515
 
1516
+ ```python
1517
+ run = Run("CheckpointsTestsFlow/8992")
1518
+ with artifact_store_from(run=run, config={
1519
+ "client_params": {
1520
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1521
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1522
+ },
1523
+ }):
1524
+ with Checkpoint() as cp:
1525
+ latest = cp.list(
1526
+ task=run["start"].task
1527
+ )[0]
1528
+ print(latest)
1529
+ cp.load(
1530
+ latest,
1531
+ "test-checkpoints"
1532
+ )
1484
1533
 
1485
- Parameters
1534
+ task = Task("TorchTuneFlow/8484/train/53673")
1535
+ with artifact_store_from(run=run, config={
1536
+ "client_params": {
1537
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1538
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1539
+ },
1540
+ }):
1541
+ load_model(
1542
+ task.data.model_ref,
1543
+ "test-models"
1544
+ )
1545
+ ```
1546
+ Parameters:
1486
1547
  ----------
1487
- hourly : bool, default False
1488
- Run the workflow hourly.
1489
- daily : bool, default True
1490
- 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).
1548
+
1549
+ type: str
1550
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1551
+
1552
+ config: dict or Callable
1553
+ Dictionary of configuration options for the datastore. The following keys are required:
1554
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1555
+ - example: 's3://bucket-name/path/to/root'
1556
+ - example: 'gs://bucket-name/path/to/root'
1557
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1558
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1559
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1560
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1499
1561
  """
1500
1562
  ...
1501
1563
 
@@ -1686,139 +1748,87 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1686
1748
  """
1687
1749
  ...
1688
1750
 
1689
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1751
+ @typing.overload
1752
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1690
1753
  """
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)`.
1754
+ Specifies the PyPI packages for all steps of the flow.
1695
1755
 
1756
+ Use `@pypi_base` to set common packages required by all
1757
+ steps and use `@pypi` to specify step-specific overrides.
1696
1758
 
1697
1759
  Parameters
1698
1760
  ----------
1699
- name : str
1700
- Project name. Make sure that the name is unique amongst all
1701
- projects that use the same production scheduler. The name may
1702
- contain only lowercase alphanumeric characters and underscores.
1703
-
1704
- branch : Optional[str], default None
1705
- The branch to use. If not specified, the branch is set to
1706
- `user.<username>` unless `production` is set to `True`. This can
1707
- also be set on the command line using `--branch` as a top-level option.
1708
- It is an error to specify `branch` in the decorator and on the command line.
1709
-
1710
- production : bool, default False
1711
- Whether or not the branch is the production branch. This can also be set on the
1712
- command line using `--production` as a top-level option. It is an error to specify
1713
- `production` in the decorator and on the command line.
1714
- The project branch name will be:
1715
- - if `branch` is specified:
1716
- - if `production` is True: `prod.<branch>`
1717
- - if `production` is False: `test.<branch>`
1718
- - if `branch` is not specified:
1719
- - if `production` is True: `prod`
1720
- - if `production` is False: `user.<username>`
1761
+ packages : Dict[str, str], default: {}
1762
+ Packages to use for this flow. The key is the name of the package
1763
+ and the value is the version to use.
1764
+ python : str, optional, default: None
1765
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1766
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1721
1767
  """
1722
1768
  ...
1723
1769
 
1724
1770
  @typing.overload
1725
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1771
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1772
+ ...
1773
+
1774
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1726
1775
  """
1727
- Specifies the flow(s) that this flow depends on.
1728
-
1729
- ```
1730
- @trigger_on_finish(flow='FooFlow')
1731
- ```
1732
- or
1733
- ```
1734
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1735
- ```
1736
- This decorator respects the @project decorator and triggers the flow
1737
- when upstream runs within the same namespace complete successfully
1738
-
1739
- Additionally, you can specify project aware upstream flow dependencies
1740
- by specifying the fully qualified project_flow_name.
1741
- ```
1742
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1743
- ```
1744
- or
1745
- ```
1746
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1747
- ```
1748
-
1749
- You can also specify just the project or project branch (other values will be
1750
- inferred from the current project or project branch):
1751
- ```
1752
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1753
- ```
1754
-
1755
- Note that `branch` is typically one of:
1756
- - `prod`
1757
- - `user.bob`
1758
- - `test.my_experiment`
1759
- - `prod.staging`
1776
+ Specifies the PyPI packages for all steps of the flow.
1760
1777
 
1778
+ Use `@pypi_base` to set common packages required by all
1779
+ steps and use `@pypi` to specify step-specific overrides.
1761
1780
 
1762
1781
  Parameters
1763
1782
  ----------
1764
- flow : Union[str, Dict[str, str]], optional, default None
1765
- Upstream flow dependency for this flow.
1766
- flows : List[Union[str, Dict[str, str]]], default []
1767
- Upstream flow dependencies for this flow.
1768
- options : Dict[str, Any], default {}
1769
- Backend-specific configuration for tuning eventing behavior.
1783
+ packages : Dict[str, str], default: {}
1784
+ Packages to use for this flow. The key is the name of the package
1785
+ and the value is the version to use.
1786
+ python : str, optional, default: None
1787
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1788
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1770
1789
  """
1771
1790
  ...
1772
1791
 
1773
- @typing.overload
1774
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1775
- ...
1776
-
1777
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1792
+ 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]]:
1778
1793
  """
1779
- Specifies the flow(s) that this flow depends on.
1780
-
1781
- ```
1782
- @trigger_on_finish(flow='FooFlow')
1783
- ```
1784
- or
1785
- ```
1786
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1787
- ```
1788
- This decorator respects the @project decorator and triggers the flow
1789
- when upstream runs within the same namespace complete successfully
1790
-
1791
- Additionally, you can specify project aware upstream flow dependencies
1792
- by specifying the fully qualified project_flow_name.
1793
- ```
1794
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1795
- ```
1796
- or
1797
- ```
1798
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1799
- ```
1800
-
1801
- You can also specify just the project or project branch (other values will be
1802
- inferred from the current project or project branch):
1803
- ```
1804
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1805
- ```
1806
-
1807
- Note that `branch` is typically one of:
1808
- - `prod`
1809
- - `user.bob`
1810
- - `test.my_experiment`
1811
- - `prod.staging`
1794
+ 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)
1795
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1796
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1797
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1798
+ starts only after all sensors finish.
1812
1799
 
1813
1800
 
1814
1801
  Parameters
1815
1802
  ----------
1816
- flow : Union[str, Dict[str, str]], optional, default None
1817
- Upstream flow dependency for this flow.
1818
- flows : List[Union[str, Dict[str, str]]], default []
1819
- Upstream flow dependencies for this flow.
1820
- options : Dict[str, Any], default {}
1821
- Backend-specific configuration for tuning eventing behavior.
1803
+ timeout : int
1804
+ Time, in seconds before the task times out and fails. (Default: 3600)
1805
+ poke_interval : int
1806
+ Time in seconds that the job should wait in between each try. (Default: 60)
1807
+ mode : str
1808
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1809
+ exponential_backoff : bool
1810
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1811
+ pool : str
1812
+ the slot pool this task should run in,
1813
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1814
+ soft_fail : bool
1815
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1816
+ name : str
1817
+ Name of the sensor on Airflow
1818
+ description : str
1819
+ Description of sensor in the Airflow UI
1820
+ bucket_key : Union[str, List[str]]
1821
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1822
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1823
+ bucket_name : str
1824
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1825
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1826
+ wildcard_match : bool
1827
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1828
+ aws_conn_id : str
1829
+ a reference to the s3 connection on Airflow. (Default: None)
1830
+ verify : bool
1831
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1822
1832
  """
1823
1833
  ...
1824
1834