ob-metaflow-stubs 6.0.6.0__py2.py3-none-any.whl → 6.0.6.2__py2.py3-none-any.whl

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