ob-metaflow-stubs 6.0.10.2__py2.py3-none-any.whl → 6.0.10.3__py2.py3-none-any.whl

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

Potentially problematic release.


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

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