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