ob-metaflow-stubs 6.0.6.0__py2.py3-none-any.whl → 6.0.6.1__py2.py3-none-any.whl

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