ob-metaflow-stubs 6.0.4.7__py2.py3-none-any.whl → 6.0.4.8rc1__py2.py3-none-any.whl

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