ob-metaflow-stubs 6.0.10.7__py2.py3-none-any.whl → 6.0.10.9__py2.py3-none-any.whl

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

Potentially problematic release.


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

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