ob-metaflow-stubs 6.0.10.15__py2.py3-none-any.whl → 6.0.10.17__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 (266) hide show
  1. metaflow-stubs/__init__.pyi +1053 -1053
  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 +2 -2
  9. metaflow-stubs/events.pyi +3 -3
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  14. metaflow-stubs/meta_files.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +47 -47
  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 +3 -3
  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 +2 -2
  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 +3 -3
  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/cards/__init__.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/cards/hf_hub_card.pyi +4 -4
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +3 -3
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  64. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  65. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/__init__.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_cli.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +3 -3
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +2 -2
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +2 -2
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +2 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +2 -2
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +3 -3
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +3 -3
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +4 -4
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +2 -2
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +2 -2
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +3 -3
  91. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  92. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +2 -2
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +3 -3
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  97. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +2 -2
  98. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +2 -2
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +2 -2
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +3 -3
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +3 -3
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +2 -2
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  110. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  111. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  112. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  114. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +3 -3
  115. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  116. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  117. metaflow-stubs/mf_extensions/outerbounds/toplevel/s3_proxy.pyi +2 -2
  118. metaflow-stubs/multicore_utils.pyi +2 -2
  119. metaflow-stubs/ob_internal.pyi +2 -2
  120. metaflow-stubs/packaging_sys/__init__.pyi +5 -5
  121. metaflow-stubs/packaging_sys/backend.pyi +4 -4
  122. metaflow-stubs/packaging_sys/distribution_support.pyi +5 -5
  123. metaflow-stubs/packaging_sys/tar_backend.pyi +5 -5
  124. metaflow-stubs/packaging_sys/utils.pyi +2 -2
  125. metaflow-stubs/packaging_sys/v1.pyi +2 -2
  126. metaflow-stubs/parameters.pyi +2 -2
  127. metaflow-stubs/plugins/__init__.pyi +10 -10
  128. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  129. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  130. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  131. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  132. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  133. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  134. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  135. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  136. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  137. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  138. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  139. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  140. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  141. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  142. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  143. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  144. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  145. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  146. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  147. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  148. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  149. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  150. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  151. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  152. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  153. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  154. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  155. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  156. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  157. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  158. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  159. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  160. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  161. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  162. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  163. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  164. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  165. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  166. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  167. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  168. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  169. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  170. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  171. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/cards/card_modules/basic.pyi +5 -3
  173. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  174. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  175. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +3 -3
  176. metaflow-stubs/plugins/cards/card_modules/json_viewer.pyi +2 -2
  177. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  178. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  179. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  180. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  181. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  182. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  183. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  184. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  185. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  186. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  187. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  188. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  189. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  190. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  191. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  192. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  193. metaflow-stubs/plugins/exit_hook/__init__.pyi +2 -2
  194. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +2 -2
  195. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  196. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  197. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  198. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  199. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  200. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  201. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  202. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  203. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  204. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  205. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  206. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  207. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  208. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  209. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  210. metaflow-stubs/plugins/ollama/__init__.pyi +3 -3
  211. metaflow-stubs/plugins/optuna/__init__.pyi +2 -2
  212. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  213. metaflow-stubs/plugins/parsers.pyi +2 -2
  214. metaflow-stubs/plugins/perimeters.pyi +2 -2
  215. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  216. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  217. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  218. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  219. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  220. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  221. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  222. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  223. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  224. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  225. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  226. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  227. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  228. metaflow-stubs/plugins/secrets/secrets_func.pyi +2 -2
  229. metaflow-stubs/plugins/secrets/secrets_spec.pyi +2 -2
  230. metaflow-stubs/plugins/secrets/utils.pyi +2 -2
  231. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  232. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  233. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  234. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  235. metaflow-stubs/plugins/torchtune/__init__.pyi +2 -2
  236. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  237. metaflow-stubs/plugins/uv/uv_environment.pyi +3 -3
  238. metaflow-stubs/profilers/__init__.pyi +2 -2
  239. metaflow-stubs/pylint_wrapper.pyi +2 -2
  240. metaflow-stubs/runner/__init__.pyi +2 -2
  241. metaflow-stubs/runner/deployer.pyi +5 -5
  242. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  243. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  244. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  245. metaflow-stubs/runner/nbrun.pyi +2 -2
  246. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  247. metaflow-stubs/runner/utils.pyi +4 -4
  248. metaflow-stubs/system/__init__.pyi +2 -2
  249. metaflow-stubs/system/system_logger.pyi +2 -2
  250. metaflow-stubs/system/system_monitor.pyi +2 -2
  251. metaflow-stubs/tagging_util.pyi +2 -2
  252. metaflow-stubs/tuple_util.pyi +2 -2
  253. metaflow-stubs/user_configs/__init__.pyi +2 -2
  254. metaflow-stubs/user_configs/config_options.pyi +2 -2
  255. metaflow-stubs/user_configs/config_parameters.pyi +7 -7
  256. metaflow-stubs/user_decorators/__init__.pyi +2 -2
  257. metaflow-stubs/user_decorators/common.pyi +2 -2
  258. metaflow-stubs/user_decorators/mutable_flow.pyi +6 -6
  259. metaflow-stubs/user_decorators/mutable_step.pyi +5 -5
  260. metaflow-stubs/user_decorators/user_flow_decorator.pyi +3 -3
  261. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  262. {ob_metaflow_stubs-6.0.10.15.dist-info → ob_metaflow_stubs-6.0.10.17.dist-info}/METADATA +1 -1
  263. ob_metaflow_stubs-6.0.10.17.dist-info/RECORD +266 -0
  264. ob_metaflow_stubs-6.0.10.15.dist-info/RECORD +0 -266
  265. {ob_metaflow_stubs-6.0.10.15.dist-info → ob_metaflow_stubs-6.0.10.17.dist-info}/WHEEL +0 -0
  266. {ob_metaflow_stubs-6.0.10.15.dist-info → ob_metaflow_stubs-6.0.10.17.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.18.10.1+obcheckpoint(0.2.8);ob(v1) #
4
- # Generated on 2025-10-09T09:15:42.339267 #
3
+ # MF version: 2.18.11.1+obcheckpoint(0.2.8);ob(v1) #
4
+ # Generated on 2025-10-13T21:06:57.979951 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -40,17 +40,17 @@ from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
42
  from . import metaflow_git as metaflow_git
43
- from . import cards as cards
44
43
  from . import tuple_util as tuple_util
44
+ from . import cards as cards
45
45
  from . import events as events
46
46
  from . import runner as runner
47
47
  from . import plugins as plugins
48
48
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
49
49
  from . import includefile as includefile
50
50
  from .includefile import IncludeFile as IncludeFile
51
+ from .plugins.parsers import yaml_parser as yaml_parser
51
52
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
52
53
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
53
- from .plugins.parsers import yaml_parser as yaml_parser
54
54
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
55
55
  from . import client as client
56
56
  from .client.core import namespace as namespace
@@ -170,188 +170,63 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
170
170
  ...
171
171
 
172
172
  @typing.overload
173
- 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]]]:
174
- """
175
- Specifies the resources needed when executing this step.
176
-
177
- Use `@resources` to specify the resource requirements
178
- independently of the specific compute layer (`@batch`, `@kubernetes`).
179
-
180
- You can choose the compute layer on the command line by executing e.g.
181
- ```
182
- python myflow.py run --with batch
183
- ```
184
- or
185
- ```
186
- python myflow.py run --with kubernetes
187
- ```
188
- which executes the flow on the desired system using the
189
- requirements specified in `@resources`.
190
-
191
-
192
- Parameters
193
- ----------
194
- cpu : int, default 1
195
- Number of CPUs required for this step.
196
- gpu : int, optional, default None
197
- Number of GPUs required for this step.
198
- disk : int, optional, default None
199
- Disk size (in MB) required for this step. Only applies on Kubernetes.
200
- memory : int, default 4096
201
- Memory size (in MB) required for this step.
202
- shared_memory : int, optional, default None
203
- The value for the size (in MiB) of the /dev/shm volume for this step.
204
- This parameter maps to the `--shm-size` option in Docker.
205
- """
206
- ...
207
-
208
- @typing.overload
209
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
210
- ...
211
-
212
- @typing.overload
213
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
214
- ...
215
-
216
- 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):
217
- """
218
- Specifies the resources needed when executing this step.
219
-
220
- Use `@resources` to specify the resource requirements
221
- independently of the specific compute layer (`@batch`, `@kubernetes`).
222
-
223
- You can choose the compute layer on the command line by executing e.g.
224
- ```
225
- python myflow.py run --with batch
226
- ```
227
- or
228
- ```
229
- python myflow.py run --with kubernetes
230
- ```
231
- which executes the flow on the desired system using the
232
- requirements specified in `@resources`.
233
-
234
-
235
- Parameters
236
- ----------
237
- cpu : int, default 1
238
- Number of CPUs required for this step.
239
- gpu : int, optional, default None
240
- Number of GPUs required for this step.
241
- disk : int, optional, default None
242
- Disk size (in MB) required for this step. Only applies on Kubernetes.
243
- memory : int, default 4096
244
- Memory size (in MB) required for this step.
245
- shared_memory : int, optional, default None
246
- The value for the size (in MiB) of the /dev/shm volume for this step.
247
- This parameter maps to the `--shm-size` option in Docker.
248
- """
249
- ...
250
-
251
- @typing.overload
252
- 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]]]:
173
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
253
174
  """
254
- Specifies the Conda environment for the step.
175
+ Specifies the PyPI packages for the step.
255
176
 
256
177
  Information in this decorator will augment any
257
- attributes set in the `@conda_base` flow-level decorator. Hence,
258
- you can use `@conda_base` to set packages required by all
259
- steps and use `@conda` to specify step-specific overrides.
178
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
179
+ you can use `@pypi_base` to set packages required by all
180
+ steps and use `@pypi` to specify step-specific overrides.
260
181
 
261
182
 
262
183
  Parameters
263
184
  ----------
264
- packages : Dict[str, str], default {}
185
+ packages : Dict[str, str], default: {}
265
186
  Packages to use for this step. The key is the name of the package
266
187
  and the value is the version to use.
267
- libraries : Dict[str, str], default {}
268
- Supported for backward compatibility. When used with packages, packages will take precedence.
269
- python : str, optional, default None
188
+ python : str, optional, default: None
270
189
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
271
190
  that the version used will correspond to the version of the Python interpreter used to start the run.
272
- disabled : bool, default False
273
- If set to True, disables @conda.
274
191
  """
275
192
  ...
276
193
 
277
194
  @typing.overload
278
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
195
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
279
196
  ...
280
197
 
281
198
  @typing.overload
282
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
199
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
283
200
  ...
284
201
 
285
- 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):
202
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
286
203
  """
287
- Specifies the Conda environment for the step.
204
+ Specifies the PyPI packages for the step.
288
205
 
289
206
  Information in this decorator will augment any
290
- attributes set in the `@conda_base` flow-level decorator. Hence,
291
- you can use `@conda_base` to set packages required by all
292
- steps and use `@conda` to specify step-specific overrides.
207
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
208
+ you can use `@pypi_base` to set packages required by all
209
+ steps and use `@pypi` to specify step-specific overrides.
293
210
 
294
211
 
295
212
  Parameters
296
213
  ----------
297
- packages : Dict[str, str], default {}
214
+ packages : Dict[str, str], default: {}
298
215
  Packages to use for this step. The key is the name of the package
299
216
  and the value is the version to use.
300
- libraries : Dict[str, str], default {}
301
- Supported for backward compatibility. When used with packages, packages will take precedence.
302
- python : str, optional, default None
217
+ python : str, optional, default: None
303
218
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
304
219
  that the version used will correspond to the version of the Python interpreter used to start the run.
305
- disabled : bool, default False
306
- If set to True, disables @conda.
307
220
  """
308
221
  ...
309
222
 
310
- 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]]]:
223
+ def nebius_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]]]:
311
224
  """
312
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
313
-
314
- User code call
315
- --------------
316
- @ollama(
317
- models=[...],
318
- ...
319
- )
320
-
321
- Valid backend options
322
- ---------------------
323
- - 'local': Run as a separate process on the local task machine.
324
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
325
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
326
-
327
- Valid model options
328
- -------------------
329
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
225
+ `@nebius_s3_proxy` is a Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
226
+ It exists to make it easier for users to know that this decorator should only be used with
227
+ a Neo Cloud like Nebius. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
330
228
 
331
229
 
332
- Parameters
333
- ----------
334
- models: list[str]
335
- List of Ollama containers running models in sidecars.
336
- backend: str
337
- Determines where and how to run the Ollama process.
338
- force_pull: bool
339
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
340
- cache_update_policy: str
341
- Cache update policy: "auto", "force", or "never".
342
- force_cache_update: bool
343
- Simple override for "force" cache update policy.
344
- debug: bool
345
- Whether to turn on verbose debugging logs.
346
- circuit_breaker_config: dict
347
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
348
- timeout_config: dict
349
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
350
- """
351
- ...
352
-
353
- 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]]]:
354
- """
355
230
  Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
356
231
  for S3 read and write requests.
357
232
 
@@ -410,375 +285,90 @@ def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typin
410
285
  ...
411
286
 
412
287
  @typing.overload
413
- 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]]]:
288
+ 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]]]:
414
289
  """
415
- Creates a human-readable report, a Metaflow Card, after this step completes.
290
+ Enables checkpointing for a step.
416
291
 
417
- Note that you may add multiple `@card` decorators in a step with different parameters.
292
+ > Examples
293
+
294
+ - Saving Checkpoints
295
+
296
+ ```python
297
+ @checkpoint
298
+ @step
299
+ def train(self):
300
+ model = create_model(self.parameters, checkpoint_path = None)
301
+ for i in range(self.epochs):
302
+ # some training logic
303
+ loss = model.train(self.dataset)
304
+ if i % 10 == 0:
305
+ model.save(
306
+ current.checkpoint.directory,
307
+ )
308
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
309
+ # and returns a reference dictionary to the checkpoint saved in the datastore
310
+ self.latest_checkpoint = current.checkpoint.save(
311
+ name="epoch_checkpoint",
312
+ metadata={
313
+ "epoch": i,
314
+ "loss": loss,
315
+ }
316
+ )
317
+ ```
318
+
319
+ - Using Loaded Checkpoints
320
+
321
+ ```python
322
+ @retry(times=3)
323
+ @checkpoint
324
+ @step
325
+ def train(self):
326
+ # Assume that the task has restarted and the previous attempt of the task
327
+ # saved a checkpoint
328
+ checkpoint_path = None
329
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
330
+ print("Loaded checkpoint from the previous attempt")
331
+ checkpoint_path = current.checkpoint.directory
332
+
333
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
334
+ for i in range(self.epochs):
335
+ ...
336
+ ```
418
337
 
419
338
 
420
339
  Parameters
421
340
  ----------
422
- type : str, default 'default'
423
- Card type.
424
- id : str, optional, default None
425
- If multiple cards are present, use this id to identify this card.
426
- options : Dict[str, Any], default {}
427
- Options passed to the card. The contents depend on the card type.
428
- timeout : int, default 45
429
- Interrupt reporting if it takes more than this many seconds.
341
+ load_policy : str, default: "fresh"
342
+ The policy for loading the checkpoint. The following policies are supported:
343
+ - "eager": Loads the the latest available checkpoint within the namespace.
344
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
345
+ will be loaded at the start of the task.
346
+ - "none": Do not load any checkpoint
347
+ - "fresh": Loads the lastest checkpoint created within the running Task.
348
+ This mode helps loading checkpoints across various retry attempts of the same task.
349
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
350
+ created within the task will be loaded when the task is retries execution on failure.
351
+
352
+ temp_dir_root : str, default: None
353
+ The root directory under which `current.checkpoint.directory` will be created.
430
354
  """
431
355
  ...
432
356
 
433
357
  @typing.overload
434
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
358
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
435
359
  ...
436
360
 
437
361
  @typing.overload
438
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
362
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
439
363
  ...
440
364
 
441
- 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):
365
+ 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):
442
366
  """
443
- Creates a human-readable report, a Metaflow Card, after this step completes.
367
+ Enables checkpointing for a step.
444
368
 
445
- Note that you may add multiple `@card` decorators in a step with different parameters.
369
+ > Examples
446
370
 
447
-
448
- Parameters
449
- ----------
450
- type : str, default 'default'
451
- Card type.
452
- id : str, optional, default None
453
- If multiple cards are present, use this id to identify this card.
454
- options : Dict[str, Any], default {}
455
- Options passed to the card. The contents depend on the card type.
456
- timeout : int, default 45
457
- Interrupt reporting if it takes more than this many seconds.
458
- """
459
- ...
460
-
461
- def nebius_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]]]:
462
- """
463
- `@nebius_s3_proxy` is a Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
464
- It exists to make it easier for users to know that this decorator should only be used with
465
- a Neo Cloud like Nebius. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
466
-
467
-
468
- Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
469
- for S3 read and write requests.
470
-
471
- This decorator requires an integration in the Outerbounds platform that
472
- points to an external bucket. It affects S3 operations performed via
473
- Metaflow's `get_aws_client` and `S3` within a `@step`.
474
-
475
- Read operations
476
- ---------------
477
- All read operations pass through the proxy. If an object does not already
478
- exist in the external bucket, it is cached there. For example, if code reads
479
- from buckets `FOO` and `BAR` using the `S3` interface, objects from both
480
- buckets are cached in the external bucket.
481
-
482
- During task execution, all S3‑related read requests are routed through the
483
- proxy:
484
- - If the object is present in the external object store, the proxy
485
- streams it directly from there without accessing the requested origin
486
- bucket.
487
- - If the object is not present in the external storage, the proxy
488
- fetches it from the requested bucket, caches it in the external
489
- storage, and streams the response from the origin bucket.
490
-
491
- Warning
492
- -------
493
- All READ operations (e.g., GetObject, HeadObject) pass through the external
494
- bucket regardless of the bucket specified in user code. Even
495
- `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
496
- external bucket cache.
497
-
498
- Write operations
499
- ----------------
500
- Write behavior is controlled by the `write_mode` parameter, which determines
501
- whether writes also persist objects in the cache.
502
-
503
- `write_mode` values:
504
- - `origin-and-cache`: objects are written both to the cache and to their
505
- intended origin bucket.
506
- - `origin`: objects are written only to their intended origin bucket.
507
-
508
-
509
- Parameters
510
- ----------
511
- integration_name : str, optional
512
- [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
513
- that holds the configuration for the external, S3‑compatible object
514
- storage bucket. If not specified, the only available S3 proxy
515
- integration in the namespace is used (fails if multiple exist).
516
- write_mode : str, optional
517
- Controls whether writes also go to the external bucket.
518
- - `origin` (default)
519
- - `origin-and-cache`
520
- debug : bool, optional
521
- Enables debug logging for proxy operations.
522
- """
523
- ...
524
-
525
- @typing.overload
526
- 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]]]:
527
- """
528
- Specifies secrets to be retrieved and injected as environment variables prior to
529
- the execution of a step.
530
-
531
-
532
- Parameters
533
- ----------
534
- sources : List[Union[str, Dict[str, Any]]], default: []
535
- List of secret specs, defining how the secrets are to be retrieved
536
- role : str, optional, default: None
537
- Role to use for fetching secrets
538
- """
539
- ...
540
-
541
- @typing.overload
542
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
543
- ...
544
-
545
- @typing.overload
546
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
- ...
548
-
549
- 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):
550
- """
551
- Specifies secrets to be retrieved and injected as environment variables prior to
552
- the execution of a step.
553
-
554
-
555
- Parameters
556
- ----------
557
- sources : List[Union[str, Dict[str, Any]]], default: []
558
- List of secret specs, defining how the secrets are to be retrieved
559
- role : str, optional, default: None
560
- Role to use for fetching secrets
561
- """
562
- ...
563
-
564
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, cache_scope: 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]]]:
565
- """
566
- Decorator that helps cache, version, and store models/datasets from the Hugging Face Hub.
567
-
568
- Examples
569
- --------
570
-
571
- ```python
572
- # **Usage: creating references to models from the Hugging Face Hub that may be loaded in downstream steps**
573
- @huggingface_hub
574
- @step
575
- def pull_model_from_huggingface(self):
576
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
577
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
578
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
579
- # value of the function is a reference to the model in the backend storage.
580
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
581
-
582
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
583
- self.llama_model = current.huggingface_hub.snapshot_download(
584
- repo_id=self.model_id,
585
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
586
- )
587
- self.next(self.train)
588
-
589
- # **Usage: explicitly loading models at runtime from the Hugging Face Hub or from cache (from Metaflow's datastore)**
590
- @huggingface_hub
591
- @step
592
- def run_training(self):
593
- # Temporary directory (auto-cleaned on exit)
594
- with current.huggingface_hub.load(
595
- repo_id="google-bert/bert-base-uncased",
596
- allow_patterns=["*.bin"],
597
- ) as local_path:
598
- # Use files under local_path
599
- train_model(local_path)
600
- ...
601
-
602
- # **Usage: loading models directly from the Hugging Face Hub or from cache (from Metaflow's datastore)**
603
-
604
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
605
- @step
606
- def pull_model_from_huggingface(self):
607
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
608
-
609
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora", "/my-lora-directory")])
610
- @step
611
- def finetune_model(self):
612
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
613
- # path_to_model will be /my-directory
614
-
615
-
616
- # Takes all the arguments passed to `snapshot_download`
617
- # except for `local_dir`
618
- @huggingface_hub(load=[
619
- {
620
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
621
- },
622
- {
623
- "repo_id": "myorg/mistral-lora",
624
- "repo_type": "model",
625
- },
626
- ])
627
- @step
628
- def finetune_model(self):
629
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
630
- # path_to_model will be /my-directory
631
- ```
632
-
633
-
634
- Parameters
635
- ----------
636
- temp_dir_root : str, optional
637
- The root directory that will hold the temporary directory where objects will be downloaded.
638
-
639
- cache_scope : str, optional
640
- The scope of the cache. Can be `checkpoint` / `flow` / `global`.
641
- - `checkpoint` (default): All repos are stored like objects saved by `@checkpoint`.
642
- i.e., the cached path is derived from the namespace, flow, step, and Metaflow foreach iteration.
643
- Any repo downloaded under this scope will only be retrieved from the cache when the step runs under the same namespace in the same flow (at the same foreach index).
644
-
645
- - `flow`: All repos are cached under the flow, regardless of namespace.
646
- i.e., the cached path is derived solely from the flow name.
647
- When to use this mode: (1) Multiple users are executing the same flow and want shared access to the repos cached by the decorator. (2) Multiple versions of a flow are deployed, all needing access to the same repos cached by the decorator.
648
-
649
- - `global`: All repos are cached under a globally static path.
650
- i.e., the base path of the cache is static and all repos are stored under it.
651
- When to use this mode:
652
- - All repos from the Hugging Face Hub need to be shared by users across all flow executions.
653
- - Each caching scope comes with its own trade-offs:
654
- - `checkpoint`:
655
- - Has explicit control over when caches are populated (controlled by the same flow that has the `@huggingface_hub` decorator) but ends up hitting the Hugging Face Hub more often if there are many users/namespaces/steps.
656
- - Since objects are written on a `namespace/flow/step` basis, the blast radius of a bad checkpoint is limited to a particular flow in a namespace.
657
- - `flow`:
658
- - Has less control over when caches are populated (can be written by any execution instance of a flow from any namespace) but results in more cache hits.
659
- - The blast radius of a bad checkpoint is limited to all runs of a particular flow.
660
- - It doesn't promote cache reuse across flows.
661
- - `global`:
662
- - Has no control over when caches are populated (can be written by any flow execution) but has the highest cache hit rate.
663
- - It promotes cache reuse across flows.
664
- - The blast radius of a bad checkpoint spans every flow that could be using a particular repo.
665
-
666
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
667
- The list of repos (models/datasets) to load.
668
-
669
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
670
-
671
- - If repo (model/dataset) is not found in the datastore:
672
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
673
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
674
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
675
-
676
- - If repo is found in the datastore:
677
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
678
- """
679
- ...
680
-
681
- 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]]]:
682
- """
683
- Specifies that this step should execute on DGX cloud.
684
-
685
-
686
- Parameters
687
- ----------
688
- gpu : int
689
- Number of GPUs to use.
690
- gpu_type : str
691
- Type of Nvidia GPU to use.
692
- queue_timeout : int
693
- Time to keep the job in NVCF's queue.
694
- """
695
- ...
696
-
697
- @typing.overload
698
- 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]]]:
699
- """
700
- Enables checkpointing for a step.
701
-
702
- > Examples
703
-
704
- - Saving Checkpoints
705
-
706
- ```python
707
- @checkpoint
708
- @step
709
- def train(self):
710
- model = create_model(self.parameters, checkpoint_path = None)
711
- for i in range(self.epochs):
712
- # some training logic
713
- loss = model.train(self.dataset)
714
- if i % 10 == 0:
715
- model.save(
716
- current.checkpoint.directory,
717
- )
718
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
719
- # and returns a reference dictionary to the checkpoint saved in the datastore
720
- self.latest_checkpoint = current.checkpoint.save(
721
- name="epoch_checkpoint",
722
- metadata={
723
- "epoch": i,
724
- "loss": loss,
725
- }
726
- )
727
- ```
728
-
729
- - Using Loaded Checkpoints
730
-
731
- ```python
732
- @retry(times=3)
733
- @checkpoint
734
- @step
735
- def train(self):
736
- # Assume that the task has restarted and the previous attempt of the task
737
- # saved a checkpoint
738
- checkpoint_path = None
739
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
740
- print("Loaded checkpoint from the previous attempt")
741
- checkpoint_path = current.checkpoint.directory
742
-
743
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
744
- for i in range(self.epochs):
745
- ...
746
- ```
747
-
748
-
749
- Parameters
750
- ----------
751
- load_policy : str, default: "fresh"
752
- The policy for loading the checkpoint. The following policies are supported:
753
- - "eager": Loads the the latest available checkpoint within the namespace.
754
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
755
- will be loaded at the start of the task.
756
- - "none": Do not load any checkpoint
757
- - "fresh": Loads the lastest checkpoint created within the running Task.
758
- This mode helps loading checkpoints across various retry attempts of the same task.
759
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
760
- created within the task will be loaded when the task is retries execution on failure.
761
-
762
- temp_dir_root : str, default: None
763
- The root directory under which `current.checkpoint.directory` will be created.
764
- """
765
- ...
766
-
767
- @typing.overload
768
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
769
- ...
770
-
771
- @typing.overload
772
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
773
- ...
774
-
775
- 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):
776
- """
777
- Enables checkpointing for a step.
778
-
779
- > Examples
780
-
781
- - Saving Checkpoints
371
+ - Saving Checkpoints
782
372
 
783
373
  ```python
784
374
  @checkpoint
@@ -842,21 +432,209 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
842
432
  ...
843
433
 
844
434
  @typing.overload
845
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
435
+ 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]]]:
846
436
  """
847
- A simple decorator that demonstrates using CardDecoratorInjector
848
- to inject a card and render simple markdown content.
437
+ Specifies the Conda environment for the step.
438
+
439
+ Information in this decorator will augment any
440
+ attributes set in the `@conda_base` flow-level decorator. Hence,
441
+ you can use `@conda_base` to set packages required by all
442
+ steps and use `@conda` to specify step-specific overrides.
443
+
444
+
445
+ Parameters
446
+ ----------
447
+ packages : Dict[str, str], default {}
448
+ Packages to use for this step. The key is the name of the package
449
+ and the value is the version to use.
450
+ libraries : Dict[str, str], default {}
451
+ Supported for backward compatibility. When used with packages, packages will take precedence.
452
+ python : str, optional, default None
453
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
454
+ that the version used will correspond to the version of the Python interpreter used to start the run.
455
+ disabled : bool, default False
456
+ If set to True, disables @conda.
849
457
  """
850
458
  ...
851
459
 
852
460
  @typing.overload
853
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
461
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
854
462
  ...
855
463
 
856
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
464
+ @typing.overload
465
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
466
+ ...
467
+
468
+ 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):
857
469
  """
858
- A simple decorator that demonstrates using CardDecoratorInjector
859
- to inject a card and render simple markdown content.
470
+ Specifies the Conda environment for the step.
471
+
472
+ Information in this decorator will augment any
473
+ attributes set in the `@conda_base` flow-level decorator. Hence,
474
+ you can use `@conda_base` to set packages required by all
475
+ steps and use `@conda` to specify step-specific overrides.
476
+
477
+
478
+ Parameters
479
+ ----------
480
+ packages : Dict[str, str], default {}
481
+ Packages to use for this step. The key is the name of the package
482
+ and the value is the version to use.
483
+ libraries : Dict[str, str], default {}
484
+ Supported for backward compatibility. When used with packages, packages will take precedence.
485
+ python : str, optional, default None
486
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
487
+ that the version used will correspond to the version of the Python interpreter used to start the run.
488
+ disabled : bool, default False
489
+ If set to True, disables @conda.
490
+ """
491
+ ...
492
+
493
+ @typing.overload
494
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
495
+ """
496
+ Internal decorator to support Fast bakery
497
+ """
498
+ ...
499
+
500
+ @typing.overload
501
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
502
+ ...
503
+
504
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
505
+ """
506
+ Internal decorator to support Fast bakery
507
+ """
508
+ ...
509
+
510
+ @typing.overload
511
+ 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]]]:
512
+ """
513
+ Specifies secrets to be retrieved and injected as environment variables prior to
514
+ the execution of a step.
515
+
516
+
517
+ Parameters
518
+ ----------
519
+ sources : List[Union[str, Dict[str, Any]]], default: []
520
+ List of secret specs, defining how the secrets are to be retrieved
521
+ role : str, optional, default: None
522
+ Role to use for fetching secrets
523
+ """
524
+ ...
525
+
526
+ @typing.overload
527
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
528
+ ...
529
+
530
+ @typing.overload
531
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
532
+ ...
533
+
534
+ 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):
535
+ """
536
+ Specifies secrets to be retrieved and injected as environment variables prior to
537
+ the execution of a step.
538
+
539
+
540
+ Parameters
541
+ ----------
542
+ sources : List[Union[str, Dict[str, Any]]], default: []
543
+ List of secret specs, defining how the secrets are to be retrieved
544
+ role : str, optional, default: None
545
+ Role to use for fetching secrets
546
+ """
547
+ ...
548
+
549
+ @typing.overload
550
+ 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]]]:
551
+ """
552
+ Specifies a timeout for your step.
553
+
554
+ This decorator is useful if this step may hang indefinitely.
555
+
556
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
557
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
558
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
559
+
560
+ Note that all the values specified in parameters are added together so if you specify
561
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
562
+
563
+
564
+ Parameters
565
+ ----------
566
+ seconds : int, default 0
567
+ Number of seconds to wait prior to timing out.
568
+ minutes : int, default 0
569
+ Number of minutes to wait prior to timing out.
570
+ hours : int, default 0
571
+ Number of hours to wait prior to timing out.
572
+ """
573
+ ...
574
+
575
+ @typing.overload
576
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
577
+ ...
578
+
579
+ @typing.overload
580
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
581
+ ...
582
+
583
+ 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):
584
+ """
585
+ Specifies a timeout for your step.
586
+
587
+ This decorator is useful if this step may hang indefinitely.
588
+
589
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
590
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
591
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
592
+
593
+ Note that all the values specified in parameters are added together so if you specify
594
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
595
+
596
+
597
+ Parameters
598
+ ----------
599
+ seconds : int, default 0
600
+ Number of seconds to wait prior to timing out.
601
+ minutes : int, default 0
602
+ Number of minutes to wait prior to timing out.
603
+ hours : int, default 0
604
+ Number of hours to wait prior to timing out.
605
+ """
606
+ ...
607
+
608
+ @typing.overload
609
+ 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]]]:
610
+ """
611
+ Specifies environment variables to be set prior to the execution of a step.
612
+
613
+
614
+ Parameters
615
+ ----------
616
+ vars : Dict[str, str], default {}
617
+ Dictionary of environment variables to set.
618
+ """
619
+ ...
620
+
621
+ @typing.overload
622
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
623
+ ...
624
+
625
+ @typing.overload
626
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
627
+ ...
628
+
629
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
630
+ """
631
+ Specifies environment variables to be set prior to the execution of a step.
632
+
633
+
634
+ Parameters
635
+ ----------
636
+ vars : Dict[str, str], default {}
637
+ Dictionary of environment variables to set.
860
638
  """
861
639
  ...
862
640
 
@@ -924,67 +702,131 @@ def coreweave_s3_proxy(*, integration_name: typing.Optional[str] = None, write_m
924
702
  """
925
703
  ...
926
704
 
927
- 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]]]:
705
+ 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]]]:
706
+ """
707
+ Specifies that this step should execute on DGX cloud.
708
+
709
+
710
+ Parameters
711
+ ----------
712
+ gpu : int
713
+ Number of GPUs to use.
714
+ gpu_type : str
715
+ Type of Nvidia GPU to use.
716
+ queue_timeout : int
717
+ Time to keep the job in NVCF's queue.
718
+ """
719
+ ...
720
+
721
+ @typing.overload
722
+ 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]]]:
723
+ """
724
+ Specifies the number of times the task corresponding
725
+ to a step needs to be retried.
726
+
727
+ This decorator is useful for handling transient errors, such as networking issues.
728
+ If your task contains operations that can't be retried safely, e.g. database updates,
729
+ it is advisable to annotate it with `@retry(times=0)`.
730
+
731
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
732
+ decorator will execute a no-op task after all retries have been exhausted,
733
+ ensuring that the flow execution can continue.
734
+
735
+
736
+ Parameters
737
+ ----------
738
+ times : int, default 3
739
+ Number of times to retry this task.
740
+ minutes_between_retries : int, default 2
741
+ Number of minutes between retries.
742
+ """
743
+ ...
744
+
745
+ @typing.overload
746
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
747
+ ...
748
+
749
+ @typing.overload
750
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
751
+ ...
752
+
753
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
928
754
  """
929
- Specifies that this step should execute on DGX cloud.
755
+ Specifies the number of times the task corresponding
756
+ to a step needs to be retried.
757
+
758
+ This decorator is useful for handling transient errors, such as networking issues.
759
+ If your task contains operations that can't be retried safely, e.g. database updates,
760
+ it is advisable to annotate it with `@retry(times=0)`.
761
+
762
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
763
+ decorator will execute a no-op task after all retries have been exhausted,
764
+ ensuring that the flow execution can continue.
930
765
 
931
766
 
932
767
  Parameters
933
768
  ----------
934
- gpu : int
935
- Number of GPUs to use.
936
- gpu_type : str
937
- Type of Nvidia GPU to use.
769
+ times : int, default 3
770
+ Number of times to retry this task.
771
+ minutes_between_retries : int, default 2
772
+ Number of minutes between retries.
938
773
  """
939
774
  ...
940
775
 
941
- 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]]]:
776
+ 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]]]:
942
777
  """
943
- This decorator is used to run vllm APIs as Metaflow task sidecars.
778
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
944
779
 
945
780
  User code call
946
781
  --------------
947
- @vllm(
948
- model="...",
782
+ @ollama(
783
+ models=[...],
949
784
  ...
950
785
  )
951
786
 
952
787
  Valid backend options
953
788
  ---------------------
954
789
  - 'local': Run as a separate process on the local task machine.
790
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
791
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
955
792
 
956
793
  Valid model options
957
794
  -------------------
958
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
959
-
960
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
961
- If you need multiple models, you must create multiple @vllm decorators.
795
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
962
796
 
963
797
 
964
798
  Parameters
965
799
  ----------
966
- model: str
967
- HuggingFace model identifier to be served by vLLM.
800
+ models: list[str]
801
+ List of Ollama containers running models in sidecars.
968
802
  backend: str
969
- Determines where and how to run the vLLM process.
970
- openai_api_server: bool
971
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
972
- Default is False (uses native engine).
973
- Set to True for backward compatibility with existing code.
803
+ Determines where and how to run the Ollama process.
804
+ force_pull: bool
805
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
806
+ cache_update_policy: str
807
+ Cache update policy: "auto", "force", or "never".
808
+ force_cache_update: bool
809
+ Simple override for "force" cache update policy.
974
810
  debug: bool
975
811
  Whether to turn on verbose debugging logs.
976
- card_refresh_interval: int
977
- Interval in seconds for refreshing the vLLM status card.
978
- Only used when openai_api_server=True.
979
- max_retries: int
980
- Maximum number of retries checking for vLLM server startup.
981
- Only used when openai_api_server=True.
982
- retry_alert_frequency: int
983
- Frequency of alert logs for vLLM server startup retries.
984
- Only used when openai_api_server=True.
985
- engine_args : dict
986
- Additional keyword arguments to pass to the vLLM engine.
987
- For example, `tensor_parallel_size=2`.
812
+ circuit_breaker_config: dict
813
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
814
+ timeout_config: dict
815
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
816
+ """
817
+ ...
818
+
819
+ 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]]]:
820
+ """
821
+ Specifies that this step should execute on DGX cloud.
822
+
823
+
824
+ Parameters
825
+ ----------
826
+ gpu : int
827
+ Number of GPUs to use.
828
+ gpu_type : str
829
+ Type of Nvidia GPU to use.
988
830
  """
989
831
  ...
990
832
 
@@ -1007,132 +849,199 @@ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
1007
849
  """
1008
850
  ...
1009
851
 
1010
- @typing.overload
1011
- 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]]]:
852
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, cache_scope: 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]]]:
1012
853
  """
1013
- Enables loading / saving of models within a step.
854
+ Decorator that helps cache, version, and store models/datasets from the Hugging Face Hub.
855
+
856
+ Examples
857
+ --------
1014
858
 
1015
- > Examples
1016
- - Saving Models
1017
859
  ```python
1018
- @model
860
+ # **Usage: creating references to models from the Hugging Face Hub that may be loaded in downstream steps**
861
+ @huggingface_hub
1019
862
  @step
1020
- def train(self):
1021
- # current.model.save returns a dictionary reference to the model saved
1022
- self.my_model = current.model.save(
1023
- path_to_my_model,
1024
- label="my_model",
1025
- metadata={
1026
- "epochs": 10,
1027
- "batch-size": 32,
1028
- "learning-rate": 0.001,
1029
- }
863
+ def pull_model_from_huggingface(self):
864
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
865
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
866
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
867
+ # value of the function is a reference to the model in the backend storage.
868
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
869
+
870
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
871
+ self.llama_model = current.huggingface_hub.snapshot_download(
872
+ repo_id=self.model_id,
873
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1030
874
  )
1031
- self.next(self.test)
875
+ self.next(self.train)
1032
876
 
1033
- @model(load="my_model")
877
+ # **Usage: explicitly loading models at runtime from the Hugging Face Hub or from cache (from Metaflow's datastore)**
878
+ @huggingface_hub
1034
879
  @step
1035
- def test(self):
1036
- # `current.model.loaded` returns a dictionary of the loaded models
1037
- # where the key is the name of the artifact and the value is the path to the model
1038
- print(os.listdir(current.model.loaded["my_model"]))
1039
- self.next(self.end)
1040
- ```
880
+ def run_training(self):
881
+ # Temporary directory (auto-cleaned on exit)
882
+ with current.huggingface_hub.load(
883
+ repo_id="google-bert/bert-base-uncased",
884
+ allow_patterns=["*.bin"],
885
+ ) as local_path:
886
+ # Use files under local_path
887
+ train_model(local_path)
888
+ ...
1041
889
 
1042
- - Loading models
1043
- ```python
890
+ # **Usage: loading models directly from the Hugging Face Hub or from cache (from Metaflow's datastore)**
891
+
892
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1044
893
  @step
1045
- def train(self):
1046
- # current.model.load returns the path to the model loaded
1047
- checkpoint_path = current.model.load(
1048
- self.checkpoint_key,
1049
- )
1050
- model_path = current.model.load(
1051
- self.model,
1052
- )
1053
- self.next(self.test)
894
+ def pull_model_from_huggingface(self):
895
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
896
+
897
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora", "/my-lora-directory")])
898
+ @step
899
+ def finetune_model(self):
900
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
901
+ # path_to_model will be /my-directory
902
+
903
+
904
+ # Takes all the arguments passed to `snapshot_download`
905
+ # except for `local_dir`
906
+ @huggingface_hub(load=[
907
+ {
908
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
909
+ },
910
+ {
911
+ "repo_id": "myorg/mistral-lora",
912
+ "repo_type": "model",
913
+ },
914
+ ])
915
+ @step
916
+ def finetune_model(self):
917
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
918
+ # path_to_model will be /my-directory
919
+ ```
920
+
921
+
922
+ Parameters
923
+ ----------
924
+ temp_dir_root : str, optional
925
+ The root directory that will hold the temporary directory where objects will be downloaded.
926
+
927
+ cache_scope : str, optional
928
+ The scope of the cache. Can be `checkpoint` / `flow` / `global`.
929
+ - `checkpoint` (default): All repos are stored like objects saved by `@checkpoint`.
930
+ i.e., the cached path is derived from the namespace, flow, step, and Metaflow foreach iteration.
931
+ Any repo downloaded under this scope will only be retrieved from the cache when the step runs under the same namespace in the same flow (at the same foreach index).
932
+
933
+ - `flow`: All repos are cached under the flow, regardless of namespace.
934
+ i.e., the cached path is derived solely from the flow name.
935
+ When to use this mode: (1) Multiple users are executing the same flow and want shared access to the repos cached by the decorator. (2) Multiple versions of a flow are deployed, all needing access to the same repos cached by the decorator.
936
+
937
+ - `global`: All repos are cached under a globally static path.
938
+ i.e., the base path of the cache is static and all repos are stored under it.
939
+ When to use this mode:
940
+ - All repos from the Hugging Face Hub need to be shared by users across all flow executions.
941
+ - Each caching scope comes with its own trade-offs:
942
+ - `checkpoint`:
943
+ - Has explicit control over when caches are populated (controlled by the same flow that has the `@huggingface_hub` decorator) but ends up hitting the Hugging Face Hub more often if there are many users/namespaces/steps.
944
+ - Since objects are written on a `namespace/flow/step` basis, the blast radius of a bad checkpoint is limited to a particular flow in a namespace.
945
+ - `flow`:
946
+ - Has less control over when caches are populated (can be written by any execution instance of a flow from any namespace) but results in more cache hits.
947
+ - The blast radius of a bad checkpoint is limited to all runs of a particular flow.
948
+ - It doesn't promote cache reuse across flows.
949
+ - `global`:
950
+ - Has no control over when caches are populated (can be written by any flow execution) but has the highest cache hit rate.
951
+ - It promotes cache reuse across flows.
952
+ - The blast radius of a bad checkpoint spans every flow that could be using a particular repo.
953
+
954
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
955
+ The list of repos (models/datasets) to load.
956
+
957
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
958
+
959
+ - If repo (model/dataset) is not found in the datastore:
960
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
961
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
962
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
963
+
964
+ - If repo is found in the datastore:
965
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
966
+ """
967
+ ...
968
+
969
+ @typing.overload
970
+ 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]]]:
971
+ """
972
+ Specifies the resources needed when executing this step.
973
+
974
+ Use `@resources` to specify the resource requirements
975
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
976
+
977
+ You can choose the compute layer on the command line by executing e.g.
1054
978
  ```
979
+ python myflow.py run --with batch
980
+ ```
981
+ or
982
+ ```
983
+ python myflow.py run --with kubernetes
984
+ ```
985
+ which executes the flow on the desired system using the
986
+ requirements specified in `@resources`.
1055
987
 
1056
988
 
1057
989
  Parameters
1058
990
  ----------
1059
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1060
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1061
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by `current.checkpoint` / `current.model` / `current.huggingface_hub`.
1062
- If a list of tuples is provided, the first element is the artifact name and the second element is the path the artifact needs be unpacked on
1063
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1064
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1065
-
1066
- temp_dir_root : str, default: None
1067
- The root directory under which `current.model.loaded` will store loaded models
991
+ cpu : int, default 1
992
+ Number of CPUs required for this step.
993
+ gpu : int, optional, default None
994
+ Number of GPUs required for this step.
995
+ disk : int, optional, default None
996
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
997
+ memory : int, default 4096
998
+ Memory size (in MB) required for this step.
999
+ shared_memory : int, optional, default None
1000
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1001
+ This parameter maps to the `--shm-size` option in Docker.
1068
1002
  """
1069
1003
  ...
1070
1004
 
1071
1005
  @typing.overload
1072
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1006
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1073
1007
  ...
1074
1008
 
1075
1009
  @typing.overload
1076
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1010
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1077
1011
  ...
1078
1012
 
1079
- 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):
1013
+ 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):
1080
1014
  """
1081
- Enables loading / saving of models within a step.
1015
+ Specifies the resources needed when executing this step.
1082
1016
 
1083
- > Examples
1084
- - Saving Models
1085
- ```python
1086
- @model
1087
- @step
1088
- def train(self):
1089
- # current.model.save returns a dictionary reference to the model saved
1090
- self.my_model = current.model.save(
1091
- path_to_my_model,
1092
- label="my_model",
1093
- metadata={
1094
- "epochs": 10,
1095
- "batch-size": 32,
1096
- "learning-rate": 0.001,
1097
- }
1098
- )
1099
- self.next(self.test)
1017
+ Use `@resources` to specify the resource requirements
1018
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1100
1019
 
1101
- @model(load="my_model")
1102
- @step
1103
- def test(self):
1104
- # `current.model.loaded` returns a dictionary of the loaded models
1105
- # where the key is the name of the artifact and the value is the path to the model
1106
- print(os.listdir(current.model.loaded["my_model"]))
1107
- self.next(self.end)
1020
+ You can choose the compute layer on the command line by executing e.g.
1108
1021
  ```
1109
-
1110
- - Loading models
1111
- ```python
1112
- @step
1113
- def train(self):
1114
- # current.model.load returns the path to the model loaded
1115
- checkpoint_path = current.model.load(
1116
- self.checkpoint_key,
1117
- )
1118
- model_path = current.model.load(
1119
- self.model,
1120
- )
1121
- self.next(self.test)
1022
+ python myflow.py run --with batch
1023
+ ```
1024
+ or
1025
+ ```
1026
+ python myflow.py run --with kubernetes
1122
1027
  ```
1028
+ which executes the flow on the desired system using the
1029
+ requirements specified in `@resources`.
1123
1030
 
1124
1031
 
1125
1032
  Parameters
1126
1033
  ----------
1127
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1128
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1129
- 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`.
1130
- 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
1131
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1132
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1133
-
1134
- temp_dir_root : str, default: None
1135
- The root directory under which `current.model.loaded` will store loaded models
1034
+ cpu : int, default 1
1035
+ Number of CPUs required for this step.
1036
+ gpu : int, optional, default None
1037
+ Number of GPUs required for this step.
1038
+ disk : int, optional, default None
1039
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1040
+ memory : int, default 4096
1041
+ Memory size (in MB) required for this step.
1042
+ shared_memory : int, optional, default None
1043
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1044
+ This parameter maps to the `--shm-size` option in Docker.
1136
1045
  """
1137
1046
  ...
1138
1047
 
@@ -1226,287 +1135,479 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
1226
1135
  ...
1227
1136
 
1228
1137
  @typing.overload
1229
- 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]]]:
1138
+ 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]]]:
1230
1139
  """
1231
- Specifies a timeout for your step.
1232
-
1233
- This decorator is useful if this step may hang indefinitely.
1234
-
1235
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1236
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1237
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1140
+ Specifies that the step will success under all circumstances.
1238
1141
 
1239
- Note that all the values specified in parameters are added together so if you specify
1240
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1142
+ The decorator will create an optional artifact, specified by `var`, which
1143
+ contains the exception raised. You can use it to detect the presence
1144
+ of errors, indicating that all happy-path artifacts produced by the step
1145
+ are missing.
1241
1146
 
1242
1147
 
1243
1148
  Parameters
1244
1149
  ----------
1245
- seconds : int, default 0
1246
- Number of seconds to wait prior to timing out.
1247
- minutes : int, default 0
1248
- Number of minutes to wait prior to timing out.
1249
- hours : int, default 0
1250
- Number of hours to wait prior to timing out.
1150
+ var : str, optional, default None
1151
+ Name of the artifact in which to store the caught exception.
1152
+ If not specified, the exception is not stored.
1153
+ print_exception : bool, default True
1154
+ Determines whether or not the exception is printed to
1155
+ stdout when caught.
1251
1156
  """
1252
1157
  ...
1253
1158
 
1254
1159
  @typing.overload
1255
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1160
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1256
1161
  ...
1257
1162
 
1258
1163
  @typing.overload
1259
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1164
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1260
1165
  ...
1261
1166
 
1262
- 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):
1167
+ 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):
1263
1168
  """
1264
- Specifies a timeout for your step.
1169
+ Specifies that the step will success under all circumstances.
1265
1170
 
1266
- This decorator is useful if this step may hang indefinitely.
1171
+ The decorator will create an optional artifact, specified by `var`, which
1172
+ contains the exception raised. You can use it to detect the presence
1173
+ of errors, indicating that all happy-path artifacts produced by the step
1174
+ are missing.
1267
1175
 
1268
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1269
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1270
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1271
1176
 
1272
- Note that all the values specified in parameters are added together so if you specify
1273
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1177
+ Parameters
1178
+ ----------
1179
+ var : str, optional, default None
1180
+ Name of the artifact in which to store the caught exception.
1181
+ If not specified, the exception is not stored.
1182
+ print_exception : bool, default True
1183
+ Determines whether or not the exception is printed to
1184
+ stdout when caught.
1185
+ """
1186
+ ...
1187
+
1188
+ 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]]]:
1189
+ """
1190
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
1191
+
1192
+ User code call
1193
+ --------------
1194
+ @vllm(
1195
+ model="...",
1196
+ ...
1197
+ )
1198
+
1199
+ Valid backend options
1200
+ ---------------------
1201
+ - 'local': Run as a separate process on the local task machine.
1202
+
1203
+ Valid model options
1204
+ -------------------
1205
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1206
+
1207
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
1208
+ If you need multiple models, you must create multiple @vllm decorators.
1274
1209
 
1275
1210
 
1276
1211
  Parameters
1277
1212
  ----------
1278
- seconds : int, default 0
1279
- Number of seconds to wait prior to timing out.
1280
- minutes : int, default 0
1281
- Number of minutes to wait prior to timing out.
1282
- hours : int, default 0
1283
- Number of hours to wait prior to timing out.
1213
+ model: str
1214
+ HuggingFace model identifier to be served by vLLM.
1215
+ backend: str
1216
+ Determines where and how to run the vLLM process.
1217
+ openai_api_server: bool
1218
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
1219
+ Default is False (uses native engine).
1220
+ Set to True for backward compatibility with existing code.
1221
+ debug: bool
1222
+ Whether to turn on verbose debugging logs.
1223
+ card_refresh_interval: int
1224
+ Interval in seconds for refreshing the vLLM status card.
1225
+ Only used when openai_api_server=True.
1226
+ max_retries: int
1227
+ Maximum number of retries checking for vLLM server startup.
1228
+ Only used when openai_api_server=True.
1229
+ retry_alert_frequency: int
1230
+ Frequency of alert logs for vLLM server startup retries.
1231
+ Only used when openai_api_server=True.
1232
+ engine_args : dict
1233
+ Additional keyword arguments to pass to the vLLM engine.
1234
+ For example, `tensor_parallel_size=2`.
1284
1235
  """
1285
1236
  ...
1286
1237
 
1287
1238
  @typing.overload
1288
- 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]]]:
1239
+ 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]]]:
1289
1240
  """
1290
- Specifies that the step will success under all circumstances.
1241
+ Enables loading / saving of models within a step.
1242
+
1243
+ > Examples
1244
+ - Saving Models
1245
+ ```python
1246
+ @model
1247
+ @step
1248
+ def train(self):
1249
+ # current.model.save returns a dictionary reference to the model saved
1250
+ self.my_model = current.model.save(
1251
+ path_to_my_model,
1252
+ label="my_model",
1253
+ metadata={
1254
+ "epochs": 10,
1255
+ "batch-size": 32,
1256
+ "learning-rate": 0.001,
1257
+ }
1258
+ )
1259
+ self.next(self.test)
1260
+
1261
+ @model(load="my_model")
1262
+ @step
1263
+ def test(self):
1264
+ # `current.model.loaded` returns a dictionary of the loaded models
1265
+ # where the key is the name of the artifact and the value is the path to the model
1266
+ print(os.listdir(current.model.loaded["my_model"]))
1267
+ self.next(self.end)
1268
+ ```
1269
+
1270
+ - Loading models
1271
+ ```python
1272
+ @step
1273
+ def train(self):
1274
+ # current.model.load returns the path to the model loaded
1275
+ checkpoint_path = current.model.load(
1276
+ self.checkpoint_key,
1277
+ )
1278
+ model_path = current.model.load(
1279
+ self.model,
1280
+ )
1281
+ self.next(self.test)
1282
+ ```
1291
1283
 
1292
- The decorator will create an optional artifact, specified by `var`, which
1293
- contains the exception raised. You can use it to detect the presence
1294
- of errors, indicating that all happy-path artifacts produced by the step
1295
- are missing.
1296
1284
 
1285
+ Parameters
1286
+ ----------
1287
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1288
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1289
+ 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`.
1290
+ 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
1291
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1292
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1297
1293
 
1298
- Parameters
1299
- ----------
1300
- var : str, optional, default None
1301
- Name of the artifact in which to store the caught exception.
1302
- If not specified, the exception is not stored.
1303
- print_exception : bool, default True
1304
- Determines whether or not the exception is printed to
1305
- stdout when caught.
1294
+ temp_dir_root : str, default: None
1295
+ The root directory under which `current.model.loaded` will store loaded models
1306
1296
  """
1307
1297
  ...
1308
1298
 
1309
1299
  @typing.overload
1310
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1300
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1311
1301
  ...
1312
1302
 
1313
1303
  @typing.overload
1314
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1304
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1315
1305
  ...
1316
1306
 
1317
- 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):
1307
+ 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):
1318
1308
  """
1319
- Specifies that the step will success under all circumstances.
1309
+ Enables loading / saving of models within a step.
1320
1310
 
1321
- The decorator will create an optional artifact, specified by `var`, which
1322
- contains the exception raised. You can use it to detect the presence
1323
- of errors, indicating that all happy-path artifacts produced by the step
1324
- are missing.
1311
+ > Examples
1312
+ - Saving Models
1313
+ ```python
1314
+ @model
1315
+ @step
1316
+ def train(self):
1317
+ # current.model.save returns a dictionary reference to the model saved
1318
+ self.my_model = current.model.save(
1319
+ path_to_my_model,
1320
+ label="my_model",
1321
+ metadata={
1322
+ "epochs": 10,
1323
+ "batch-size": 32,
1324
+ "learning-rate": 0.001,
1325
+ }
1326
+ )
1327
+ self.next(self.test)
1328
+
1329
+ @model(load="my_model")
1330
+ @step
1331
+ def test(self):
1332
+ # `current.model.loaded` returns a dictionary of the loaded models
1333
+ # where the key is the name of the artifact and the value is the path to the model
1334
+ print(os.listdir(current.model.loaded["my_model"]))
1335
+ self.next(self.end)
1336
+ ```
1337
+
1338
+ - Loading models
1339
+ ```python
1340
+ @step
1341
+ def train(self):
1342
+ # current.model.load returns the path to the model loaded
1343
+ checkpoint_path = current.model.load(
1344
+ self.checkpoint_key,
1345
+ )
1346
+ model_path = current.model.load(
1347
+ self.model,
1348
+ )
1349
+ self.next(self.test)
1350
+ ```
1325
1351
 
1326
1352
 
1327
1353
  Parameters
1328
1354
  ----------
1329
- var : str, optional, default None
1330
- Name of the artifact in which to store the caught exception.
1331
- If not specified, the exception is not stored.
1332
- print_exception : bool, default True
1333
- Determines whether or not the exception is printed to
1334
- stdout when caught.
1355
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
1356
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
1357
+ 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`.
1358
+ 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
1359
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
1360
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1361
+
1362
+ temp_dir_root : str, default: None
1363
+ The root directory under which `current.model.loaded` will store loaded models
1335
1364
  """
1336
1365
  ...
1337
1366
 
1338
1367
  @typing.overload
1339
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1368
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1340
1369
  """
1341
- Internal decorator to support Fast bakery
1370
+ Decorator prototype for all step decorators. This function gets specialized
1371
+ and imported for all decorators types by _import_plugin_decorators().
1342
1372
  """
1343
1373
  ...
1344
1374
 
1345
1375
  @typing.overload
1346
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1376
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1347
1377
  ...
1348
1378
 
1349
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1379
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1350
1380
  """
1351
- Internal decorator to support Fast bakery
1381
+ Decorator prototype for all step decorators. This function gets specialized
1382
+ and imported for all decorators types by _import_plugin_decorators().
1352
1383
  """
1353
1384
  ...
1354
1385
 
1355
1386
  @typing.overload
1356
- 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]]]:
1387
+ 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]]]:
1357
1388
  """
1358
- Specifies environment variables to be set prior to the execution of a step.
1389
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1390
+
1391
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1359
1392
 
1360
1393
 
1361
1394
  Parameters
1362
1395
  ----------
1363
- vars : Dict[str, str], default {}
1364
- Dictionary of environment variables to set.
1396
+ type : str, default 'default'
1397
+ Card type.
1398
+ id : str, optional, default None
1399
+ If multiple cards are present, use this id to identify this card.
1400
+ options : Dict[str, Any], default {}
1401
+ Options passed to the card. The contents depend on the card type.
1402
+ timeout : int, default 45
1403
+ Interrupt reporting if it takes more than this many seconds.
1365
1404
  """
1366
1405
  ...
1367
1406
 
1368
1407
  @typing.overload
1369
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1408
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1370
1409
  ...
1371
1410
 
1372
1411
  @typing.overload
1373
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1412
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1374
1413
  ...
1375
1414
 
1376
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
1415
+ 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):
1377
1416
  """
1378
- Specifies environment variables to be set prior to the execution of a step.
1417
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1418
+
1419
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1379
1420
 
1380
1421
 
1381
1422
  Parameters
1382
1423
  ----------
1383
- vars : Dict[str, str], default {}
1384
- Dictionary of environment variables to set.
1424
+ type : str, default 'default'
1425
+ Card type.
1426
+ id : str, optional, default None
1427
+ If multiple cards are present, use this id to identify this card.
1428
+ options : Dict[str, Any], default {}
1429
+ Options passed to the card. The contents depend on the card type.
1430
+ timeout : int, default 45
1431
+ Interrupt reporting if it takes more than this many seconds.
1385
1432
  """
1386
1433
  ...
1387
1434
 
1388
1435
  @typing.overload
1389
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1436
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1390
1437
  """
1391
- Specifies the PyPI packages for the step.
1392
-
1393
- Information in this decorator will augment any
1394
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1395
- you can use `@pypi_base` to set packages required by all
1396
- steps and use `@pypi` to specify step-specific overrides.
1397
-
1398
-
1399
- Parameters
1400
- ----------
1401
- packages : Dict[str, str], default: {}
1402
- Packages to use for this step. The key is the name of the package
1403
- and the value is the version to use.
1404
- python : str, optional, default: None
1405
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1406
- that the version used will correspond to the version of the Python interpreter used to start the run.
1438
+ A simple decorator that demonstrates using CardDecoratorInjector
1439
+ to inject a card and render simple markdown content.
1407
1440
  """
1408
1441
  ...
1409
1442
 
1410
1443
  @typing.overload
1411
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1444
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1412
1445
  ...
1413
1446
 
1414
- @typing.overload
1415
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1447
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1448
+ """
1449
+ A simple decorator that demonstrates using CardDecoratorInjector
1450
+ to inject a card and render simple markdown content.
1451
+ """
1416
1452
  ...
1417
1453
 
1418
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1454
+ 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]]]:
1419
1455
  """
1420
- Specifies the PyPI packages for the step.
1456
+ Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
1457
+ for S3 read and write requests.
1421
1458
 
1422
- Information in this decorator will augment any
1423
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1424
- you can use `@pypi_base` to set packages required by all
1425
- steps and use `@pypi` to specify step-specific overrides.
1459
+ This decorator requires an integration in the Outerbounds platform that
1460
+ points to an external bucket. It affects S3 operations performed via
1461
+ Metaflow's `get_aws_client` and `S3` within a `@step`.
1462
+
1463
+ Read operations
1464
+ ---------------
1465
+ All read operations pass through the proxy. If an object does not already
1466
+ exist in the external bucket, it is cached there. For example, if code reads
1467
+ from buckets `FOO` and `BAR` using the `S3` interface, objects from both
1468
+ buckets are cached in the external bucket.
1469
+
1470
+ During task execution, all S3‑related read requests are routed through the
1471
+ proxy:
1472
+ - If the object is present in the external object store, the proxy
1473
+ streams it directly from there without accessing the requested origin
1474
+ bucket.
1475
+ - If the object is not present in the external storage, the proxy
1476
+ fetches it from the requested bucket, caches it in the external
1477
+ storage, and streams the response from the origin bucket.
1478
+
1479
+ Warning
1480
+ -------
1481
+ All READ operations (e.g., GetObject, HeadObject) pass through the external
1482
+ bucket regardless of the bucket specified in user code. Even
1483
+ `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
1484
+ external bucket cache.
1485
+
1486
+ Write operations
1487
+ ----------------
1488
+ Write behavior is controlled by the `write_mode` parameter, which determines
1489
+ whether writes also persist objects in the cache.
1490
+
1491
+ `write_mode` values:
1492
+ - `origin-and-cache`: objects are written both to the cache and to their
1493
+ intended origin bucket.
1494
+ - `origin`: objects are written only to their intended origin bucket.
1426
1495
 
1427
1496
 
1428
1497
  Parameters
1429
1498
  ----------
1430
- packages : Dict[str, str], default: {}
1431
- Packages to use for this step. The key is the name of the package
1432
- and the value is the version to use.
1433
- python : str, optional, default: None
1434
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1435
- that the version used will correspond to the version of the Python interpreter used to start the run.
1499
+ integration_name : str, optional
1500
+ [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
1501
+ that holds the configuration for the external, S3‑compatible object
1502
+ storage bucket. If not specified, the only available S3 proxy
1503
+ integration in the namespace is used (fails if multiple exist).
1504
+ write_mode : str, optional
1505
+ Controls whether writes also go to the external bucket.
1506
+ - `origin` (default)
1507
+ - `origin-and-cache`
1508
+ debug : bool, optional
1509
+ Enables debug logging for proxy operations.
1436
1510
  """
1437
1511
  ...
1438
1512
 
1439
1513
  @typing.overload
1440
- 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]]]:
1514
+ 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]]:
1441
1515
  """
1442
- Specifies the number of times the task corresponding
1443
- to a step needs to be retried.
1516
+ Specifies the flow(s) that this flow depends on.
1444
1517
 
1445
- This decorator is useful for handling transient errors, such as networking issues.
1446
- If your task contains operations that can't be retried safely, e.g. database updates,
1447
- it is advisable to annotate it with `@retry(times=0)`.
1518
+ ```
1519
+ @trigger_on_finish(flow='FooFlow')
1520
+ ```
1521
+ or
1522
+ ```
1523
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1524
+ ```
1525
+ This decorator respects the @project decorator and triggers the flow
1526
+ when upstream runs within the same namespace complete successfully
1448
1527
 
1449
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1450
- decorator will execute a no-op task after all retries have been exhausted,
1451
- ensuring that the flow execution can continue.
1528
+ Additionally, you can specify project aware upstream flow dependencies
1529
+ by specifying the fully qualified project_flow_name.
1530
+ ```
1531
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1532
+ ```
1533
+ or
1534
+ ```
1535
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1536
+ ```
1537
+
1538
+ You can also specify just the project or project branch (other values will be
1539
+ inferred from the current project or project branch):
1540
+ ```
1541
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1542
+ ```
1543
+
1544
+ Note that `branch` is typically one of:
1545
+ - `prod`
1546
+ - `user.bob`
1547
+ - `test.my_experiment`
1548
+ - `prod.staging`
1452
1549
 
1453
1550
 
1454
1551
  Parameters
1455
1552
  ----------
1456
- times : int, default 3
1457
- Number of times to retry this task.
1458
- minutes_between_retries : int, default 2
1459
- Number of minutes between retries.
1553
+ flow : Union[str, Dict[str, str]], optional, default None
1554
+ Upstream flow dependency for this flow.
1555
+ flows : List[Union[str, Dict[str, str]]], default []
1556
+ Upstream flow dependencies for this flow.
1557
+ options : Dict[str, Any], default {}
1558
+ Backend-specific configuration for tuning eventing behavior.
1460
1559
  """
1461
1560
  ...
1462
1561
 
1463
1562
  @typing.overload
1464
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1465
- ...
1466
-
1467
- @typing.overload
1468
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1563
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1469
1564
  ...
1470
1565
 
1471
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
1566
+ 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] = {}):
1472
1567
  """
1473
- Specifies the number of times the task corresponding
1474
- to a step needs to be retried.
1568
+ Specifies the flow(s) that this flow depends on.
1475
1569
 
1476
- This decorator is useful for handling transient errors, such as networking issues.
1477
- If your task contains operations that can't be retried safely, e.g. database updates,
1478
- it is advisable to annotate it with `@retry(times=0)`.
1570
+ ```
1571
+ @trigger_on_finish(flow='FooFlow')
1572
+ ```
1573
+ or
1574
+ ```
1575
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1576
+ ```
1577
+ This decorator respects the @project decorator and triggers the flow
1578
+ when upstream runs within the same namespace complete successfully
1479
1579
 
1480
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1481
- decorator will execute a no-op task after all retries have been exhausted,
1482
- ensuring that the flow execution can continue.
1580
+ Additionally, you can specify project aware upstream flow dependencies
1581
+ by specifying the fully qualified project_flow_name.
1582
+ ```
1583
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1584
+ ```
1585
+ or
1586
+ ```
1587
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1588
+ ```
1589
+
1590
+ You can also specify just the project or project branch (other values will be
1591
+ inferred from the current project or project branch):
1592
+ ```
1593
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1594
+ ```
1595
+
1596
+ Note that `branch` is typically one of:
1597
+ - `prod`
1598
+ - `user.bob`
1599
+ - `test.my_experiment`
1600
+ - `prod.staging`
1483
1601
 
1484
1602
 
1485
1603
  Parameters
1486
1604
  ----------
1487
- times : int, default 3
1488
- Number of times to retry this task.
1489
- minutes_between_retries : int, default 2
1490
- Number of minutes between retries.
1491
- """
1492
- ...
1493
-
1494
- @typing.overload
1495
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1496
- """
1497
- Decorator prototype for all step decorators. This function gets specialized
1498
- and imported for all decorators types by _import_plugin_decorators().
1499
- """
1500
- ...
1501
-
1502
- @typing.overload
1503
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1504
- ...
1505
-
1506
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1507
- """
1508
- Decorator prototype for all step decorators. This function gets specialized
1509
- and imported for all decorators types by _import_plugin_decorators().
1605
+ flow : Union[str, Dict[str, str]], optional, default None
1606
+ Upstream flow dependency for this flow.
1607
+ flows : List[Union[str, Dict[str, str]]], default []
1608
+ Upstream flow dependencies for this flow.
1609
+ options : Dict[str, Any], default {}
1610
+ Backend-specific configuration for tuning eventing behavior.
1510
1611
  """
1511
1612
  ...
1512
1613
 
@@ -1625,135 +1726,143 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1625
1726
  ...
1626
1727
 
1627
1728
  @typing.overload
1628
- 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]]:
1729
+ 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]]:
1629
1730
  """
1630
- Specifies the flow(s) that this flow depends on.
1731
+ Specifies the event(s) that this flow depends on.
1631
1732
 
1632
1733
  ```
1633
- @trigger_on_finish(flow='FooFlow')
1734
+ @trigger(event='foo')
1634
1735
  ```
1635
1736
  or
1636
1737
  ```
1637
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1738
+ @trigger(events=['foo', 'bar'])
1638
1739
  ```
1639
- This decorator respects the @project decorator and triggers the flow
1640
- when upstream runs within the same namespace complete successfully
1641
1740
 
1642
- Additionally, you can specify project aware upstream flow dependencies
1643
- by specifying the fully qualified project_flow_name.
1741
+ Additionally, you can specify the parameter mappings
1742
+ to map event payload to Metaflow parameters for the flow.
1644
1743
  ```
1645
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1744
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1646
1745
  ```
1647
1746
  or
1648
1747
  ```
1649
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1748
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1749
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1650
1750
  ```
1651
1751
 
1652
- You can also specify just the project or project branch (other values will be
1653
- inferred from the current project or project branch):
1752
+ 'parameters' can also be a list of strings and tuples like so:
1654
1753
  ```
1655
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1754
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1755
+ ```
1756
+ This is equivalent to:
1757
+ ```
1758
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1656
1759
  ```
1657
-
1658
- Note that `branch` is typically one of:
1659
- - `prod`
1660
- - `user.bob`
1661
- - `test.my_experiment`
1662
- - `prod.staging`
1663
1760
 
1664
1761
 
1665
1762
  Parameters
1666
1763
  ----------
1667
- flow : Union[str, Dict[str, str]], optional, default None
1668
- Upstream flow dependency for this flow.
1669
- flows : List[Union[str, Dict[str, str]]], default []
1670
- Upstream flow dependencies for this flow.
1764
+ event : Union[str, Dict[str, Any]], optional, default None
1765
+ Event dependency for this flow.
1766
+ events : List[Union[str, Dict[str, Any]]], default []
1767
+ Events dependency for this flow.
1671
1768
  options : Dict[str, Any], default {}
1672
1769
  Backend-specific configuration for tuning eventing behavior.
1673
1770
  """
1674
1771
  ...
1675
1772
 
1676
1773
  @typing.overload
1677
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1774
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1678
1775
  ...
1679
1776
 
1680
- 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] = {}):
1777
+ 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] = {}):
1681
1778
  """
1682
- Specifies the flow(s) that this flow depends on.
1779
+ Specifies the event(s) that this flow depends on.
1683
1780
 
1684
1781
  ```
1685
- @trigger_on_finish(flow='FooFlow')
1782
+ @trigger(event='foo')
1686
1783
  ```
1687
1784
  or
1688
1785
  ```
1689
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1786
+ @trigger(events=['foo', 'bar'])
1690
1787
  ```
1691
- This decorator respects the @project decorator and triggers the flow
1692
- when upstream runs within the same namespace complete successfully
1693
1788
 
1694
- Additionally, you can specify project aware upstream flow dependencies
1695
- by specifying the fully qualified project_flow_name.
1789
+ Additionally, you can specify the parameter mappings
1790
+ to map event payload to Metaflow parameters for the flow.
1696
1791
  ```
1697
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1792
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1698
1793
  ```
1699
1794
  or
1700
1795
  ```
1701
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1796
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1797
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1702
1798
  ```
1703
1799
 
1704
- You can also specify just the project or project branch (other values will be
1705
- inferred from the current project or project branch):
1800
+ 'parameters' can also be a list of strings and tuples like so:
1706
1801
  ```
1707
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1802
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1803
+ ```
1804
+ This is equivalent to:
1805
+ ```
1806
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1708
1807
  ```
1709
-
1710
- Note that `branch` is typically one of:
1711
- - `prod`
1712
- - `user.bob`
1713
- - `test.my_experiment`
1714
- - `prod.staging`
1715
1808
 
1716
1809
 
1717
1810
  Parameters
1718
1811
  ----------
1719
- flow : Union[str, Dict[str, str]], optional, default None
1720
- Upstream flow dependency for this flow.
1721
- flows : List[Union[str, Dict[str, str]]], default []
1722
- Upstream flow dependencies for this flow.
1812
+ event : Union[str, Dict[str, Any]], optional, default None
1813
+ Event dependency for this flow.
1814
+ events : List[Union[str, Dict[str, Any]]], default []
1815
+ Events dependency for this flow.
1723
1816
  options : Dict[str, Any], default {}
1724
1817
  Backend-specific configuration for tuning eventing behavior.
1725
1818
  """
1726
1819
  ...
1727
1820
 
1728
- @typing.overload
1729
- 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]]:
1821
+ 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]]:
1730
1822
  """
1731
- Specifies the times when the flow should be run when running on a
1732
- production scheduler.
1823
+ 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)
1824
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1825
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1826
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1827
+ starts only after all sensors finish.
1733
1828
 
1734
1829
 
1735
1830
  Parameters
1736
1831
  ----------
1737
- hourly : bool, default False
1738
- Run the workflow hourly.
1739
- daily : bool, default True
1740
- Run the workflow daily.
1741
- weekly : bool, default False
1742
- Run the workflow weekly.
1743
- cron : str, optional, default None
1744
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1745
- specified by this expression.
1746
- timezone : str, optional, default None
1747
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1748
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1832
+ timeout : int
1833
+ Time, in seconds before the task times out and fails. (Default: 3600)
1834
+ poke_interval : int
1835
+ Time in seconds that the job should wait in between each try. (Default: 60)
1836
+ mode : str
1837
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1838
+ exponential_backoff : bool
1839
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1840
+ pool : str
1841
+ the slot pool this task should run in,
1842
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1843
+ soft_fail : bool
1844
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1845
+ name : str
1846
+ Name of the sensor on Airflow
1847
+ description : str
1848
+ Description of sensor in the Airflow UI
1849
+ bucket_key : Union[str, List[str]]
1850
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1851
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1852
+ bucket_name : str
1853
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1854
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1855
+ wildcard_match : bool
1856
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1857
+ aws_conn_id : str
1858
+ a reference to the s3 connection on Airflow. (Default: None)
1859
+ verify : bool
1860
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1749
1861
  """
1750
1862
  ...
1751
1863
 
1752
1864
  @typing.overload
1753
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1754
- ...
1755
-
1756
- 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):
1865
+ 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]]:
1757
1866
  """
1758
1867
  Specifies the times when the flow should be run when running on a
1759
1868
  production scheduler.
@@ -1776,48 +1885,37 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1776
1885
  """
1777
1886
  ...
1778
1887
 
1779
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1780
- """
1781
- Specifies what flows belong to the same project.
1782
-
1783
- A project-specific namespace is created for all flows that
1784
- use the same `@project(name)`.
1785
-
1786
-
1787
- Parameters
1788
- ----------
1789
- name : str
1790
- Project name. Make sure that the name is unique amongst all
1791
- projects that use the same production scheduler. The name may
1792
- contain only lowercase alphanumeric characters and underscores.
1793
-
1794
- branch : Optional[str], default None
1795
- The branch to use. If not specified, the branch is set to
1796
- `user.<username>` unless `production` is set to `True`. This can
1797
- also be set on the command line using `--branch` as a top-level option.
1798
- It is an error to specify `branch` in the decorator and on the command line.
1888
+ @typing.overload
1889
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1890
+ ...
1891
+
1892
+ 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):
1893
+ """
1894
+ Specifies the times when the flow should be run when running on a
1895
+ production scheduler.
1799
1896
 
1800
- production : bool, default False
1801
- Whether or not the branch is the production branch. This can also be set on the
1802
- command line using `--production` as a top-level option. It is an error to specify
1803
- `production` in the decorator and on the command line.
1804
- The project branch name will be:
1805
- - if `branch` is specified:
1806
- - if `production` is True: `prod.<branch>`
1807
- - if `production` is False: `test.<branch>`
1808
- - if `branch` is not specified:
1809
- - if `production` is True: `prod`
1810
- - if `production` is False: `user.<username>`
1897
+
1898
+ Parameters
1899
+ ----------
1900
+ hourly : bool, default False
1901
+ Run the workflow hourly.
1902
+ daily : bool, default True
1903
+ Run the workflow daily.
1904
+ weekly : bool, default False
1905
+ Run the workflow weekly.
1906
+ cron : str, optional, default None
1907
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1908
+ specified by this expression.
1909
+ timezone : str, optional, default None
1910
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1911
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1811
1912
  """
1812
1913
  ...
1813
1914
 
1814
- 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]]:
1915
+ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1815
1916
  """
1816
- 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)
1817
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1818
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1819
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1820
- starts only after all sensors finish.
1917
+ The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1918
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1821
1919
 
1822
1920
 
1823
1921
  Parameters
@@ -1839,111 +1937,62 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1839
1937
  Name of the sensor on Airflow
1840
1938
  description : str
1841
1939
  Description of sensor in the Airflow UI
1842
- bucket_key : Union[str, List[str]]
1843
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1844
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1845
- bucket_name : str
1846
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1847
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1848
- wildcard_match : bool
1849
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1850
- aws_conn_id : str
1851
- a reference to the s3 connection on Airflow. (Default: None)
1852
- verify : bool
1853
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1940
+ external_dag_id : str
1941
+ The dag_id that contains the task you want to wait for.
1942
+ external_task_ids : List[str]
1943
+ The list of task_ids that you want to wait for.
1944
+ If None (default value) the sensor waits for the DAG. (Default: None)
1945
+ allowed_states : List[str]
1946
+ Iterable of allowed states, (Default: ['success'])
1947
+ failed_states : List[str]
1948
+ Iterable of failed or dis-allowed states. (Default: None)
1949
+ execution_delta : datetime.timedelta
1950
+ time difference with the previous execution to look at,
1951
+ the default is the same logical date as the current task or DAG. (Default: None)
1952
+ check_existence: bool
1953
+ Set to True to check if the external task exists or check if
1954
+ the DAG to wait for exists. (Default: True)
1854
1955
  """
1855
1956
  ...
1856
1957
 
1857
1958
  @typing.overload
1858
- 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]]:
1959
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1859
1960
  """
1860
- Specifies the event(s) that this flow depends on.
1861
-
1862
- ```
1863
- @trigger(event='foo')
1864
- ```
1865
- or
1866
- ```
1867
- @trigger(events=['foo', 'bar'])
1868
- ```
1869
-
1870
- Additionally, you can specify the parameter mappings
1871
- to map event payload to Metaflow parameters for the flow.
1872
- ```
1873
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1874
- ```
1875
- or
1876
- ```
1877
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1878
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1879
- ```
1880
-
1881
- 'parameters' can also be a list of strings and tuples like so:
1882
- ```
1883
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1884
- ```
1885
- This is equivalent to:
1886
- ```
1887
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1888
- ```
1961
+ Specifies the PyPI packages for all steps of the flow.
1889
1962
 
1963
+ Use `@pypi_base` to set common packages required by all
1964
+ steps and use `@pypi` to specify step-specific overrides.
1890
1965
 
1891
1966
  Parameters
1892
1967
  ----------
1893
- event : Union[str, Dict[str, Any]], optional, default None
1894
- Event dependency for this flow.
1895
- events : List[Union[str, Dict[str, Any]]], default []
1896
- Events dependency for this flow.
1897
- options : Dict[str, Any], default {}
1898
- Backend-specific configuration for tuning eventing behavior.
1968
+ packages : Dict[str, str], default: {}
1969
+ Packages to use for this flow. The key is the name of the package
1970
+ and the value is the version to use.
1971
+ python : str, optional, default: None
1972
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1973
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1899
1974
  """
1900
1975
  ...
1901
1976
 
1902
1977
  @typing.overload
1903
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1978
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1904
1979
  ...
1905
1980
 
1906
- 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] = {}):
1981
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1907
1982
  """
1908
- Specifies the event(s) that this flow depends on.
1909
-
1910
- ```
1911
- @trigger(event='foo')
1912
- ```
1913
- or
1914
- ```
1915
- @trigger(events=['foo', 'bar'])
1916
- ```
1917
-
1918
- Additionally, you can specify the parameter mappings
1919
- to map event payload to Metaflow parameters for the flow.
1920
- ```
1921
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1922
- ```
1923
- or
1924
- ```
1925
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1926
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1927
- ```
1928
-
1929
- 'parameters' can also be a list of strings and tuples like so:
1930
- ```
1931
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1932
- ```
1933
- This is equivalent to:
1934
- ```
1935
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1936
- ```
1983
+ Specifies the PyPI packages for all steps of the flow.
1937
1984
 
1985
+ Use `@pypi_base` to set common packages required by all
1986
+ steps and use `@pypi` to specify step-specific overrides.
1938
1987
 
1939
1988
  Parameters
1940
1989
  ----------
1941
- event : Union[str, Dict[str, Any]], optional, default None
1942
- Event dependency for this flow.
1943
- events : List[Union[str, Dict[str, Any]]], default []
1944
- Events dependency for this flow.
1945
- options : Dict[str, Any], default {}
1946
- Backend-specific configuration for tuning eventing behavior.
1990
+ packages : Dict[str, str], default: {}
1991
+ Packages to use for this flow. The key is the name of the package
1992
+ and the value is the version to use.
1993
+ python : str, optional, default: None
1994
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1995
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1947
1996
  """
1948
1997
  ...
1949
1998
 
@@ -1998,87 +2047,38 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1998
2047
  """
1999
2048
  ...
2000
2049
 
2001
- @typing.overload
2002
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
2050
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
2003
2051
  """
2004
- Specifies the PyPI packages for all steps of the flow.
2005
-
2006
- Use `@pypi_base` to set common packages required by all
2007
- steps and use `@pypi` to specify step-specific overrides.
2052
+ Specifies what flows belong to the same project.
2008
2053
 
2009
- Parameters
2010
- ----------
2011
- packages : Dict[str, str], default: {}
2012
- Packages to use for this flow. The key is the name of the package
2013
- and the value is the version to use.
2014
- python : str, optional, default: None
2015
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
2016
- that the version used will correspond to the version of the Python interpreter used to start the run.
2017
- """
2018
- ...
2019
-
2020
- @typing.overload
2021
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
2022
- ...
2023
-
2024
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
2025
- """
2026
- Specifies the PyPI packages for all steps of the flow.
2054
+ A project-specific namespace is created for all flows that
2055
+ use the same `@project(name)`.
2027
2056
 
2028
- Use `@pypi_base` to set common packages required by all
2029
- steps and use `@pypi` to specify step-specific overrides.
2030
2057
 
2031
2058
  Parameters
2032
2059
  ----------
2033
- packages : Dict[str, str], default: {}
2034
- Packages to use for this flow. The key is the name of the package
2035
- and the value is the version to use.
2036
- python : str, optional, default: None
2037
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
2038
- that the version used will correspond to the version of the Python interpreter used to start the run.
2039
- """
2040
- ...
2041
-
2042
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
2043
- """
2044
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
2045
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
2060
+ name : str
2061
+ Project name. Make sure that the name is unique amongst all
2062
+ projects that use the same production scheduler. The name may
2063
+ contain only lowercase alphanumeric characters and underscores.
2046
2064
 
2065
+ branch : Optional[str], default None
2066
+ The branch to use. If not specified, the branch is set to
2067
+ `user.<username>` unless `production` is set to `True`. This can
2068
+ also be set on the command line using `--branch` as a top-level option.
2069
+ It is an error to specify `branch` in the decorator and on the command line.
2047
2070
 
2048
- Parameters
2049
- ----------
2050
- timeout : int
2051
- Time, in seconds before the task times out and fails. (Default: 3600)
2052
- poke_interval : int
2053
- Time in seconds that the job should wait in between each try. (Default: 60)
2054
- mode : str
2055
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
2056
- exponential_backoff : bool
2057
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
2058
- pool : str
2059
- the slot pool this task should run in,
2060
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
2061
- soft_fail : bool
2062
- Set to true to mark the task as SKIPPED on failure. (Default: False)
2063
- name : str
2064
- Name of the sensor on Airflow
2065
- description : str
2066
- Description of sensor in the Airflow UI
2067
- external_dag_id : str
2068
- The dag_id that contains the task you want to wait for.
2069
- external_task_ids : List[str]
2070
- The list of task_ids that you want to wait for.
2071
- If None (default value) the sensor waits for the DAG. (Default: None)
2072
- allowed_states : List[str]
2073
- Iterable of allowed states, (Default: ['success'])
2074
- failed_states : List[str]
2075
- Iterable of failed or dis-allowed states. (Default: None)
2076
- execution_delta : datetime.timedelta
2077
- time difference with the previous execution to look at,
2078
- the default is the same logical date as the current task or DAG. (Default: None)
2079
- check_existence: bool
2080
- Set to True to check if the external task exists or check if
2081
- the DAG to wait for exists. (Default: True)
2071
+ production : bool, default False
2072
+ Whether or not the branch is the production branch. This can also be set on the
2073
+ command line using `--production` as a top-level option. It is an error to specify
2074
+ `production` in the decorator and on the command line.
2075
+ The project branch name will be:
2076
+ - if `branch` is specified:
2077
+ - if `production` is True: `prod.<branch>`
2078
+ - if `production` is False: `test.<branch>`
2079
+ - if `branch` is not specified:
2080
+ - if `production` is True: `prod`
2081
+ - if `production` is False: `user.<username>`
2082
2082
  """
2083
2083
  ...
2084
2084