ob-metaflow-stubs 6.0.5.2__py2.py3-none-any.whl → 6.0.6.1__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (261) hide show
  1. metaflow-stubs/__init__.pyi +1051 -1051
  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 +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  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 +36 -36
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +3 -3
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +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 +3 -3
  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 +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +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 +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +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 +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +3 -3
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +2 -2
  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 +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +4 -4
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +2 -2
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +3 -3
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +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 +3 -3
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +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 +2 -2
  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 +7 -7
  119. metaflow-stubs/packaging_sys/backend.pyi +2 -2
  120. metaflow-stubs/packaging_sys/distribution_support.pyi +4 -4
  121. metaflow-stubs/packaging_sys/tar_backend.pyi +4 -4
  122. metaflow-stubs/packaging_sys/utils.pyi +2 -2
  123. metaflow-stubs/packaging_sys/v1.pyi +2 -2
  124. metaflow-stubs/parameters.pyi +3 -3
  125. metaflow-stubs/plugins/__init__.pyi +10 -10
  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 +3 -3
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  139. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  140. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  141. metaflow-stubs/plugins/aws/__init__.pyi +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 +3 -3
  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 +3 -3
  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 +3 -3
  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 +3 -3
  212. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  213. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  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 +2 -2
  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 +32 -32
  237. metaflow-stubs/runner/deployer_impl.pyi +2 -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 +2 -2
  243. metaflow-stubs/system/__init__.pyi +2 -2
  244. metaflow-stubs/system/system_logger.pyi +2 -2
  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 +2 -2
  250. metaflow-stubs/user_configs/config_parameters.pyi +4 -4
  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 +4 -4
  254. metaflow-stubs/user_decorators/mutable_step.pyi +3 -3
  255. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  256. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  257. {ob_metaflow_stubs-6.0.5.2.dist-info → ob_metaflow_stubs-6.0.6.1.dist-info}/METADATA +1 -1
  258. ob_metaflow_stubs-6.0.6.1.dist-info/RECORD +261 -0
  259. ob_metaflow_stubs-6.0.5.2.dist-info/RECORD +0 -261
  260. {ob_metaflow_stubs-6.0.5.2.dist-info → ob_metaflow_stubs-6.0.6.1.dist-info}/WHEEL +0 -0
  261. {ob_metaflow_stubs-6.0.5.2.dist-info → ob_metaflow_stubs-6.0.6.1.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.16.8.1+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-08-04T19:06:54.653206 #
3
+ # MF version: 2.17.0.1+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-08-06T15:19:48.623789 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -39,8 +39,8 @@ from .user_decorators.user_step_decorator import UserStepDecorator as UserStepDe
39
39
  from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
- from . import cards as cards
43
42
  from . import tuple_util as tuple_util
43
+ from . import cards as cards
44
44
  from . import metaflow_git as metaflow_git
45
45
  from . import events as events
46
46
  from . import runner as runner
@@ -48,9 +48,9 @@ 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 conda_environment_yml_parser as conda_environment_yml_parser
53
51
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
52
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_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
@@ -167,303 +167,195 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
167
167
  """
168
168
  ...
169
169
 
170
- 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]]]:
171
- """
172
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
173
-
174
- User code call
175
- --------------
176
- @ollama(
177
- models=[...],
178
- ...
179
- )
180
-
181
- Valid backend options
182
- ---------------------
183
- - 'local': Run as a separate process on the local task machine.
184
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
185
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
186
-
187
- Valid model options
188
- -------------------
189
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
190
-
191
-
192
- Parameters
193
- ----------
194
- models: list[str]
195
- List of Ollama containers running models in sidecars.
196
- backend: str
197
- Determines where and how to run the Ollama process.
198
- force_pull: bool
199
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
200
- cache_update_policy: str
201
- Cache update policy: "auto", "force", or "never".
202
- force_cache_update: bool
203
- Simple override for "force" cache update policy.
204
- debug: bool
205
- Whether to turn on verbose debugging logs.
206
- circuit_breaker_config: dict
207
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
208
- timeout_config: dict
209
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
210
- """
211
- ...
212
-
213
- 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]]]:
214
- """
215
- S3 Proxy decorator for routing S3 requests through a local proxy service.
216
-
217
-
218
- Parameters
219
- ----------
220
- integration_name : str, optional
221
- Name of the S3 proxy integration. If not specified, will use the only
222
- available S3 proxy integration in the namespace (fails if multiple exist).
223
- write_mode : str, optional
224
- The desired behavior during write operations to target (origin) S3 bucket.
225
- allowed options are:
226
- "origin-and-cache" -> write to both the target S3 bucket and local object
227
- storage
228
- "origin" -> only write to the target S3 bucket
229
- "cache" -> only write to the object storage service used for caching
230
- debug : bool, optional
231
- Enable debug logging for proxy operations.
232
- """
233
- ...
234
-
235
170
  @typing.overload
236
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
171
+ def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
237
172
  """
238
- Enables checkpointing for a step.
173
+ Enables loading / saving of models within a step.
239
174
 
240
175
  > Examples
241
-
242
- - Saving Checkpoints
243
-
176
+ - Saving Models
244
177
  ```python
245
- @checkpoint
178
+ @model
246
179
  @step
247
180
  def train(self):
248
- model = create_model(self.parameters, checkpoint_path = None)
249
- for i in range(self.epochs):
250
- # some training logic
251
- loss = model.train(self.dataset)
252
- if i % 10 == 0:
253
- model.save(
254
- current.checkpoint.directory,
255
- )
256
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
257
- # and returns a reference dictionary to the checkpoint saved in the datastore
258
- self.latest_checkpoint = current.checkpoint.save(
259
- name="epoch_checkpoint",
260
- metadata={
261
- "epoch": i,
262
- "loss": loss,
263
- }
264
- )
265
- ```
181
+ # current.model.save returns a dictionary reference to the model saved
182
+ self.my_model = current.model.save(
183
+ path_to_my_model,
184
+ label="my_model",
185
+ metadata={
186
+ "epochs": 10,
187
+ "batch-size": 32,
188
+ "learning-rate": 0.001,
189
+ }
190
+ )
191
+ self.next(self.test)
266
192
 
267
- - Using Loaded Checkpoints
193
+ @model(load="my_model")
194
+ @step
195
+ def test(self):
196
+ # `current.model.loaded` returns a dictionary of the loaded models
197
+ # where the key is the name of the artifact and the value is the path to the model
198
+ print(os.listdir(current.model.loaded["my_model"]))
199
+ self.next(self.end)
200
+ ```
268
201
 
202
+ - Loading models
269
203
  ```python
270
- @retry(times=3)
271
- @checkpoint
272
204
  @step
273
205
  def train(self):
274
- # Assume that the task has restarted and the previous attempt of the task
275
- # saved a checkpoint
276
- checkpoint_path = None
277
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
278
- print("Loaded checkpoint from the previous attempt")
279
- checkpoint_path = current.checkpoint.directory
280
-
281
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
282
- for i in range(self.epochs):
283
- ...
206
+ # current.model.load returns the path to the model loaded
207
+ checkpoint_path = current.model.load(
208
+ self.checkpoint_key,
209
+ )
210
+ model_path = current.model.load(
211
+ self.model,
212
+ )
213
+ self.next(self.test)
284
214
  ```
285
215
 
286
216
 
287
217
  Parameters
288
218
  ----------
289
- load_policy : str, default: "fresh"
290
- The policy for loading the checkpoint. The following policies are supported:
291
- - "eager": Loads the the latest available checkpoint within the namespace.
292
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
293
- will be loaded at the start of the task.
294
- - "none": Do not load any checkpoint
295
- - "fresh": Loads the lastest checkpoint created within the running Task.
296
- This mode helps loading checkpoints across various retry attempts of the same task.
297
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
298
- created within the task will be loaded when the task is retries execution on failure.
219
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
220
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
221
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
222
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
223
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
224
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
299
225
 
300
226
  temp_dir_root : str, default: None
301
- The root directory under which `current.checkpoint.directory` will be created.
227
+ The root directory under which `current.model.loaded` will store loaded models
302
228
  """
303
229
  ...
304
230
 
305
231
  @typing.overload
306
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
232
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
307
233
  ...
308
234
 
309
235
  @typing.overload
310
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
236
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
311
237
  ...
312
238
 
313
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
239
+ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
314
240
  """
315
- Enables checkpointing for a step.
241
+ Enables loading / saving of models within a step.
316
242
 
317
243
  > Examples
318
-
319
- - Saving Checkpoints
320
-
244
+ - Saving Models
321
245
  ```python
322
- @checkpoint
246
+ @model
323
247
  @step
324
248
  def train(self):
325
- model = create_model(self.parameters, checkpoint_path = None)
326
- for i in range(self.epochs):
327
- # some training logic
328
- loss = model.train(self.dataset)
329
- if i % 10 == 0:
330
- model.save(
331
- current.checkpoint.directory,
332
- )
333
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
334
- # and returns a reference dictionary to the checkpoint saved in the datastore
335
- self.latest_checkpoint = current.checkpoint.save(
336
- name="epoch_checkpoint",
337
- metadata={
338
- "epoch": i,
339
- "loss": loss,
340
- }
341
- )
342
- ```
249
+ # current.model.save returns a dictionary reference to the model saved
250
+ self.my_model = current.model.save(
251
+ path_to_my_model,
252
+ label="my_model",
253
+ metadata={
254
+ "epochs": 10,
255
+ "batch-size": 32,
256
+ "learning-rate": 0.001,
257
+ }
258
+ )
259
+ self.next(self.test)
343
260
 
344
- - Using Loaded Checkpoints
261
+ @model(load="my_model")
262
+ @step
263
+ def test(self):
264
+ # `current.model.loaded` returns a dictionary of the loaded models
265
+ # where the key is the name of the artifact and the value is the path to the model
266
+ print(os.listdir(current.model.loaded["my_model"]))
267
+ self.next(self.end)
268
+ ```
345
269
 
270
+ - Loading models
346
271
  ```python
347
- @retry(times=3)
348
- @checkpoint
349
272
  @step
350
273
  def train(self):
351
- # Assume that the task has restarted and the previous attempt of the task
352
- # saved a checkpoint
353
- checkpoint_path = None
354
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
355
- print("Loaded checkpoint from the previous attempt")
356
- checkpoint_path = current.checkpoint.directory
357
-
358
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
359
- for i in range(self.epochs):
360
- ...
274
+ # current.model.load returns the path to the model loaded
275
+ checkpoint_path = current.model.load(
276
+ self.checkpoint_key,
277
+ )
278
+ model_path = current.model.load(
279
+ self.model,
280
+ )
281
+ self.next(self.test)
361
282
  ```
362
283
 
363
284
 
364
285
  Parameters
365
286
  ----------
366
- load_policy : str, default: "fresh"
367
- The policy for loading the checkpoint. The following policies are supported:
368
- - "eager": Loads the the latest available checkpoint within the namespace.
369
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
370
- will be loaded at the start of the task.
371
- - "none": Do not load any checkpoint
372
- - "fresh": Loads the lastest checkpoint created within the running Task.
373
- This mode helps loading checkpoints across various retry attempts of the same task.
374
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
375
- created within the task will be loaded when the task is retries execution on failure.
287
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
288
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
289
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
290
+ If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
291
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
292
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
376
293
 
377
294
  temp_dir_root : str, default: None
378
- The root directory under which `current.checkpoint.directory` will be created.
295
+ The root directory under which `current.model.loaded` will store loaded models
379
296
  """
380
297
  ...
381
298
 
382
299
  @typing.overload
383
- 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]]]:
300
+ 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]]]:
384
301
  """
385
- Specifies environment variables to be set prior to the execution of a step.
386
-
302
+ Specifies the Conda environment for the step.
387
303
 
388
- Parameters
389
- ----------
390
- vars : Dict[str, str], default {}
391
- Dictionary of environment variables to set.
392
- """
393
- ...
394
-
395
- @typing.overload
396
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
397
- ...
398
-
399
- @typing.overload
400
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
401
- ...
402
-
403
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
404
- """
405
- Specifies environment variables to be set prior to the execution of a step.
406
-
407
-
408
- Parameters
409
- ----------
410
- vars : Dict[str, str], default {}
411
- Dictionary of environment variables to set.
412
- """
413
- ...
414
-
415
- @typing.overload
416
- 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]]]:
417
- """
418
- Specifies that the step will success under all circumstances.
419
-
420
- The decorator will create an optional artifact, specified by `var`, which
421
- contains the exception raised. You can use it to detect the presence
422
- of errors, indicating that all happy-path artifacts produced by the step
423
- are missing.
304
+ Information in this decorator will augment any
305
+ attributes set in the `@conda_base` flow-level decorator. Hence,
306
+ you can use `@conda_base` to set packages required by all
307
+ steps and use `@conda` to specify step-specific overrides.
424
308
 
425
309
 
426
310
  Parameters
427
311
  ----------
428
- var : str, optional, default None
429
- Name of the artifact in which to store the caught exception.
430
- If not specified, the exception is not stored.
431
- print_exception : bool, default True
432
- Determines whether or not the exception is printed to
433
- stdout when caught.
312
+ packages : Dict[str, str], default {}
313
+ Packages to use for this step. The key is the name of the package
314
+ and the value is the version to use.
315
+ libraries : Dict[str, str], default {}
316
+ Supported for backward compatibility. When used with packages, packages will take precedence.
317
+ python : str, optional, default None
318
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
319
+ that the version used will correspond to the version of the Python interpreter used to start the run.
320
+ disabled : bool, default False
321
+ If set to True, disables @conda.
434
322
  """
435
323
  ...
436
324
 
437
325
  @typing.overload
438
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
326
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
439
327
  ...
440
328
 
441
329
  @typing.overload
442
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
330
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
443
331
  ...
444
332
 
445
- 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):
333
+ 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):
446
334
  """
447
- Specifies that the step will success under all circumstances.
335
+ Specifies the Conda environment for the step.
448
336
 
449
- The decorator will create an optional artifact, specified by `var`, which
450
- contains the exception raised. You can use it to detect the presence
451
- of errors, indicating that all happy-path artifacts produced by the step
452
- are missing.
337
+ Information in this decorator will augment any
338
+ attributes set in the `@conda_base` flow-level decorator. Hence,
339
+ you can use `@conda_base` to set packages required by all
340
+ steps and use `@conda` to specify step-specific overrides.
453
341
 
454
342
 
455
343
  Parameters
456
344
  ----------
457
- var : str, optional, default None
458
- Name of the artifact in which to store the caught exception.
459
- If not specified, the exception is not stored.
460
- print_exception : bool, default True
461
- Determines whether or not the exception is printed to
462
- stdout when caught.
345
+ packages : Dict[str, str], default {}
346
+ Packages to use for this step. The key is the name of the package
347
+ and the value is the version to use.
348
+ libraries : Dict[str, str], default {}
349
+ Supported for backward compatibility. When used with packages, packages will take precedence.
350
+ python : str, optional, default None
351
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
352
+ that the version used will correspond to the version of the Python interpreter used to start the run.
353
+ disabled : bool, default False
354
+ If set to True, disables @conda.
463
355
  """
464
356
  ...
465
357
 
466
- 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]]]:
358
+ 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]]]:
467
359
  """
468
360
  Specifies that this step should execute on DGX cloud.
469
361
 
@@ -474,88 +366,39 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
474
366
  Number of GPUs to use.
475
367
  gpu_type : str
476
368
  Type of Nvidia GPU to use.
477
- queue_timeout : int
478
- Time to keep the job in NVCF's queue.
479
- """
480
- ...
481
-
482
- @typing.overload
483
- def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
484
- """
485
- Creates a human-readable report, a Metaflow Card, after this step completes.
486
-
487
- Note that you may add multiple `@card` decorators in a step with different parameters.
488
-
489
-
490
- Parameters
491
- ----------
492
- type : str, default 'default'
493
- Card type.
494
- id : str, optional, default None
495
- If multiple cards are present, use this id to identify this card.
496
- options : Dict[str, Any], default {}
497
- Options passed to the card. The contents depend on the card type.
498
- timeout : int, default 45
499
- Interrupt reporting if it takes more than this many seconds.
500
369
  """
501
370
  ...
502
371
 
503
372
  @typing.overload
504
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
505
- ...
506
-
507
- @typing.overload
508
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
509
- ...
510
-
511
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
373
+ 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]]]:
512
374
  """
513
- Creates a human-readable report, a Metaflow Card, after this step completes.
514
-
515
- Note that you may add multiple `@card` decorators in a step with different parameters.
375
+ Specifies environment variables to be set prior to the execution of a step.
516
376
 
517
377
 
518
378
  Parameters
519
379
  ----------
520
- type : str, default 'default'
521
- Card type.
522
- id : str, optional, default None
523
- If multiple cards are present, use this id to identify this card.
524
- options : Dict[str, Any], default {}
525
- Options passed to the card. The contents depend on the card type.
526
- timeout : int, default 45
527
- Interrupt reporting if it takes more than this many seconds.
380
+ vars : Dict[str, str], default {}
381
+ Dictionary of environment variables to set.
528
382
  """
529
383
  ...
530
384
 
531
385
  @typing.overload
532
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
533
- """
534
- Internal decorator to support Fast bakery
535
- """
386
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
536
387
  ...
537
388
 
538
389
  @typing.overload
539
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
540
- ...
541
-
542
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
543
- """
544
- Internal decorator to support Fast bakery
545
- """
390
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
546
391
  ...
547
392
 
548
- 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]]]:
393
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
549
394
  """
550
- Specifies that this step should execute on DGX cloud.
395
+ Specifies environment variables to be set prior to the execution of a step.
551
396
 
552
397
 
553
398
  Parameters
554
399
  ----------
555
- gpu : int
556
- Number of GPUs to use.
557
- gpu_type : str
558
- Type of Nvidia GPU to use.
400
+ vars : Dict[str, str], default {}
401
+ Dictionary of environment variables to set.
559
402
  """
560
403
  ...
561
404
 
@@ -611,149 +454,168 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
611
454
  ...
612
455
 
613
456
  @typing.overload
614
- 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]]]:
457
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
615
458
  """
616
- Specifies secrets to be retrieved and injected as environment variables prior to
617
- the execution of a step.
459
+ Enables checkpointing for a step.
618
460
 
461
+ > Examples
619
462
 
620
- Parameters
621
- ----------
622
- sources : List[Union[str, Dict[str, Any]]], default: []
623
- List of secret specs, defining how the secrets are to be retrieved
624
- role : str, optional, default: None
625
- Role to use for fetching secrets
626
- """
627
- ...
628
-
629
- @typing.overload
630
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
631
- ...
632
-
633
- @typing.overload
634
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
635
- ...
636
-
637
- 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):
638
- """
639
- Specifies secrets to be retrieved and injected as environment variables prior to
640
- the execution of a step.
463
+ - Saving Checkpoints
464
+
465
+ ```python
466
+ @checkpoint
467
+ @step
468
+ def train(self):
469
+ model = create_model(self.parameters, checkpoint_path = None)
470
+ for i in range(self.epochs):
471
+ # some training logic
472
+ loss = model.train(self.dataset)
473
+ if i % 10 == 0:
474
+ model.save(
475
+ current.checkpoint.directory,
476
+ )
477
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
478
+ # and returns a reference dictionary to the checkpoint saved in the datastore
479
+ self.latest_checkpoint = current.checkpoint.save(
480
+ name="epoch_checkpoint",
481
+ metadata={
482
+ "epoch": i,
483
+ "loss": loss,
484
+ }
485
+ )
486
+ ```
487
+
488
+ - Using Loaded Checkpoints
489
+
490
+ ```python
491
+ @retry(times=3)
492
+ @checkpoint
493
+ @step
494
+ def train(self):
495
+ # Assume that the task has restarted and the previous attempt of the task
496
+ # saved a checkpoint
497
+ checkpoint_path = None
498
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
499
+ print("Loaded checkpoint from the previous attempt")
500
+ checkpoint_path = current.checkpoint.directory
501
+
502
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
503
+ for i in range(self.epochs):
504
+ ...
505
+ ```
641
506
 
642
507
 
643
508
  Parameters
644
509
  ----------
645
- sources : List[Union[str, Dict[str, Any]]], default: []
646
- List of secret specs, defining how the secrets are to be retrieved
647
- role : str, optional, default: None
648
- Role to use for fetching secrets
510
+ load_policy : str, default: "fresh"
511
+ The policy for loading the checkpoint. The following policies are supported:
512
+ - "eager": Loads the the latest available checkpoint within the namespace.
513
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
514
+ will be loaded at the start of the task.
515
+ - "none": Do not load any checkpoint
516
+ - "fresh": Loads the lastest checkpoint created within the running Task.
517
+ This mode helps loading checkpoints across various retry attempts of the same task.
518
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
519
+ created within the task will be loaded when the task is retries execution on failure.
520
+
521
+ temp_dir_root : str, default: None
522
+ The root directory under which `current.checkpoint.directory` will be created.
649
523
  """
650
524
  ...
651
525
 
652
526
  @typing.overload
653
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
654
- """
655
- A simple decorator that demonstrates using CardDecoratorInjector
656
- to inject a card and render simple markdown content.
657
- """
527
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
658
528
  ...
659
529
 
660
530
  @typing.overload
661
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
531
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
662
532
  ...
663
533
 
664
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
534
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
665
535
  """
666
- A simple decorator that demonstrates using CardDecoratorInjector
667
- to inject a card and render simple markdown content.
668
- """
669
- ...
670
-
671
- 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]]]:
672
- """
673
- Specifies that this step should execute on Kubernetes.
536
+ Enables checkpointing for a step.
537
+
538
+ > Examples
539
+
540
+ - Saving Checkpoints
541
+
542
+ ```python
543
+ @checkpoint
544
+ @step
545
+ def train(self):
546
+ model = create_model(self.parameters, checkpoint_path = None)
547
+ for i in range(self.epochs):
548
+ # some training logic
549
+ loss = model.train(self.dataset)
550
+ if i % 10 == 0:
551
+ model.save(
552
+ current.checkpoint.directory,
553
+ )
554
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
555
+ # and returns a reference dictionary to the checkpoint saved in the datastore
556
+ self.latest_checkpoint = current.checkpoint.save(
557
+ name="epoch_checkpoint",
558
+ metadata={
559
+ "epoch": i,
560
+ "loss": loss,
561
+ }
562
+ )
563
+ ```
564
+
565
+ - Using Loaded Checkpoints
566
+
567
+ ```python
568
+ @retry(times=3)
569
+ @checkpoint
570
+ @step
571
+ def train(self):
572
+ # Assume that the task has restarted and the previous attempt of the task
573
+ # saved a checkpoint
574
+ checkpoint_path = None
575
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
576
+ print("Loaded checkpoint from the previous attempt")
577
+ checkpoint_path = current.checkpoint.directory
578
+
579
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
580
+ for i in range(self.epochs):
581
+ ...
582
+ ```
674
583
 
675
584
 
676
585
  Parameters
677
586
  ----------
678
- cpu : int, default 1
679
- Number of CPUs required for this step. If `@resources` is
680
- also present, the maximum value from all decorators is used.
681
- memory : int, default 4096
682
- Memory size (in MB) required for this step. If
683
- `@resources` is also present, the maximum value from all decorators is
684
- used.
685
- disk : int, default 10240
686
- Disk size (in MB) required for this step. If
687
- `@resources` is also present, the maximum value from all decorators is
688
- used.
689
- image : str, optional, default None
690
- Docker image to use when launching on Kubernetes. If not specified, and
691
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
692
- not, a default Docker image mapping to the current version of Python is used.
693
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
694
- If given, the imagePullPolicy to be applied to the Docker image of the step.
695
- image_pull_secrets: List[str], default []
696
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
697
- Kubernetes image pull secrets to use when pulling container images
698
- in Kubernetes.
699
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
700
- Kubernetes service account to use when launching pod in Kubernetes.
701
- secrets : List[str], optional, default None
702
- Kubernetes secrets to use when launching pod in Kubernetes. These
703
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
704
- in Metaflow configuration.
705
- node_selector: Union[Dict[str,str], str], optional, default None
706
- Kubernetes node selector(s) to apply to the pod running the task.
707
- Can be passed in as a comma separated string of values e.g.
708
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
709
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
710
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
711
- Kubernetes namespace to use when launching pod in Kubernetes.
712
- gpu : int, optional, default None
713
- Number of GPUs required for this step. A value of zero implies that
714
- the scheduled node should not have GPUs.
715
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
716
- The vendor of the GPUs to be used for this step.
717
- tolerations : List[Dict[str,str]], default []
718
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
719
- Kubernetes tolerations to use when launching pod in Kubernetes.
720
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
721
- Kubernetes labels to use when launching pod in Kubernetes.
722
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
723
- Kubernetes annotations to use when launching pod in Kubernetes.
724
- use_tmpfs : bool, default False
725
- This enables an explicit tmpfs mount for this step.
726
- tmpfs_tempdir : bool, default True
727
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
728
- tmpfs_size : int, optional, default: None
729
- The value for the size (in MiB) of the tmpfs mount for this step.
730
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
731
- memory allocated for this step.
732
- tmpfs_path : str, optional, default /metaflow_temp
733
- Path to tmpfs mount for this step.
734
- persistent_volume_claims : Dict[str, str], optional, default None
735
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
736
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
737
- shared_memory: int, optional
738
- Shared memory size (in MiB) required for this step
739
- port: int, optional
740
- Port number to specify in the Kubernetes job object
741
- compute_pool : str, optional, default None
742
- Compute pool to be used for for this step.
743
- If not specified, any accessible compute pool within the perimeter is used.
744
- hostname_resolution_timeout: int, default 10 * 60
745
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
746
- Only applicable when @parallel is used.
747
- qos: str, default: Burstable
748
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
587
+ load_policy : str, default: "fresh"
588
+ The policy for loading the checkpoint. The following policies are supported:
589
+ - "eager": Loads the the latest available checkpoint within the namespace.
590
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
591
+ will be loaded at the start of the task.
592
+ - "none": Do not load any checkpoint
593
+ - "fresh": Loads the lastest checkpoint created within the running Task.
594
+ This mode helps loading checkpoints across various retry attempts of the same task.
595
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
596
+ created within the task will be loaded when the task is retries execution on failure.
749
597
 
750
- security_context: Dict[str, Any], optional, default None
751
- Container security context. Applies to the task container. Allows the following keys:
752
- - privileged: bool, optional, default None
753
- - allow_privilege_escalation: bool, optional, default None
754
- - run_as_user: int, optional, default None
755
- - run_as_group: int, optional, default None
756
- - run_as_non_root: bool, optional, default None
598
+ temp_dir_root : str, default: None
599
+ The root directory under which `current.checkpoint.directory` will be created.
600
+ """
601
+ ...
602
+
603
+ @typing.overload
604
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
605
+ """
606
+ A simple decorator that demonstrates using CardDecoratorInjector
607
+ to inject a card and render simple markdown content.
608
+ """
609
+ ...
610
+
611
+ @typing.overload
612
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
613
+ ...
614
+
615
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
616
+ """
617
+ A simple decorator that demonstrates using CardDecoratorInjector
618
+ to inject a card and render simple markdown content.
757
619
  """
758
620
  ...
759
621
 
@@ -778,6 +640,81 @@ def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag],
778
640
  """
779
641
  ...
780
642
 
643
+ @typing.overload
644
+ 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]]]:
645
+ """
646
+ Specifies a timeout for your step.
647
+
648
+ This decorator is useful if this step may hang indefinitely.
649
+
650
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
651
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
652
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
653
+
654
+ Note that all the values specified in parameters are added together so if you specify
655
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
656
+
657
+
658
+ Parameters
659
+ ----------
660
+ seconds : int, default 0
661
+ Number of seconds to wait prior to timing out.
662
+ minutes : int, default 0
663
+ Number of minutes to wait prior to timing out.
664
+ hours : int, default 0
665
+ Number of hours to wait prior to timing out.
666
+ """
667
+ ...
668
+
669
+ @typing.overload
670
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
671
+ ...
672
+
673
+ @typing.overload
674
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
675
+ ...
676
+
677
+ 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):
678
+ """
679
+ Specifies a timeout for your step.
680
+
681
+ This decorator is useful if this step may hang indefinitely.
682
+
683
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
684
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
685
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
686
+
687
+ Note that all the values specified in parameters are added together so if you specify
688
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
689
+
690
+
691
+ Parameters
692
+ ----------
693
+ seconds : int, default 0
694
+ Number of seconds to wait prior to timing out.
695
+ minutes : int, default 0
696
+ Number of minutes to wait prior to timing out.
697
+ hours : int, default 0
698
+ Number of hours to wait prior to timing out.
699
+ """
700
+ ...
701
+
702
+ 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]]]:
703
+ """
704
+ Specifies that this step should execute on DGX cloud.
705
+
706
+
707
+ Parameters
708
+ ----------
709
+ gpu : int
710
+ Number of GPUs to use.
711
+ gpu_type : str
712
+ Type of Nvidia GPU to use.
713
+ queue_timeout : int
714
+ Time to keep the job in NVCF's queue.
715
+ """
716
+ ...
717
+
781
718
  @typing.overload
782
719
  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]]]:
783
720
  """
@@ -833,24 +770,114 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
833
770
  """
834
771
  ...
835
772
 
836
- 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]]]:
773
+ @typing.overload
774
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
837
775
  """
838
- Decorator that helps cache, version and store models/datasets from huggingface hub.
839
-
840
- > Examples
841
-
842
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
843
- ```python
844
- @huggingface_hub
845
- @step
846
- def pull_model_from_huggingface(self):
847
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
848
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
849
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
850
- # value of the function is a reference to the model in the backend storage.
851
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
852
-
853
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
776
+ Decorator prototype for all step decorators. This function gets specialized
777
+ and imported for all decorators types by _import_plugin_decorators().
778
+ """
779
+ ...
780
+
781
+ @typing.overload
782
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
783
+ ...
784
+
785
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
786
+ """
787
+ Decorator prototype for all step decorators. This function gets specialized
788
+ and imported for all decorators types by _import_plugin_decorators().
789
+ """
790
+ ...
791
+
792
+ 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]]]:
793
+ """
794
+ S3 Proxy decorator for routing S3 requests through a local proxy service.
795
+
796
+
797
+ Parameters
798
+ ----------
799
+ integration_name : str, optional
800
+ Name of the S3 proxy integration. If not specified, will use the only
801
+ available S3 proxy integration in the namespace (fails if multiple exist).
802
+ write_mode : str, optional
803
+ The desired behavior during write operations to target (origin) S3 bucket.
804
+ allowed options are:
805
+ "origin-and-cache" -> write to both the target S3 bucket and local object
806
+ storage
807
+ "origin" -> only write to the target S3 bucket
808
+ "cache" -> only write to the object storage service used for caching
809
+ debug : bool, optional
810
+ Enable debug logging for proxy operations.
811
+ """
812
+ ...
813
+
814
+ @typing.overload
815
+ def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
816
+ """
817
+ Creates a human-readable report, a Metaflow Card, after this step completes.
818
+
819
+ Note that you may add multiple `@card` decorators in a step with different parameters.
820
+
821
+
822
+ Parameters
823
+ ----------
824
+ type : str, default 'default'
825
+ Card type.
826
+ id : str, optional, default None
827
+ If multiple cards are present, use this id to identify this card.
828
+ options : Dict[str, Any], default {}
829
+ Options passed to the card. The contents depend on the card type.
830
+ timeout : int, default 45
831
+ Interrupt reporting if it takes more than this many seconds.
832
+ """
833
+ ...
834
+
835
+ @typing.overload
836
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
837
+ ...
838
+
839
+ @typing.overload
840
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
841
+ ...
842
+
843
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
844
+ """
845
+ Creates a human-readable report, a Metaflow Card, after this step completes.
846
+
847
+ Note that you may add multiple `@card` decorators in a step with different parameters.
848
+
849
+
850
+ Parameters
851
+ ----------
852
+ type : str, default 'default'
853
+ Card type.
854
+ id : str, optional, default None
855
+ If multiple cards are present, use this id to identify this card.
856
+ options : Dict[str, Any], default {}
857
+ Options passed to the card. The contents depend on the card type.
858
+ timeout : int, default 45
859
+ Interrupt reporting if it takes more than this many seconds.
860
+ """
861
+ ...
862
+
863
+ 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]]]:
864
+ """
865
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
866
+
867
+ > Examples
868
+
869
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
870
+ ```python
871
+ @huggingface_hub
872
+ @step
873
+ def pull_model_from_huggingface(self):
874
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
875
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
876
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
877
+ # value of the function is a reference to the model in the backend storage.
878
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
879
+
880
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
854
881
  self.llama_model = current.huggingface_hub.snapshot_download(
855
882
  repo_id=self.model_id,
856
883
  allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
@@ -913,489 +940,454 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.
913
940
  """
914
941
  ...
915
942
 
916
- @typing.overload
917
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
918
- """
919
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
920
- It exists to make it easier for users to know that this decorator should only be used with
921
- a Neo Cloud like CoreWeave.
922
- """
923
- ...
924
-
925
- @typing.overload
926
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
927
- ...
928
-
929
- def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
930
- """
931
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
932
- It exists to make it easier for users to know that this decorator should only be used with
933
- a Neo Cloud like CoreWeave.
934
- """
935
- ...
936
-
937
- @typing.overload
938
- 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]]]:
943
+ 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]]]:
939
944
  """
940
- Specifies the resources needed when executing this step.
945
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
941
946
 
942
- Use `@resources` to specify the resource requirements
943
- independently of the specific compute layer (`@batch`, `@kubernetes`).
947
+ User code call
948
+ --------------
949
+ @vllm(
950
+ model="...",
951
+ ...
952
+ )
944
953
 
945
- You can choose the compute layer on the command line by executing e.g.
946
- ```
947
- python myflow.py run --with batch
948
- ```
949
- or
950
- ```
951
- python myflow.py run --with kubernetes
952
- ```
953
- which executes the flow on the desired system using the
954
- requirements specified in `@resources`.
954
+ Valid backend options
955
+ ---------------------
956
+ - 'local': Run as a separate process on the local task machine.
957
+
958
+ Valid model options
959
+ -------------------
960
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
961
+
962
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
963
+ If you need multiple models, you must create multiple @vllm decorators.
955
964
 
956
965
 
957
966
  Parameters
958
967
  ----------
959
- cpu : int, default 1
960
- Number of CPUs required for this step.
961
- gpu : int, optional, default None
962
- Number of GPUs required for this step.
963
- disk : int, optional, default None
964
- Disk size (in MB) required for this step. Only applies on Kubernetes.
965
- memory : int, default 4096
966
- Memory size (in MB) required for this step.
967
- shared_memory : int, optional, default None
968
- The value for the size (in MiB) of the /dev/shm volume for this step.
969
- This parameter maps to the `--shm-size` option in Docker.
968
+ model: str
969
+ HuggingFace model identifier to be served by vLLM.
970
+ backend: str
971
+ Determines where and how to run the vLLM process.
972
+ openai_api_server: bool
973
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
974
+ Default is False (uses native engine).
975
+ Set to True for backward compatibility with existing code.
976
+ debug: bool
977
+ Whether to turn on verbose debugging logs.
978
+ card_refresh_interval: int
979
+ Interval in seconds for refreshing the vLLM status card.
980
+ Only used when openai_api_server=True.
981
+ max_retries: int
982
+ Maximum number of retries checking for vLLM server startup.
983
+ Only used when openai_api_server=True.
984
+ retry_alert_frequency: int
985
+ Frequency of alert logs for vLLM server startup retries.
986
+ Only used when openai_api_server=True.
987
+ engine_args : dict
988
+ Additional keyword arguments to pass to the vLLM engine.
989
+ For example, `tensor_parallel_size=2`.
970
990
  """
971
991
  ...
972
992
 
973
- @typing.overload
974
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
975
- ...
976
-
977
- @typing.overload
978
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
979
- ...
980
-
981
- 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):
993
+ 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]]]:
982
994
  """
983
- Specifies the resources needed when executing this step.
984
-
985
- Use `@resources` to specify the resource requirements
986
- independently of the specific compute layer (`@batch`, `@kubernetes`).
987
-
988
- You can choose the compute layer on the command line by executing e.g.
989
- ```
990
- python myflow.py run --with batch
991
- ```
992
- or
993
- ```
994
- python myflow.py run --with kubernetes
995
- ```
996
- which executes the flow on the desired system using the
997
- requirements specified in `@resources`.
995
+ Specifies that this step should execute on Kubernetes.
998
996
 
999
997
 
1000
998
  Parameters
1001
999
  ----------
1002
1000
  cpu : int, default 1
1003
- Number of CPUs required for this step.
1004
- gpu : int, optional, default None
1005
- Number of GPUs required for this step.
1006
- disk : int, optional, default None
1007
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1001
+ Number of CPUs required for this step. If `@resources` is
1002
+ also present, the maximum value from all decorators is used.
1008
1003
  memory : int, default 4096
1009
- Memory size (in MB) required for this step.
1010
- shared_memory : int, optional, default None
1011
- The value for the size (in MiB) of the /dev/shm volume for this step.
1012
- This parameter maps to the `--shm-size` option in Docker.
1013
- """
1014
- ...
1015
-
1016
- @typing.overload
1017
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1018
- """
1019
- Decorator prototype for all step decorators. This function gets specialized
1020
- and imported for all decorators types by _import_plugin_decorators().
1021
- """
1022
- ...
1023
-
1024
- @typing.overload
1025
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1026
- ...
1027
-
1028
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1029
- """
1030
- Decorator prototype for all step decorators. This function gets specialized
1031
- and imported for all decorators types by _import_plugin_decorators().
1032
- """
1033
- ...
1034
-
1035
- @typing.overload
1036
- def model(*, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1037
- """
1038
- Enables loading / saving of models within a step.
1039
-
1040
- > Examples
1041
- - Saving Models
1042
- ```python
1043
- @model
1044
- @step
1045
- def train(self):
1046
- # current.model.save returns a dictionary reference to the model saved
1047
- self.my_model = current.model.save(
1048
- path_to_my_model,
1049
- label="my_model",
1050
- metadata={
1051
- "epochs": 10,
1052
- "batch-size": 32,
1053
- "learning-rate": 0.001,
1054
- }
1055
- )
1056
- self.next(self.test)
1057
-
1058
- @model(load="my_model")
1059
- @step
1060
- def test(self):
1061
- # `current.model.loaded` returns a dictionary of the loaded models
1062
- # where the key is the name of the artifact and the value is the path to the model
1063
- print(os.listdir(current.model.loaded["my_model"]))
1064
- self.next(self.end)
1065
- ```
1066
-
1067
- - Loading models
1068
- ```python
1069
- @step
1070
- def train(self):
1071
- # current.model.load returns the path to the model loaded
1072
- checkpoint_path = current.model.load(
1073
- self.checkpoint_key,
1074
- )
1075
- model_path = current.model.load(
1076
- self.model,
1077
- )
1078
- self.next(self.test)
1079
- ```
1080
-
1081
-
1082
- Parameters
1083
- ----------
1084
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1085
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1086
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1087
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1088
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1089
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1004
+ Memory size (in MB) required for this step. If
1005
+ `@resources` is also present, the maximum value from all decorators is
1006
+ used.
1007
+ disk : int, default 10240
1008
+ Disk size (in MB) required for this step. If
1009
+ `@resources` is also present, the maximum value from all decorators is
1010
+ used.
1011
+ image : str, optional, default None
1012
+ Docker image to use when launching on Kubernetes. If not specified, and
1013
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1014
+ not, a default Docker image mapping to the current version of Python is used.
1015
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1016
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
1017
+ image_pull_secrets: List[str], default []
1018
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
1019
+ Kubernetes image pull secrets to use when pulling container images
1020
+ in Kubernetes.
1021
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1022
+ Kubernetes service account to use when launching pod in Kubernetes.
1023
+ secrets : List[str], optional, default None
1024
+ Kubernetes secrets to use when launching pod in Kubernetes. These
1025
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1026
+ in Metaflow configuration.
1027
+ node_selector: Union[Dict[str,str], str], optional, default None
1028
+ Kubernetes node selector(s) to apply to the pod running the task.
1029
+ Can be passed in as a comma separated string of values e.g.
1030
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
1031
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
1032
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1033
+ Kubernetes namespace to use when launching pod in Kubernetes.
1034
+ gpu : int, optional, default None
1035
+ Number of GPUs required for this step. A value of zero implies that
1036
+ the scheduled node should not have GPUs.
1037
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1038
+ The vendor of the GPUs to be used for this step.
1039
+ tolerations : List[Dict[str,str]], default []
1040
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1041
+ Kubernetes tolerations to use when launching pod in Kubernetes.
1042
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
1043
+ Kubernetes labels to use when launching pod in Kubernetes.
1044
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
1045
+ Kubernetes annotations to use when launching pod in Kubernetes.
1046
+ use_tmpfs : bool, default False
1047
+ This enables an explicit tmpfs mount for this step.
1048
+ tmpfs_tempdir : bool, default True
1049
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1050
+ tmpfs_size : int, optional, default: None
1051
+ The value for the size (in MiB) of the tmpfs mount for this step.
1052
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1053
+ memory allocated for this step.
1054
+ tmpfs_path : str, optional, default /metaflow_temp
1055
+ Path to tmpfs mount for this step.
1056
+ persistent_volume_claims : Dict[str, str], optional, default None
1057
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1058
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1059
+ shared_memory: int, optional
1060
+ Shared memory size (in MiB) required for this step
1061
+ port: int, optional
1062
+ Port number to specify in the Kubernetes job object
1063
+ compute_pool : str, optional, default None
1064
+ Compute pool to be used for for this step.
1065
+ If not specified, any accessible compute pool within the perimeter is used.
1066
+ hostname_resolution_timeout: int, default 10 * 60
1067
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
1068
+ Only applicable when @parallel is used.
1069
+ qos: str, default: Burstable
1070
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
1090
1071
 
1091
- temp_dir_root : str, default: None
1092
- The root directory under which `current.model.loaded` will store loaded models
1072
+ security_context: Dict[str, Any], optional, default None
1073
+ Container security context. Applies to the task container. Allows the following keys:
1074
+ - privileged: bool, optional, default None
1075
+ - allow_privilege_escalation: bool, optional, default None
1076
+ - run_as_user: int, optional, default None
1077
+ - run_as_group: int, optional, default None
1078
+ - run_as_non_root: bool, optional, default None
1093
1079
  """
1094
1080
  ...
1095
1081
 
1096
1082
  @typing.overload
1097
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1083
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1084
+ """
1085
+ Decorator prototype for all step decorators. This function gets specialized
1086
+ and imported for all decorators types by _import_plugin_decorators().
1087
+ """
1098
1088
  ...
1099
1089
 
1100
1090
  @typing.overload
1101
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1091
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1102
1092
  ...
1103
1093
 
1104
- def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load: typing.Union[typing.List[str], str, typing.List[typing.Tuple[str, typing.Optional[str]]]] = None, temp_dir_root: str = None):
1094
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1105
1095
  """
1106
- Enables loading / saving of models within a step.
1107
-
1108
- > Examples
1109
- - Saving Models
1110
- ```python
1111
- @model
1112
- @step
1113
- def train(self):
1114
- # current.model.save returns a dictionary reference to the model saved
1115
- self.my_model = current.model.save(
1116
- path_to_my_model,
1117
- label="my_model",
1118
- metadata={
1119
- "epochs": 10,
1120
- "batch-size": 32,
1121
- "learning-rate": 0.001,
1122
- }
1123
- )
1124
- self.next(self.test)
1125
-
1126
- @model(load="my_model")
1127
- @step
1128
- def test(self):
1129
- # `current.model.loaded` returns a dictionary of the loaded models
1130
- # where the key is the name of the artifact and the value is the path to the model
1131
- print(os.listdir(current.model.loaded["my_model"]))
1132
- self.next(self.end)
1133
- ```
1134
-
1135
- - Loading models
1136
- ```python
1137
- @step
1138
- def train(self):
1139
- # current.model.load returns the path to the model loaded
1140
- checkpoint_path = current.model.load(
1141
- self.checkpoint_key,
1142
- )
1143
- model_path = current.model.load(
1144
- self.model,
1145
- )
1146
- self.next(self.test)
1147
- ```
1096
+ Decorator prototype for all step decorators. This function gets specialized
1097
+ and imported for all decorators types by _import_plugin_decorators().
1098
+ """
1099
+ ...
1100
+
1101
+ @typing.overload
1102
+ 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]]]:
1103
+ """
1104
+ Specifies secrets to be retrieved and injected as environment variables prior to
1105
+ the execution of a step.
1148
1106
 
1149
1107
 
1150
1108
  Parameters
1151
1109
  ----------
1152
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1153
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1154
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1155
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1156
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1157
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1158
-
1159
- temp_dir_root : str, default: None
1160
- The root directory under which `current.model.loaded` will store loaded models
1110
+ sources : List[Union[str, Dict[str, Any]]], default: []
1111
+ List of secret specs, defining how the secrets are to be retrieved
1112
+ role : str, optional, default: None
1113
+ Role to use for fetching secrets
1161
1114
  """
1162
1115
  ...
1163
1116
 
1164
- 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]]]:
1117
+ @typing.overload
1118
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1119
+ ...
1120
+
1121
+ @typing.overload
1122
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1123
+ ...
1124
+
1125
+ 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):
1165
1126
  """
1166
- This decorator is used to run vllm APIs as Metaflow task sidecars.
1167
-
1168
- User code call
1169
- --------------
1170
- @vllm(
1171
- model="...",
1172
- ...
1173
- )
1174
-
1175
- Valid backend options
1176
- ---------------------
1177
- - 'local': Run as a separate process on the local task machine.
1178
-
1179
- Valid model options
1180
- -------------------
1181
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1182
-
1183
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
1184
- If you need multiple models, you must create multiple @vllm decorators.
1127
+ Specifies secrets to be retrieved and injected as environment variables prior to
1128
+ the execution of a step.
1185
1129
 
1186
1130
 
1187
1131
  Parameters
1188
1132
  ----------
1189
- model: str
1190
- HuggingFace model identifier to be served by vLLM.
1191
- backend: str
1192
- Determines where and how to run the vLLM process.
1193
- openai_api_server: bool
1194
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
1195
- Default is False (uses native engine).
1196
- Set to True for backward compatibility with existing code.
1197
- debug: bool
1198
- Whether to turn on verbose debugging logs.
1199
- card_refresh_interval: int
1200
- Interval in seconds for refreshing the vLLM status card.
1201
- Only used when openai_api_server=True.
1202
- max_retries: int
1203
- Maximum number of retries checking for vLLM server startup.
1204
- Only used when openai_api_server=True.
1205
- retry_alert_frequency: int
1206
- Frequency of alert logs for vLLM server startup retries.
1207
- Only used when openai_api_server=True.
1208
- engine_args : dict
1209
- Additional keyword arguments to pass to the vLLM engine.
1210
- For example, `tensor_parallel_size=2`.
1133
+ sources : List[Union[str, Dict[str, Any]]], default: []
1134
+ List of secret specs, defining how the secrets are to be retrieved
1135
+ role : str, optional, default: None
1136
+ Role to use for fetching secrets
1211
1137
  """
1212
1138
  ...
1213
1139
 
1214
1140
  @typing.overload
1215
- 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]]]:
1141
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1216
1142
  """
1217
- Specifies the Conda environment for the step.
1218
-
1219
- Information in this decorator will augment any
1220
- attributes set in the `@conda_base` flow-level decorator. Hence,
1221
- you can use `@conda_base` to set packages required by all
1222
- steps and use `@conda` to specify step-specific overrides.
1143
+ Internal decorator to support Fast bakery
1144
+ """
1145
+ ...
1146
+
1147
+ @typing.overload
1148
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1149
+ ...
1150
+
1151
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1152
+ """
1153
+ Internal decorator to support Fast bakery
1154
+ """
1155
+ ...
1156
+
1157
+ @typing.overload
1158
+ 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]]]:
1159
+ """
1160
+ Specifies the resources needed when executing this step.
1223
1161
 
1162
+ Use `@resources` to specify the resource requirements
1163
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1224
1164
 
1225
- Parameters
1226
- ----------
1227
- packages : Dict[str, str], default {}
1228
- Packages to use for this step. The key is the name of the package
1229
- and the value is the version to use.
1230
- libraries : Dict[str, str], default {}
1231
- Supported for backward compatibility. When used with packages, packages will take precedence.
1232
- python : str, optional, default None
1233
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1234
- that the version used will correspond to the version of the Python interpreter used to start the run.
1235
- disabled : bool, default False
1236
- If set to True, disables @conda.
1165
+ You can choose the compute layer on the command line by executing e.g.
1166
+ ```
1167
+ python myflow.py run --with batch
1168
+ ```
1169
+ or
1170
+ ```
1171
+ python myflow.py run --with kubernetes
1172
+ ```
1173
+ which executes the flow on the desired system using the
1174
+ requirements specified in `@resources`.
1175
+
1176
+
1177
+ Parameters
1178
+ ----------
1179
+ cpu : int, default 1
1180
+ Number of CPUs required for this step.
1181
+ gpu : int, optional, default None
1182
+ Number of GPUs required for this step.
1183
+ disk : int, optional, default None
1184
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1185
+ memory : int, default 4096
1186
+ Memory size (in MB) required for this step.
1187
+ shared_memory : int, optional, default None
1188
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1189
+ This parameter maps to the `--shm-size` option in Docker.
1237
1190
  """
1238
1191
  ...
1239
1192
 
1240
1193
  @typing.overload
1241
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1194
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1242
1195
  ...
1243
1196
 
1244
1197
  @typing.overload
1245
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1198
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1246
1199
  ...
1247
1200
 
1248
- 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):
1201
+ 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):
1249
1202
  """
1250
- Specifies the Conda environment for the step.
1203
+ Specifies the resources needed when executing this step.
1251
1204
 
1252
- Information in this decorator will augment any
1253
- attributes set in the `@conda_base` flow-level decorator. Hence,
1254
- you can use `@conda_base` to set packages required by all
1255
- steps and use `@conda` to specify step-specific overrides.
1205
+ Use `@resources` to specify the resource requirements
1206
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1207
+
1208
+ You can choose the compute layer on the command line by executing e.g.
1209
+ ```
1210
+ python myflow.py run --with batch
1211
+ ```
1212
+ or
1213
+ ```
1214
+ python myflow.py run --with kubernetes
1215
+ ```
1216
+ which executes the flow on the desired system using the
1217
+ requirements specified in `@resources`.
1256
1218
 
1257
1219
 
1258
1220
  Parameters
1259
1221
  ----------
1260
- packages : Dict[str, str], default {}
1261
- Packages to use for this step. The key is the name of the package
1262
- and the value is the version to use.
1263
- libraries : Dict[str, str], default {}
1264
- Supported for backward compatibility. When used with packages, packages will take precedence.
1265
- python : str, optional, default None
1266
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1267
- that the version used will correspond to the version of the Python interpreter used to start the run.
1268
- disabled : bool, default False
1269
- If set to True, disables @conda.
1222
+ cpu : int, default 1
1223
+ Number of CPUs required for this step.
1224
+ gpu : int, optional, default None
1225
+ Number of GPUs required for this step.
1226
+ disk : int, optional, default None
1227
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1228
+ memory : int, default 4096
1229
+ Memory size (in MB) required for this step.
1230
+ shared_memory : int, optional, default None
1231
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1232
+ This parameter maps to the `--shm-size` option in Docker.
1270
1233
  """
1271
1234
  ...
1272
1235
 
1273
1236
  @typing.overload
1274
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1237
+ 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]]]:
1275
1238
  """
1276
- Decorator prototype for all step decorators. This function gets specialized
1277
- and imported for all decorators types by _import_plugin_decorators().
1239
+ Specifies that the step will success under all circumstances.
1240
+
1241
+ The decorator will create an optional artifact, specified by `var`, which
1242
+ contains the exception raised. You can use it to detect the presence
1243
+ of errors, indicating that all happy-path artifacts produced by the step
1244
+ are missing.
1245
+
1246
+
1247
+ Parameters
1248
+ ----------
1249
+ var : str, optional, default None
1250
+ Name of the artifact in which to store the caught exception.
1251
+ If not specified, the exception is not stored.
1252
+ print_exception : bool, default True
1253
+ Determines whether or not the exception is printed to
1254
+ stdout when caught.
1278
1255
  """
1279
1256
  ...
1280
1257
 
1281
1258
  @typing.overload
1282
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1259
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1283
1260
  ...
1284
1261
 
1285
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1286
- """
1287
- Decorator prototype for all step decorators. This function gets specialized
1288
- and imported for all decorators types by _import_plugin_decorators().
1289
- """
1262
+ @typing.overload
1263
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1290
1264
  ...
1291
1265
 
1292
- @typing.overload
1293
- 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]]]:
1266
+ 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):
1294
1267
  """
1295
- Specifies a timeout for your step.
1296
-
1297
- This decorator is useful if this step may hang indefinitely.
1298
-
1299
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1300
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1301
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1268
+ Specifies that the step will success under all circumstances.
1302
1269
 
1303
- Note that all the values specified in parameters are added together so if you specify
1304
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1270
+ The decorator will create an optional artifact, specified by `var`, which
1271
+ contains the exception raised. You can use it to detect the presence
1272
+ of errors, indicating that all happy-path artifacts produced by the step
1273
+ are missing.
1305
1274
 
1306
1275
 
1307
1276
  Parameters
1308
1277
  ----------
1309
- seconds : int, default 0
1310
- Number of seconds to wait prior to timing out.
1311
- minutes : int, default 0
1312
- Number of minutes to wait prior to timing out.
1313
- hours : int, default 0
1314
- Number of hours to wait prior to timing out.
1278
+ var : str, optional, default None
1279
+ Name of the artifact in which to store the caught exception.
1280
+ If not specified, the exception is not stored.
1281
+ print_exception : bool, default True
1282
+ Determines whether or not the exception is printed to
1283
+ stdout when caught.
1315
1284
  """
1316
1285
  ...
1317
1286
 
1318
1287
  @typing.overload
1319
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1288
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1289
+ """
1290
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1291
+ It exists to make it easier for users to know that this decorator should only be used with
1292
+ a Neo Cloud like CoreWeave.
1293
+ """
1320
1294
  ...
1321
1295
 
1322
1296
  @typing.overload
1323
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1297
+ def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1324
1298
  ...
1325
1299
 
1326
- 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):
1300
+ def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1327
1301
  """
1328
- Specifies a timeout for your step.
1329
-
1330
- This decorator is useful if this step may hang indefinitely.
1331
-
1332
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1333
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1334
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1335
-
1336
- Note that all the values specified in parameters are added together so if you specify
1337
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1338
-
1339
-
1340
- Parameters
1341
- ----------
1342
- seconds : int, default 0
1343
- Number of seconds to wait prior to timing out.
1344
- minutes : int, default 0
1345
- Number of minutes to wait prior to timing out.
1346
- hours : int, default 0
1347
- Number of hours to wait prior to timing out.
1302
+ CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1303
+ It exists to make it easier for users to know that this decorator should only be used with
1304
+ a Neo Cloud like CoreWeave.
1348
1305
  """
1349
1306
  ...
1350
1307
 
1351
- @typing.overload
1352
- 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]]:
1308
+ 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]]]:
1353
1309
  """
1354
- Specifies the times when the flow should be run when running on a
1355
- production scheduler.
1310
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
1311
+
1312
+ User code call
1313
+ --------------
1314
+ @ollama(
1315
+ models=[...],
1316
+ ...
1317
+ )
1318
+
1319
+ Valid backend options
1320
+ ---------------------
1321
+ - 'local': Run as a separate process on the local task machine.
1322
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
1323
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
1324
+
1325
+ Valid model options
1326
+ -------------------
1327
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
1356
1328
 
1357
1329
 
1358
1330
  Parameters
1359
1331
  ----------
1360
- hourly : bool, default False
1361
- Run the workflow hourly.
1362
- daily : bool, default True
1363
- Run the workflow daily.
1364
- weekly : bool, default False
1365
- Run the workflow weekly.
1366
- cron : str, optional, default None
1367
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1368
- specified by this expression.
1369
- timezone : str, optional, default None
1370
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1371
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1332
+ models: list[str]
1333
+ List of Ollama containers running models in sidecars.
1334
+ backend: str
1335
+ Determines where and how to run the Ollama process.
1336
+ force_pull: bool
1337
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
1338
+ cache_update_policy: str
1339
+ Cache update policy: "auto", "force", or "never".
1340
+ force_cache_update: bool
1341
+ Simple override for "force" cache update policy.
1342
+ debug: bool
1343
+ Whether to turn on verbose debugging logs.
1344
+ circuit_breaker_config: dict
1345
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
1346
+ timeout_config: dict
1347
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
1372
1348
  """
1373
1349
  ...
1374
1350
 
1375
- @typing.overload
1376
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1377
- ...
1378
-
1379
- 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):
1351
+ 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]]:
1380
1352
  """
1381
- Specifies the times when the flow should be run when running on a
1382
- production scheduler.
1353
+ 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)
1354
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1355
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1356
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1357
+ starts only after all sensors finish.
1383
1358
 
1384
1359
 
1385
1360
  Parameters
1386
1361
  ----------
1387
- hourly : bool, default False
1388
- Run the workflow hourly.
1389
- daily : bool, default True
1390
- Run the workflow daily.
1391
- weekly : bool, default False
1392
- Run the workflow weekly.
1393
- cron : str, optional, default None
1394
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1395
- specified by this expression.
1396
- timezone : str, optional, default None
1397
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1398
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1362
+ timeout : int
1363
+ Time, in seconds before the task times out and fails. (Default: 3600)
1364
+ poke_interval : int
1365
+ Time in seconds that the job should wait in between each try. (Default: 60)
1366
+ mode : str
1367
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1368
+ exponential_backoff : bool
1369
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1370
+ pool : str
1371
+ the slot pool this task should run in,
1372
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1373
+ soft_fail : bool
1374
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1375
+ name : str
1376
+ Name of the sensor on Airflow
1377
+ description : str
1378
+ Description of sensor in the Airflow UI
1379
+ bucket_key : Union[str, List[str]]
1380
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1381
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1382
+ bucket_name : str
1383
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1384
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1385
+ wildcard_match : bool
1386
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1387
+ aws_conn_id : str
1388
+ a reference to the s3 connection on Airflow. (Default: None)
1389
+ verify : bool
1390
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1399
1391
  """
1400
1392
  ...
1401
1393
 
@@ -1442,142 +1434,6 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1442
1434
  """
1443
1435
  ...
1444
1436
 
1445
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1446
- """
1447
- Specifies what flows belong to the same project.
1448
-
1449
- A project-specific namespace is created for all flows that
1450
- use the same `@project(name)`.
1451
-
1452
-
1453
- Parameters
1454
- ----------
1455
- name : str
1456
- Project name. Make sure that the name is unique amongst all
1457
- projects that use the same production scheduler. The name may
1458
- contain only lowercase alphanumeric characters and underscores.
1459
-
1460
- branch : Optional[str], default None
1461
- The branch to use. If not specified, the branch is set to
1462
- `user.<username>` unless `production` is set to `True`. This can
1463
- also be set on the command line using `--branch` as a top-level option.
1464
- It is an error to specify `branch` in the decorator and on the command line.
1465
-
1466
- production : bool, default False
1467
- Whether or not the branch is the production branch. This can also be set on the
1468
- command line using `--production` as a top-level option. It is an error to specify
1469
- `production` in the decorator and on the command line.
1470
- The project branch name will be:
1471
- - if `branch` is specified:
1472
- - if `production` is True: `prod.<branch>`
1473
- - if `production` is False: `test.<branch>`
1474
- - if `branch` is not specified:
1475
- - if `production` is True: `prod`
1476
- - if `production` is False: `user.<username>`
1477
- """
1478
- ...
1479
-
1480
- @typing.overload
1481
- 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]]:
1482
- """
1483
- Specifies the flow(s) that this flow depends on.
1484
-
1485
- ```
1486
- @trigger_on_finish(flow='FooFlow')
1487
- ```
1488
- or
1489
- ```
1490
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1491
- ```
1492
- This decorator respects the @project decorator and triggers the flow
1493
- when upstream runs within the same namespace complete successfully
1494
-
1495
- Additionally, you can specify project aware upstream flow dependencies
1496
- by specifying the fully qualified project_flow_name.
1497
- ```
1498
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1499
- ```
1500
- or
1501
- ```
1502
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1503
- ```
1504
-
1505
- You can also specify just the project or project branch (other values will be
1506
- inferred from the current project or project branch):
1507
- ```
1508
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1509
- ```
1510
-
1511
- Note that `branch` is typically one of:
1512
- - `prod`
1513
- - `user.bob`
1514
- - `test.my_experiment`
1515
- - `prod.staging`
1516
-
1517
-
1518
- Parameters
1519
- ----------
1520
- flow : Union[str, Dict[str, str]], optional, default None
1521
- Upstream flow dependency for this flow.
1522
- flows : List[Union[str, Dict[str, str]]], default []
1523
- Upstream flow dependencies for this flow.
1524
- options : Dict[str, Any], default {}
1525
- Backend-specific configuration for tuning eventing behavior.
1526
- """
1527
- ...
1528
-
1529
- @typing.overload
1530
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1531
- ...
1532
-
1533
- 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] = {}):
1534
- """
1535
- Specifies the flow(s) that this flow depends on.
1536
-
1537
- ```
1538
- @trigger_on_finish(flow='FooFlow')
1539
- ```
1540
- or
1541
- ```
1542
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1543
- ```
1544
- This decorator respects the @project decorator and triggers the flow
1545
- when upstream runs within the same namespace complete successfully
1546
-
1547
- Additionally, you can specify project aware upstream flow dependencies
1548
- by specifying the fully qualified project_flow_name.
1549
- ```
1550
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1551
- ```
1552
- or
1553
- ```
1554
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1555
- ```
1556
-
1557
- You can also specify just the project or project branch (other values will be
1558
- inferred from the current project or project branch):
1559
- ```
1560
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1561
- ```
1562
-
1563
- Note that `branch` is typically one of:
1564
- - `prod`
1565
- - `user.bob`
1566
- - `test.my_experiment`
1567
- - `prod.staging`
1568
-
1569
-
1570
- Parameters
1571
- ----------
1572
- flow : Union[str, Dict[str, str]], optional, default None
1573
- Upstream flow dependency for this flow.
1574
- flows : List[Union[str, Dict[str, str]]], default []
1575
- Upstream flow dependencies for this flow.
1576
- options : Dict[str, Any], default {}
1577
- Backend-specific configuration for tuning eventing behavior.
1578
- """
1579
- ...
1580
-
1581
1437
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1582
1438
  """
1583
1439
  Allows setting external datastores to save data for the
@@ -1677,157 +1533,249 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1677
1533
  Parameters:
1678
1534
  ----------
1679
1535
 
1680
- type: str
1681
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1536
+ type: str
1537
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1538
+
1539
+ config: dict or Callable
1540
+ Dictionary of configuration options for the datastore. The following keys are required:
1541
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1542
+ - example: 's3://bucket-name/path/to/root'
1543
+ - example: 'gs://bucket-name/path/to/root'
1544
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1545
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1546
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1547
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1548
+ """
1549
+ ...
1550
+
1551
+ @typing.overload
1552
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1553
+ """
1554
+ Specifies the PyPI packages for all steps of the flow.
1555
+
1556
+ Use `@pypi_base` to set common packages required by all
1557
+ steps and use `@pypi` to specify step-specific overrides.
1558
+
1559
+ Parameters
1560
+ ----------
1561
+ packages : Dict[str, str], default: {}
1562
+ Packages to use for this flow. The key is the name of the package
1563
+ and the value is the version to use.
1564
+ python : str, optional, default: None
1565
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1566
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1567
+ """
1568
+ ...
1569
+
1570
+ @typing.overload
1571
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1572
+ ...
1573
+
1574
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1575
+ """
1576
+ Specifies the PyPI packages for all steps of the flow.
1577
+
1578
+ Use `@pypi_base` to set common packages required by all
1579
+ steps and use `@pypi` to specify step-specific overrides.
1580
+
1581
+ Parameters
1582
+ ----------
1583
+ packages : Dict[str, str], default: {}
1584
+ Packages to use for this flow. The key is the name of the package
1585
+ and the value is the version to use.
1586
+ python : str, optional, default: None
1587
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1588
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1589
+ """
1590
+ ...
1591
+
1592
+ @typing.overload
1593
+ 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]]:
1594
+ """
1595
+ Specifies the times when the flow should be run when running on a
1596
+ production scheduler.
1597
+
1598
+
1599
+ Parameters
1600
+ ----------
1601
+ hourly : bool, default False
1602
+ Run the workflow hourly.
1603
+ daily : bool, default True
1604
+ Run the workflow daily.
1605
+ weekly : bool, default False
1606
+ Run the workflow weekly.
1607
+ cron : str, optional, default None
1608
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1609
+ specified by this expression.
1610
+ timezone : str, optional, default None
1611
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1612
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1613
+ """
1614
+ ...
1615
+
1616
+ @typing.overload
1617
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1618
+ ...
1619
+
1620
+ 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):
1621
+ """
1622
+ Specifies the times when the flow should be run when running on a
1623
+ production scheduler.
1624
+
1625
+
1626
+ Parameters
1627
+ ----------
1628
+ hourly : bool, default False
1629
+ Run the workflow hourly.
1630
+ daily : bool, default True
1631
+ Run the workflow daily.
1632
+ weekly : bool, default False
1633
+ Run the workflow weekly.
1634
+ cron : str, optional, default None
1635
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1636
+ specified by this expression.
1637
+ timezone : str, optional, default None
1638
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1639
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1640
+ """
1641
+ ...
1642
+
1643
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1644
+ """
1645
+ Specifies what flows belong to the same project.
1646
+
1647
+ A project-specific namespace is created for all flows that
1648
+ use the same `@project(name)`.
1649
+
1650
+
1651
+ Parameters
1652
+ ----------
1653
+ name : str
1654
+ Project name. Make sure that the name is unique amongst all
1655
+ projects that use the same production scheduler. The name may
1656
+ contain only lowercase alphanumeric characters and underscores.
1657
+
1658
+ branch : Optional[str], default None
1659
+ The branch to use. If not specified, the branch is set to
1660
+ `user.<username>` unless `production` is set to `True`. This can
1661
+ also be set on the command line using `--branch` as a top-level option.
1662
+ It is an error to specify `branch` in the decorator and on the command line.
1682
1663
 
1683
- config: dict or Callable
1684
- Dictionary of configuration options for the datastore. The following keys are required:
1685
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1686
- - example: 's3://bucket-name/path/to/root'
1687
- - example: 'gs://bucket-name/path/to/root'
1688
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1689
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1690
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1691
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1664
+ production : bool, default False
1665
+ Whether or not the branch is the production branch. This can also be set on the
1666
+ command line using `--production` as a top-level option. It is an error to specify
1667
+ `production` in the decorator and on the command line.
1668
+ The project branch name will be:
1669
+ - if `branch` is specified:
1670
+ - if `production` is True: `prod.<branch>`
1671
+ - if `production` is False: `test.<branch>`
1672
+ - if `branch` is not specified:
1673
+ - if `production` is True: `prod`
1674
+ - if `production` is False: `user.<username>`
1692
1675
  """
1693
1676
  ...
1694
1677
 
1695
1678
  @typing.overload
1696
- def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1679
+ 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]]:
1697
1680
  """
1698
- Specifies the event(s) that this flow depends on.
1681
+ Specifies the flow(s) that this flow depends on.
1699
1682
 
1700
1683
  ```
1701
- @trigger(event='foo')
1684
+ @trigger_on_finish(flow='FooFlow')
1702
1685
  ```
1703
1686
  or
1704
1687
  ```
1705
- @trigger(events=['foo', 'bar'])
1688
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1706
1689
  ```
1690
+ This decorator respects the @project decorator and triggers the flow
1691
+ when upstream runs within the same namespace complete successfully
1707
1692
 
1708
- Additionally, you can specify the parameter mappings
1709
- to map event payload to Metaflow parameters for the flow.
1693
+ Additionally, you can specify project aware upstream flow dependencies
1694
+ by specifying the fully qualified project_flow_name.
1710
1695
  ```
1711
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1696
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1712
1697
  ```
1713
1698
  or
1714
1699
  ```
1715
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1716
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1700
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1717
1701
  ```
1718
1702
 
1719
- 'parameters' can also be a list of strings and tuples like so:
1720
- ```
1721
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1722
- ```
1723
- This is equivalent to:
1703
+ You can also specify just the project or project branch (other values will be
1704
+ inferred from the current project or project branch):
1724
1705
  ```
1725
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1706
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1726
1707
  ```
1727
1708
 
1709
+ Note that `branch` is typically one of:
1710
+ - `prod`
1711
+ - `user.bob`
1712
+ - `test.my_experiment`
1713
+ - `prod.staging`
1714
+
1728
1715
 
1729
1716
  Parameters
1730
1717
  ----------
1731
- event : Union[str, Dict[str, Any]], optional, default None
1732
- Event dependency for this flow.
1733
- events : List[Union[str, Dict[str, Any]]], default []
1734
- Events dependency for this flow.
1718
+ flow : Union[str, Dict[str, str]], optional, default None
1719
+ Upstream flow dependency for this flow.
1720
+ flows : List[Union[str, Dict[str, str]]], default []
1721
+ Upstream flow dependencies for this flow.
1735
1722
  options : Dict[str, Any], default {}
1736
1723
  Backend-specific configuration for tuning eventing behavior.
1737
1724
  """
1738
1725
  ...
1739
1726
 
1740
1727
  @typing.overload
1741
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1728
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1742
1729
  ...
1743
1730
 
1744
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1731
+ 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] = {}):
1745
1732
  """
1746
- Specifies the event(s) that this flow depends on.
1733
+ Specifies the flow(s) that this flow depends on.
1747
1734
 
1748
1735
  ```
1749
- @trigger(event='foo')
1736
+ @trigger_on_finish(flow='FooFlow')
1750
1737
  ```
1751
1738
  or
1752
1739
  ```
1753
- @trigger(events=['foo', 'bar'])
1740
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1754
1741
  ```
1742
+ This decorator respects the @project decorator and triggers the flow
1743
+ when upstream runs within the same namespace complete successfully
1755
1744
 
1756
- Additionally, you can specify the parameter mappings
1757
- to map event payload to Metaflow parameters for the flow.
1745
+ Additionally, you can specify project aware upstream flow dependencies
1746
+ by specifying the fully qualified project_flow_name.
1758
1747
  ```
1759
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1748
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1760
1749
  ```
1761
1750
  or
1762
1751
  ```
1763
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1764
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1752
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1765
1753
  ```
1766
1754
 
1767
- 'parameters' can also be a list of strings and tuples like so:
1768
- ```
1769
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1770
- ```
1771
- This is equivalent to:
1755
+ You can also specify just the project or project branch (other values will be
1756
+ inferred from the current project or project branch):
1772
1757
  ```
1773
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1758
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1774
1759
  ```
1775
1760
 
1761
+ Note that `branch` is typically one of:
1762
+ - `prod`
1763
+ - `user.bob`
1764
+ - `test.my_experiment`
1765
+ - `prod.staging`
1766
+
1776
1767
 
1777
1768
  Parameters
1778
1769
  ----------
1779
- event : Union[str, Dict[str, Any]], optional, default None
1780
- Event dependency for this flow.
1781
- events : List[Union[str, Dict[str, Any]]], default []
1782
- Events dependency for this flow.
1770
+ flow : Union[str, Dict[str, str]], optional, default None
1771
+ Upstream flow dependency for this flow.
1772
+ flows : List[Union[str, Dict[str, str]]], default []
1773
+ Upstream flow dependencies for this flow.
1783
1774
  options : Dict[str, Any], default {}
1784
1775
  Backend-specific configuration for tuning eventing behavior.
1785
1776
  """
1786
1777
  ...
1787
1778
 
1788
- 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]]:
1789
- """
1790
- 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)
1791
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1792
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1793
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1794
- starts only after all sensors finish.
1795
-
1796
-
1797
- Parameters
1798
- ----------
1799
- timeout : int
1800
- Time, in seconds before the task times out and fails. (Default: 3600)
1801
- poke_interval : int
1802
- Time in seconds that the job should wait in between each try. (Default: 60)
1803
- mode : str
1804
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1805
- exponential_backoff : bool
1806
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1807
- pool : str
1808
- the slot pool this task should run in,
1809
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1810
- soft_fail : bool
1811
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1812
- name : str
1813
- Name of the sensor on Airflow
1814
- description : str
1815
- Description of sensor in the Airflow UI
1816
- bucket_key : Union[str, List[str]]
1817
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1818
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1819
- bucket_name : str
1820
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1821
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1822
- wildcard_match : bool
1823
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1824
- aws_conn_id : str
1825
- a reference to the s3 connection on Airflow. (Default: None)
1826
- verify : bool
1827
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1828
- """
1829
- ...
1830
-
1831
1779
  @typing.overload
1832
1780
  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]]:
1833
1781
  """
@@ -1880,43 +1828,95 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1880
1828
  ...
1881
1829
 
1882
1830
  @typing.overload
1883
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1831
+ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1884
1832
  """
1885
- Specifies the PyPI packages for all steps of the flow.
1833
+ Specifies the event(s) that this flow depends on.
1834
+
1835
+ ```
1836
+ @trigger(event='foo')
1837
+ ```
1838
+ or
1839
+ ```
1840
+ @trigger(events=['foo', 'bar'])
1841
+ ```
1842
+
1843
+ Additionally, you can specify the parameter mappings
1844
+ to map event payload to Metaflow parameters for the flow.
1845
+ ```
1846
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1847
+ ```
1848
+ or
1849
+ ```
1850
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1851
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1852
+ ```
1853
+
1854
+ 'parameters' can also be a list of strings and tuples like so:
1855
+ ```
1856
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1857
+ ```
1858
+ This is equivalent to:
1859
+ ```
1860
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1861
+ ```
1886
1862
 
1887
- Use `@pypi_base` to set common packages required by all
1888
- steps and use `@pypi` to specify step-specific overrides.
1889
1863
 
1890
1864
  Parameters
1891
1865
  ----------
1892
- packages : Dict[str, str], default: {}
1893
- Packages to use for this flow. The key is the name of the package
1894
- and the value is the version to use.
1895
- python : str, optional, default: None
1896
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1897
- that the version used will correspond to the version of the Python interpreter used to start the run.
1866
+ event : Union[str, Dict[str, Any]], optional, default None
1867
+ Event dependency for this flow.
1868
+ events : List[Union[str, Dict[str, Any]]], default []
1869
+ Events dependency for this flow.
1870
+ options : Dict[str, Any], default {}
1871
+ Backend-specific configuration for tuning eventing behavior.
1898
1872
  """
1899
1873
  ...
1900
1874
 
1901
1875
  @typing.overload
1902
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1876
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1903
1877
  ...
1904
1878
 
1905
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1879
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1906
1880
  """
1907
- Specifies the PyPI packages for all steps of the flow.
1881
+ Specifies the event(s) that this flow depends on.
1882
+
1883
+ ```
1884
+ @trigger(event='foo')
1885
+ ```
1886
+ or
1887
+ ```
1888
+ @trigger(events=['foo', 'bar'])
1889
+ ```
1890
+
1891
+ Additionally, you can specify the parameter mappings
1892
+ to map event payload to Metaflow parameters for the flow.
1893
+ ```
1894
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1895
+ ```
1896
+ or
1897
+ ```
1898
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1899
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1900
+ ```
1901
+
1902
+ 'parameters' can also be a list of strings and tuples like so:
1903
+ ```
1904
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1905
+ ```
1906
+ This is equivalent to:
1907
+ ```
1908
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1909
+ ```
1908
1910
 
1909
- Use `@pypi_base` to set common packages required by all
1910
- steps and use `@pypi` to specify step-specific overrides.
1911
1911
 
1912
1912
  Parameters
1913
1913
  ----------
1914
- packages : Dict[str, str], default: {}
1915
- Packages to use for this flow. The key is the name of the package
1916
- and the value is the version to use.
1917
- python : str, optional, default: None
1918
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1919
- that the version used will correspond to the version of the Python interpreter used to start the run.
1914
+ event : Union[str, Dict[str, Any]], optional, default None
1915
+ Event dependency for this flow.
1916
+ events : List[Union[str, Dict[str, Any]]], default []
1917
+ Events dependency for this flow.
1918
+ options : Dict[str, Any], default {}
1919
+ Backend-specific configuration for tuning eventing behavior.
1920
1920
  """
1921
1921
  ...
1922
1922