ob-metaflow-stubs 6.0.7.0__py2.py3-none-any.whl → 6.0.7.1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of ob-metaflow-stubs might be problematic. Click here for more details.

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