ob-metaflow-stubs 6.0.4.9__py2.py3-none-any.whl → 6.0.5.0__py2.py3-none-any.whl

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

Potentially problematic release.


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

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