ob-metaflow-stubs 6.0.10.13__py2.py3-none-any.whl → 6.0.10.14__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 (266) hide show
  1. metaflow-stubs/__init__.pyi +1064 -1064
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +2 -2
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +3 -3
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +2 -2
  14. metaflow-stubs/meta_files.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +69 -69
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +1 -1
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +3 -3
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +1 -1
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +1 -1
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +1 -1
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/cards/__init__.pyi +1 -1
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/cards/hf_hub_card.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +1 -1
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  64. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +1 -1
  65. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_cli.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +3 -3
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +2 -2
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +4 -4
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +2 -2
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +4 -4
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +1 -1
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +3 -3
  91. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  92. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  97. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  98. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +1 -1
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +2 -2
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  110. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  111. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  112. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  113. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  114. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  115. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  116. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  117. metaflow-stubs/mf_extensions/outerbounds/toplevel/s3_proxy.pyi +1 -1
  118. metaflow-stubs/multicore_utils.pyi +1 -1
  119. metaflow-stubs/ob_internal.pyi +1 -1
  120. metaflow-stubs/packaging_sys/__init__.pyi +4 -4
  121. metaflow-stubs/packaging_sys/backend.pyi +4 -4
  122. metaflow-stubs/packaging_sys/distribution_support.pyi +3 -3
  123. metaflow-stubs/packaging_sys/tar_backend.pyi +5 -5
  124. metaflow-stubs/packaging_sys/utils.pyi +1 -1
  125. metaflow-stubs/packaging_sys/v1.pyi +2 -2
  126. metaflow-stubs/parameters.pyi +2 -2
  127. metaflow-stubs/plugins/__init__.pyi +12 -12
  128. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  129. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  130. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  131. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  132. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  133. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  134. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  135. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  136. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  137. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  138. metaflow-stubs/plugins/argo/argo_workflows.pyi +1 -1
  139. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  140. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  141. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  142. metaflow-stubs/plugins/argo/exit_hooks.pyi +1 -1
  143. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  144. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  145. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  146. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  147. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  148. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  149. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  150. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  151. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  152. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  153. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  154. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  155. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  156. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  157. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  158. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +1 -1
  159. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  160. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  161. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  162. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  163. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  164. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  165. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  166. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  167. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  168. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  169. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  170. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  171. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  172. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  173. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  174. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  175. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  176. metaflow-stubs/plugins/cards/card_modules/json_viewer.pyi +1 -1
  177. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  178. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  179. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  180. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  181. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  182. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  183. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  184. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  185. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  186. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  187. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  188. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  189. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  190. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  191. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  192. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  193. metaflow-stubs/plugins/exit_hook/__init__.pyi +1 -1
  194. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +1 -1
  195. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  196. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  197. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  198. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  199. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  200. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  201. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  202. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  203. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  204. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  205. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  206. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  207. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  208. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  209. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  210. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  211. metaflow-stubs/plugins/optuna/__init__.pyi +1 -1
  212. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  213. metaflow-stubs/plugins/parsers.pyi +1 -1
  214. metaflow-stubs/plugins/perimeters.pyi +1 -1
  215. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  216. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  217. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  218. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  219. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  220. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  221. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  222. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  223. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  224. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  225. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  226. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  227. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  228. metaflow-stubs/plugins/secrets/secrets_func.pyi +1 -1
  229. metaflow-stubs/plugins/secrets/secrets_spec.pyi +1 -1
  230. metaflow-stubs/plugins/secrets/utils.pyi +1 -1
  231. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  232. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  233. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  234. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  235. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  236. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  237. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  238. metaflow-stubs/profilers/__init__.pyi +1 -1
  239. metaflow-stubs/pylint_wrapper.pyi +1 -1
  240. metaflow-stubs/runner/__init__.pyi +1 -1
  241. metaflow-stubs/runner/deployer.pyi +32 -32
  242. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  243. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  244. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  245. metaflow-stubs/runner/nbrun.pyi +1 -1
  246. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  247. metaflow-stubs/runner/utils.pyi +1 -1
  248. metaflow-stubs/system/__init__.pyi +1 -1
  249. metaflow-stubs/system/system_logger.pyi +1 -1
  250. metaflow-stubs/system/system_monitor.pyi +1 -1
  251. metaflow-stubs/tagging_util.pyi +1 -1
  252. metaflow-stubs/tuple_util.pyi +1 -1
  253. metaflow-stubs/user_configs/__init__.pyi +1 -1
  254. metaflow-stubs/user_configs/config_options.pyi +2 -2
  255. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  256. metaflow-stubs/user_decorators/__init__.pyi +1 -1
  257. metaflow-stubs/user_decorators/common.pyi +1 -1
  258. metaflow-stubs/user_decorators/mutable_flow.pyi +4 -4
  259. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  260. metaflow-stubs/user_decorators/user_flow_decorator.pyi +3 -3
  261. metaflow-stubs/user_decorators/user_step_decorator.pyi +4 -4
  262. {ob_metaflow_stubs-6.0.10.13.dist-info → ob_metaflow_stubs-6.0.10.14.dist-info}/METADATA +1 -1
  263. ob_metaflow_stubs-6.0.10.14.dist-info/RECORD +266 -0
  264. ob_metaflow_stubs-6.0.10.13.dist-info/RECORD +0 -266
  265. {ob_metaflow_stubs-6.0.10.13.dist-info → ob_metaflow_stubs-6.0.10.14.dist-info}/WHEEL +0 -0
  266. {ob_metaflow_stubs-6.0.10.13.dist-info → ob_metaflow_stubs-6.0.10.14.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.18.10.1+obcheckpoint(0.2.8);ob(v1) #
4
- # Generated on 2025-10-02T16:55:23.853861 #
4
+ # Generated on 2025-10-08T21:13:44.622359 #
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
 
@@ -39,19 +39,19 @@ from .user_decorators.user_step_decorator import UserStepDecorator as UserStepDe
39
39
  from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
- from . import tuple_util as tuple_util
43
42
  from . import cards as cards
44
- from . import metaflow_git as metaflow_git
45
43
  from . import events as events
44
+ from . import tuple_util as tuple_util
45
+ from . import metaflow_git as metaflow_git
46
46
  from . import runner as runner
47
47
  from . import plugins as plugins
48
48
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
49
49
  from . import includefile as includefile
50
50
  from .includefile import IncludeFile as IncludeFile
51
- from .plugins.parsers import yaml_parser as yaml_parser
52
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
53
51
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
54
52
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
53
+ from .plugins.parsers import yaml_parser as yaml_parser
54
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
55
55
  from . import client as client
56
56
  from .client.core import namespace as namespace
57
57
  from .client.core import get_namespace as get_namespace
@@ -170,191 +170,83 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
170
170
  ...
171
171
 
172
172
  @typing.overload
173
- 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]]]:
173
+ 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]]]:
174
174
  """
175
- Specifies the number of times the task corresponding
176
- to a step needs to be retried.
175
+ Specifies a timeout for your step.
177
176
 
178
- This decorator is useful for handling transient errors, such as networking issues.
179
- If your task contains operations that can't be retried safely, e.g. database updates,
180
- it is advisable to annotate it with `@retry(times=0)`.
177
+ This decorator is useful if this step may hang indefinitely.
181
178
 
182
- This can be used in conjunction with the `@catch` decorator. The `@catch`
183
- decorator will execute a no-op task after all retries have been exhausted,
184
- ensuring that the flow execution can continue.
179
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
180
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
181
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
182
+
183
+ Note that all the values specified in parameters are added together so if you specify
184
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
185
185
 
186
186
 
187
187
  Parameters
188
188
  ----------
189
- times : int, default 3
190
- Number of times to retry this task.
191
- minutes_between_retries : int, default 2
192
- Number of minutes between retries.
189
+ seconds : int, default 0
190
+ Number of seconds to wait prior to timing out.
191
+ minutes : int, default 0
192
+ Number of minutes to wait prior to timing out.
193
+ hours : int, default 0
194
+ Number of hours to wait prior to timing out.
193
195
  """
194
196
  ...
195
197
 
196
198
  @typing.overload
197
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
199
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
198
200
  ...
199
201
 
200
202
  @typing.overload
201
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
203
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
202
204
  ...
203
205
 
204
- 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):
206
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
205
207
  """
206
- Specifies the number of times the task corresponding
207
- to a step needs to be retried.
208
+ Specifies a timeout for your step.
208
209
 
209
- This decorator is useful for handling transient errors, such as networking issues.
210
- If your task contains operations that can't be retried safely, e.g. database updates,
211
- it is advisable to annotate it with `@retry(times=0)`.
210
+ This decorator is useful if this step may hang indefinitely.
212
211
 
213
- This can be used in conjunction with the `@catch` decorator. The `@catch`
214
- decorator will execute a no-op task after all retries have been exhausted,
215
- ensuring that the flow execution can continue.
212
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
213
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
214
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
215
+
216
+ Note that all the values specified in parameters are added together so if you specify
217
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
216
218
 
217
219
 
218
220
  Parameters
219
221
  ----------
220
- times : int, default 3
221
- Number of times to retry this task.
222
- minutes_between_retries : int, default 2
223
- Number of minutes between retries.
222
+ seconds : int, default 0
223
+ Number of seconds to wait prior to timing out.
224
+ minutes : int, default 0
225
+ Number of minutes to wait prior to timing out.
226
+ hours : int, default 0
227
+ Number of hours to wait prior to timing out.
224
228
  """
225
229
  ...
226
230
 
227
231
  @typing.overload
228
- 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]]]:
232
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
229
233
  """
230
- Enables loading / saving of models within a step.
231
-
232
- > Examples
233
- - Saving Models
234
- ```python
235
- @model
236
- @step
237
- def train(self):
238
- # current.model.save returns a dictionary reference to the model saved
239
- self.my_model = current.model.save(
240
- path_to_my_model,
241
- label="my_model",
242
- metadata={
243
- "epochs": 10,
244
- "batch-size": 32,
245
- "learning-rate": 0.001,
246
- }
247
- )
248
- self.next(self.test)
249
-
250
- @model(load="my_model")
251
- @step
252
- def test(self):
253
- # `current.model.loaded` returns a dictionary of the loaded models
254
- # where the key is the name of the artifact and the value is the path to the model
255
- print(os.listdir(current.model.loaded["my_model"]))
256
- self.next(self.end)
257
- ```
258
-
259
- - Loading models
260
- ```python
261
- @step
262
- def train(self):
263
- # current.model.load returns the path to the model loaded
264
- checkpoint_path = current.model.load(
265
- self.checkpoint_key,
266
- )
267
- model_path = current.model.load(
268
- self.model,
269
- )
270
- self.next(self.test)
271
- ```
272
-
273
-
274
- Parameters
275
- ----------
276
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
277
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
278
- 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`.
279
- 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
280
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
281
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
282
-
283
- temp_dir_root : str, default: None
284
- The root directory under which `current.model.loaded` will store loaded models
234
+ Internal decorator to support Fast bakery
285
235
  """
286
236
  ...
287
237
 
288
238
  @typing.overload
289
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
290
- ...
291
-
292
- @typing.overload
293
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
239
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
294
240
  ...
295
241
 
296
- 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):
242
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
297
243
  """
298
- Enables loading / saving of models within a step.
299
-
300
- > Examples
301
- - Saving Models
302
- ```python
303
- @model
304
- @step
305
- def train(self):
306
- # current.model.save returns a dictionary reference to the model saved
307
- self.my_model = current.model.save(
308
- path_to_my_model,
309
- label="my_model",
310
- metadata={
311
- "epochs": 10,
312
- "batch-size": 32,
313
- "learning-rate": 0.001,
314
- }
315
- )
316
- self.next(self.test)
317
-
318
- @model(load="my_model")
319
- @step
320
- def test(self):
321
- # `current.model.loaded` returns a dictionary of the loaded models
322
- # where the key is the name of the artifact and the value is the path to the model
323
- print(os.listdir(current.model.loaded["my_model"]))
324
- self.next(self.end)
325
- ```
326
-
327
- - Loading models
328
- ```python
329
- @step
330
- def train(self):
331
- # current.model.load returns the path to the model loaded
332
- checkpoint_path = current.model.load(
333
- self.checkpoint_key,
334
- )
335
- model_path = current.model.load(
336
- self.model,
337
- )
338
- self.next(self.test)
339
- ```
340
-
341
-
342
- Parameters
343
- ----------
344
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
345
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
346
- 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`.
347
- 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
348
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
349
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
350
-
351
- temp_dir_root : str, default: None
352
- The root directory under which `current.model.loaded` will store loaded models
244
+ Internal decorator to support Fast bakery
353
245
  """
354
246
  ...
355
247
 
356
248
  @typing.overload
357
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
249
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
358
250
  """
359
251
  Decorator prototype for all step decorators. This function gets specialized
360
252
  and imported for all decorators types by _import_plugin_decorators().
@@ -362,163 +254,224 @@ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.
362
254
  ...
363
255
 
364
256
  @typing.overload
365
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
257
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
366
258
  ...
367
259
 
368
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
260
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
369
261
  """
370
262
  Decorator prototype for all step decorators. This function gets specialized
371
263
  and imported for all decorators types by _import_plugin_decorators().
372
264
  """
373
265
  ...
374
266
 
375
- @typing.overload
376
- 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]]]:
377
- """
378
- Specifies environment variables to be set prior to the execution of a step.
379
-
380
-
381
- Parameters
382
- ----------
383
- vars : Dict[str, str], default {}
384
- Dictionary of environment variables to set.
385
- """
386
- ...
387
-
388
- @typing.overload
389
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
390
- ...
391
-
392
- @typing.overload
393
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
394
- ...
395
-
396
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
267
+ def coreweave_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]]]:
397
268
  """
398
- Specifies environment variables to be set prior to the execution of a step.
269
+ `@coreweave_s3_proxy` is a CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
270
+ It exists to make it easier for users to know that this decorator should only be used with
271
+ a Neo Cloud like CoreWeave. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
399
272
 
400
273
 
401
- Parameters
402
- ----------
403
- vars : Dict[str, str], default {}
404
- Dictionary of environment variables to set.
405
- """
406
- ...
407
-
408
- @typing.overload
409
- 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]]]:
410
- """
411
- Specifies the Conda environment for the step.
274
+ Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
275
+ for S3 read and write requests.
412
276
 
413
- Information in this decorator will augment any
414
- attributes set in the `@conda_base` flow-level decorator. Hence,
415
- you can use `@conda_base` to set packages required by all
416
- steps and use `@conda` to specify step-specific overrides.
277
+ This decorator requires an integration in the Outerbounds platform that
278
+ points to an external bucket. It affects S3 operations performed via
279
+ Metaflow's `get_aws_client` and `S3` within a `@step`.
417
280
 
281
+ Read operations
282
+ ---------------
283
+ All read operations pass through the proxy. If an object does not already
284
+ exist in the external bucket, it is cached there. For example, if code reads
285
+ from buckets `FOO` and `BAR` using the `S3` interface, objects from both
286
+ buckets are cached in the external bucket.
418
287
 
419
- Parameters
420
- ----------
421
- packages : Dict[str, str], default {}
422
- Packages to use for this step. The key is the name of the package
423
- and the value is the version to use.
424
- libraries : Dict[str, str], default {}
425
- Supported for backward compatibility. When used with packages, packages will take precedence.
426
- python : str, optional, default None
427
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
428
- that the version used will correspond to the version of the Python interpreter used to start the run.
429
- disabled : bool, default False
430
- If set to True, disables @conda.
288
+ During task execution, all S3‑related read requests are routed through the
289
+ proxy:
290
+ - If the object is present in the external object store, the proxy
291
+ streams it directly from there without accessing the requested origin
292
+ bucket.
293
+ - If the object is not present in the external storage, the proxy
294
+ fetches it from the requested bucket, caches it in the external
295
+ storage, and streams the response from the origin bucket.
296
+
297
+ Warning
298
+ -------
299
+ All READ operations (e.g., GetObject, HeadObject) pass through the external
300
+ bucket regardless of the bucket specified in user code. Even
301
+ `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
302
+ external bucket cache.
303
+
304
+ Write operations
305
+ ----------------
306
+ Write behavior is controlled by the `write_mode` parameter, which determines
307
+ whether writes also persist objects in the cache.
308
+
309
+ `write_mode` values:
310
+ - `origin-and-cache`: objects are written both to the cache and to their
311
+ intended origin bucket.
312
+ - `origin`: objects are written only to their intended origin bucket.
313
+
314
+
315
+ Parameters
316
+ ----------
317
+ integration_name : str, optional
318
+ [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
319
+ that holds the configuration for the external, S3‑compatible object
320
+ storage bucket. If not specified, the only available S3 proxy
321
+ integration in the namespace is used (fails if multiple exist).
322
+ write_mode : str, optional
323
+ Controls whether writes also go to the external bucket.
324
+ - `origin` (default)
325
+ - `origin-and-cache`
326
+ debug : bool, optional
327
+ Enables debug logging for proxy operations.
431
328
  """
432
329
  ...
433
330
 
434
331
  @typing.overload
435
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
332
+ 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]]]:
333
+ """
334
+ Specifies environment variables to be set prior to the execution of a step.
335
+
336
+
337
+ Parameters
338
+ ----------
339
+ vars : Dict[str, str], default {}
340
+ Dictionary of environment variables to set.
341
+ """
436
342
  ...
437
343
 
438
344
  @typing.overload
439
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
345
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
440
346
  ...
441
347
 
442
- 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):
348
+ @typing.overload
349
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
350
+ ...
351
+
352
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
443
353
  """
444
- Specifies the Conda environment for the step.
445
-
446
- Information in this decorator will augment any
447
- attributes set in the `@conda_base` flow-level decorator. Hence,
448
- you can use `@conda_base` to set packages required by all
449
- steps and use `@conda` to specify step-specific overrides.
354
+ Specifies environment variables to be set prior to the execution of a step.
450
355
 
451
356
 
452
357
  Parameters
453
358
  ----------
454
- packages : Dict[str, str], default {}
455
- Packages to use for this step. The key is the name of the package
456
- and the value is the version to use.
457
- libraries : Dict[str, str], default {}
458
- Supported for backward compatibility. When used with packages, packages will take precedence.
459
- python : str, optional, default None
460
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
461
- that the version used will correspond to the version of the Python interpreter used to start the run.
462
- disabled : bool, default False
463
- If set to True, disables @conda.
359
+ vars : Dict[str, str], default {}
360
+ Dictionary of environment variables to set.
464
361
  """
465
362
  ...
466
363
 
467
- @typing.overload
468
- 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]]]:
364
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
469
365
  """
470
- Specifies secrets to be retrieved and injected as environment variables prior to
471
- the execution of a step.
366
+ Specifies that this step should execute on DGX cloud.
472
367
 
473
368
 
474
369
  Parameters
475
370
  ----------
476
- sources : List[Union[str, Dict[str, Any]]], default: []
477
- List of secret specs, defining how the secrets are to be retrieved
478
- role : str, optional, default: None
479
- Role to use for fetching secrets
371
+ gpu : int
372
+ Number of GPUs to use.
373
+ gpu_type : str
374
+ Type of Nvidia GPU to use.
375
+ queue_timeout : int
376
+ Time to keep the job in NVCF's queue.
480
377
  """
481
378
  ...
482
379
 
483
380
  @typing.overload
484
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
381
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
382
+ """
383
+ A simple decorator that demonstrates using CardDecoratorInjector
384
+ to inject a card and render simple markdown content.
385
+ """
485
386
  ...
486
387
 
487
388
  @typing.overload
488
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
389
+ def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
489
390
  ...
490
391
 
491
- 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):
392
+ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
492
393
  """
493
- Specifies secrets to be retrieved and injected as environment variables prior to
494
- the execution of a step.
394
+ A simple decorator that demonstrates using CardDecoratorInjector
395
+ to inject a card and render simple markdown content.
396
+ """
397
+ ...
398
+
399
+ @typing.overload
400
+ 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]]]:
401
+ """
402
+ Specifies the resources needed when executing this step.
403
+
404
+ Use `@resources` to specify the resource requirements
405
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
406
+
407
+ You can choose the compute layer on the command line by executing e.g.
408
+ ```
409
+ python myflow.py run --with batch
410
+ ```
411
+ or
412
+ ```
413
+ python myflow.py run --with kubernetes
414
+ ```
415
+ which executes the flow on the desired system using the
416
+ requirements specified in `@resources`.
495
417
 
496
418
 
497
419
  Parameters
498
420
  ----------
499
- sources : List[Union[str, Dict[str, Any]]], default: []
500
- List of secret specs, defining how the secrets are to be retrieved
501
- role : str, optional, default: None
502
- Role to use for fetching secrets
421
+ cpu : int, default 1
422
+ Number of CPUs required for this step.
423
+ gpu : int, optional, default None
424
+ Number of GPUs required for this step.
425
+ disk : int, optional, default None
426
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
427
+ memory : int, default 4096
428
+ Memory size (in MB) required for this step.
429
+ shared_memory : int, optional, default None
430
+ The value for the size (in MiB) of the /dev/shm volume for this step.
431
+ This parameter maps to the `--shm-size` option in Docker.
503
432
  """
504
433
  ...
505
434
 
506
435
  @typing.overload
507
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
508
- """
509
- Decorator prototype for all step decorators. This function gets specialized
510
- and imported for all decorators types by _import_plugin_decorators().
511
- """
436
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
512
437
  ...
513
438
 
514
439
  @typing.overload
515
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
440
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
516
441
  ...
517
442
 
518
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
443
+ 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):
519
444
  """
520
- Decorator prototype for all step decorators. This function gets specialized
521
- and imported for all decorators types by _import_plugin_decorators().
445
+ Specifies the resources needed when executing this step.
446
+
447
+ Use `@resources` to specify the resource requirements
448
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
449
+
450
+ You can choose the compute layer on the command line by executing e.g.
451
+ ```
452
+ python myflow.py run --with batch
453
+ ```
454
+ or
455
+ ```
456
+ python myflow.py run --with kubernetes
457
+ ```
458
+ which executes the flow on the desired system using the
459
+ requirements specified in `@resources`.
460
+
461
+
462
+ Parameters
463
+ ----------
464
+ cpu : int, default 1
465
+ Number of CPUs required for this step.
466
+ gpu : int, optional, default None
467
+ Number of GPUs required for this step.
468
+ disk : int, optional, default None
469
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
470
+ memory : int, default 4096
471
+ Memory size (in MB) required for this step.
472
+ shared_memory : int, optional, default None
473
+ The value for the size (in MiB) of the /dev/shm volume for this step.
474
+ This parameter maps to the `--shm-size` option in Docker.
522
475
  """
523
476
  ...
524
477
 
@@ -572,8 +525,13 @@ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card
572
525
  """
573
526
  ...
574
527
 
575
- 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]]]:
528
+ def nebius_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]]]:
576
529
  """
530
+ `@nebius_s3_proxy` is a Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
531
+ It exists to make it easier for users to know that this decorator should only be used with
532
+ a Neo Cloud like Nebius. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
533
+
534
+
577
535
  Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
578
536
  for S3 read and write requests.
579
537
 
@@ -632,559 +590,445 @@ def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typin
632
590
  ...
633
591
 
634
592
  @typing.overload
635
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
593
+ 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]]]:
636
594
  """
637
- Internal decorator to support Fast bakery
595
+ Specifies the number of times the task corresponding
596
+ to a step needs to be retried.
597
+
598
+ This decorator is useful for handling transient errors, such as networking issues.
599
+ If your task contains operations that can't be retried safely, e.g. database updates,
600
+ it is advisable to annotate it with `@retry(times=0)`.
601
+
602
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
603
+ decorator will execute a no-op task after all retries have been exhausted,
604
+ ensuring that the flow execution can continue.
605
+
606
+
607
+ Parameters
608
+ ----------
609
+ times : int, default 3
610
+ Number of times to retry this task.
611
+ minutes_between_retries : int, default 2
612
+ Number of minutes between retries.
638
613
  """
639
614
  ...
640
615
 
641
616
  @typing.overload
642
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
617
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
643
618
  ...
644
619
 
645
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
620
+ @typing.overload
621
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
622
+ ...
623
+
624
+ 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):
646
625
  """
647
- Internal decorator to support Fast bakery
626
+ Specifies the number of times the task corresponding
627
+ to a step needs to be retried.
628
+
629
+ This decorator is useful for handling transient errors, such as networking issues.
630
+ If your task contains operations that can't be retried safely, e.g. database updates,
631
+ it is advisable to annotate it with `@retry(times=0)`.
632
+
633
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
634
+ decorator will execute a no-op task after all retries have been exhausted,
635
+ ensuring that the flow execution can continue.
636
+
637
+
638
+ Parameters
639
+ ----------
640
+ times : int, default 3
641
+ Number of times to retry this task.
642
+ minutes_between_retries : int, default 2
643
+ Number of minutes between retries.
648
644
  """
649
645
  ...
650
646
 
651
- 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]]]:
647
+ @typing.overload
648
+ 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]]]:
652
649
  """
653
- Decorator that helps cache, version, and store models/datasets from the Hugging Face Hub.
654
-
655
- Examples
656
- --------
657
-
658
- ```python
659
- # **Usage: creating references to models from the Hugging Face Hub that may be loaded in downstream steps**
660
- @huggingface_hub
661
- @step
662
- def pull_model_from_huggingface(self):
663
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
664
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
665
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
666
- # value of the function is a reference to the model in the backend storage.
667
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
668
-
669
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
670
- self.llama_model = current.huggingface_hub.snapshot_download(
671
- repo_id=self.model_id,
672
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
673
- )
674
- self.next(self.train)
675
-
676
- # **Usage: explicitly loading models at runtime from the Hugging Face Hub or from cache (from Metaflow's datastore)**
677
- @huggingface_hub
678
- @step
679
- def run_training(self):
680
- # Temporary directory (auto-cleaned on exit)
681
- with current.huggingface_hub.load(
682
- repo_id="google-bert/bert-base-uncased",
683
- allow_patterns=["*.bin"],
684
- ) as local_path:
685
- # Use files under local_path
686
- train_model(local_path)
687
- ...
688
-
689
- # **Usage: loading models directly from the Hugging Face Hub or from cache (from Metaflow's datastore)**
690
-
691
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
692
- @step
693
- def pull_model_from_huggingface(self):
694
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
695
-
696
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora", "/my-lora-directory")])
697
- @step
698
- def finetune_model(self):
699
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
700
- # path_to_model will be /my-directory
701
-
702
-
703
- # Takes all the arguments passed to `snapshot_download`
704
- # except for `local_dir`
705
- @huggingface_hub(load=[
706
- {
707
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
708
- },
709
- {
710
- "repo_id": "myorg/mistral-lora",
711
- "repo_type": "model",
712
- },
713
- ])
714
- @step
715
- def finetune_model(self):
716
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
717
- # path_to_model will be /my-directory
718
- ```
650
+ Specifies secrets to be retrieved and injected as environment variables prior to
651
+ the execution of a step.
719
652
 
720
653
 
721
654
  Parameters
722
655
  ----------
723
- temp_dir_root : str, optional
724
- The root directory that will hold the temporary directory where objects will be downloaded.
725
-
726
- cache_scope : str, optional
727
- The scope of the cache. Can be `checkpoint` / `flow` / `global`.
728
- - `checkpoint` (default): All repos are stored like objects saved by `@checkpoint`.
729
- i.e., the cached path is derived from the namespace, flow, step, and Metaflow foreach iteration.
730
- 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).
731
-
732
- - `flow`: All repos are cached under the flow, regardless of namespace.
733
- i.e., the cached path is derived solely from the flow name.
734
- 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.
735
-
736
- - `global`: All repos are cached under a globally static path.
737
- i.e., the base path of the cache is static and all repos are stored under it.
738
- When to use this mode:
739
- - All repos from the Hugging Face Hub need to be shared by users across all flow executions.
740
- - Each caching scope comes with its own trade-offs:
741
- - `checkpoint`:
742
- - 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.
743
- - 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.
744
- - `flow`:
745
- - 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.
746
- - The blast radius of a bad checkpoint is limited to all runs of a particular flow.
747
- - It doesn't promote cache reuse across flows.
748
- - `global`:
749
- - Has no control over when caches are populated (can be written by any flow execution) but has the highest cache hit rate.
750
- - It promotes cache reuse across flows.
751
- - The blast radius of a bad checkpoint spans every flow that could be using a particular repo.
752
-
753
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
754
- The list of repos (models/datasets) to load.
755
-
756
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
656
+ sources : List[Union[str, Dict[str, Any]]], default: []
657
+ List of secret specs, defining how the secrets are to be retrieved
658
+ role : str, optional, default: None
659
+ Role to use for fetching secrets
660
+ """
661
+ ...
662
+
663
+ @typing.overload
664
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
665
+ ...
666
+
667
+ @typing.overload
668
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
669
+ ...
670
+
671
+ 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):
672
+ """
673
+ Specifies secrets to be retrieved and injected as environment variables prior to
674
+ the execution of a step.
757
675
 
758
- - If repo (model/dataset) is not found in the datastore:
759
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
760
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
761
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
762
676
 
763
- - If repo is found in the datastore:
764
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
677
+ Parameters
678
+ ----------
679
+ sources : List[Union[str, Dict[str, Any]]], default: []
680
+ List of secret specs, defining how the secrets are to be retrieved
681
+ role : str, optional, default: None
682
+ Role to use for fetching secrets
765
683
  """
766
684
  ...
767
685
 
768
686
  @typing.overload
769
- 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]]]:
687
+ 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]]]:
770
688
  """
771
- Specifies the resources needed when executing this step.
689
+ Enables loading / saving of models within a step.
772
690
 
773
- Use `@resources` to specify the resource requirements
774
- independently of the specific compute layer (`@batch`, `@kubernetes`).
691
+ > Examples
692
+ - Saving Models
693
+ ```python
694
+ @model
695
+ @step
696
+ def train(self):
697
+ # current.model.save returns a dictionary reference to the model saved
698
+ self.my_model = current.model.save(
699
+ path_to_my_model,
700
+ label="my_model",
701
+ metadata={
702
+ "epochs": 10,
703
+ "batch-size": 32,
704
+ "learning-rate": 0.001,
705
+ }
706
+ )
707
+ self.next(self.test)
775
708
 
776
- You can choose the compute layer on the command line by executing e.g.
777
- ```
778
- python myflow.py run --with batch
779
- ```
780
- or
709
+ @model(load="my_model")
710
+ @step
711
+ def test(self):
712
+ # `current.model.loaded` returns a dictionary of the loaded models
713
+ # where the key is the name of the artifact and the value is the path to the model
714
+ print(os.listdir(current.model.loaded["my_model"]))
715
+ self.next(self.end)
781
716
  ```
782
- python myflow.py run --with kubernetes
717
+
718
+ - Loading models
719
+ ```python
720
+ @step
721
+ def train(self):
722
+ # current.model.load returns the path to the model loaded
723
+ checkpoint_path = current.model.load(
724
+ self.checkpoint_key,
725
+ )
726
+ model_path = current.model.load(
727
+ self.model,
728
+ )
729
+ self.next(self.test)
783
730
  ```
784
- which executes the flow on the desired system using the
785
- requirements specified in `@resources`.
786
731
 
787
732
 
788
733
  Parameters
789
734
  ----------
790
- cpu : int, default 1
791
- Number of CPUs required for this step.
792
- gpu : int, optional, default None
793
- Number of GPUs required for this step.
794
- disk : int, optional, default None
795
- Disk size (in MB) required for this step. Only applies on Kubernetes.
796
- memory : int, default 4096
797
- Memory size (in MB) required for this step.
798
- shared_memory : int, optional, default None
799
- The value for the size (in MiB) of the /dev/shm volume for this step.
800
- This parameter maps to the `--shm-size` option in Docker.
735
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
736
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
737
+ 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`.
738
+ 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
739
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
740
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
741
+
742
+ temp_dir_root : str, default: None
743
+ The root directory under which `current.model.loaded` will store loaded models
801
744
  """
802
745
  ...
803
746
 
804
747
  @typing.overload
805
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
748
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
806
749
  ...
807
750
 
808
751
  @typing.overload
809
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
752
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
810
753
  ...
811
754
 
812
- 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):
755
+ 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):
813
756
  """
814
- Specifies the resources needed when executing this step.
757
+ Enables loading / saving of models within a step.
815
758
 
816
- Use `@resources` to specify the resource requirements
817
- independently of the specific compute layer (`@batch`, `@kubernetes`).
759
+ > Examples
760
+ - Saving Models
761
+ ```python
762
+ @model
763
+ @step
764
+ def train(self):
765
+ # current.model.save returns a dictionary reference to the model saved
766
+ self.my_model = current.model.save(
767
+ path_to_my_model,
768
+ label="my_model",
769
+ metadata={
770
+ "epochs": 10,
771
+ "batch-size": 32,
772
+ "learning-rate": 0.001,
773
+ }
774
+ )
775
+ self.next(self.test)
818
776
 
819
- You can choose the compute layer on the command line by executing e.g.
820
- ```
821
- python myflow.py run --with batch
822
- ```
823
- or
777
+ @model(load="my_model")
778
+ @step
779
+ def test(self):
780
+ # `current.model.loaded` returns a dictionary of the loaded models
781
+ # where the key is the name of the artifact and the value is the path to the model
782
+ print(os.listdir(current.model.loaded["my_model"]))
783
+ self.next(self.end)
824
784
  ```
825
- python myflow.py run --with kubernetes
785
+
786
+ - Loading models
787
+ ```python
788
+ @step
789
+ def train(self):
790
+ # current.model.load returns the path to the model loaded
791
+ checkpoint_path = current.model.load(
792
+ self.checkpoint_key,
793
+ )
794
+ model_path = current.model.load(
795
+ self.model,
796
+ )
797
+ self.next(self.test)
826
798
  ```
827
- which executes the flow on the desired system using the
828
- requirements specified in `@resources`.
829
799
 
830
800
 
831
801
  Parameters
832
802
  ----------
833
- cpu : int, default 1
834
- Number of CPUs required for this step.
835
- gpu : int, optional, default None
836
- Number of GPUs required for this step.
837
- disk : int, optional, default None
838
- Disk size (in MB) required for this step. Only applies on Kubernetes.
839
- memory : int, default 4096
840
- Memory size (in MB) required for this step.
841
- shared_memory : int, optional, default None
842
- The value for the size (in MiB) of the /dev/shm volume for this step.
843
- This parameter maps to the `--shm-size` option in Docker.
844
- """
845
- ...
846
-
847
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
848
- """
849
- Specifies that this step should execute on Kubernetes.
850
-
851
-
852
- Parameters
853
- ----------
854
- cpu : int, default 1
855
- Number of CPUs required for this step. If `@resources` is
856
- also present, the maximum value from all decorators is used.
857
- memory : int, default 4096
858
- Memory size (in MB) required for this step. If
859
- `@resources` is also present, the maximum value from all decorators is
860
- used.
861
- disk : int, default 10240
862
- Disk size (in MB) required for this step. If
863
- `@resources` is also present, the maximum value from all decorators is
864
- used.
865
- image : str, optional, default None
866
- Docker image to use when launching on Kubernetes. If not specified, and
867
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
868
- not, a default Docker image mapping to the current version of Python is used.
869
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
870
- If given, the imagePullPolicy to be applied to the Docker image of the step.
871
- image_pull_secrets: List[str], default []
872
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
873
- Kubernetes image pull secrets to use when pulling container images
874
- in Kubernetes.
875
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
876
- Kubernetes service account to use when launching pod in Kubernetes.
877
- secrets : List[str], optional, default None
878
- Kubernetes secrets to use when launching pod in Kubernetes. These
879
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
880
- in Metaflow configuration.
881
- node_selector: Union[Dict[str,str], str], optional, default None
882
- Kubernetes node selector(s) to apply to the pod running the task.
883
- Can be passed in as a comma separated string of values e.g.
884
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
885
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
886
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
887
- Kubernetes namespace to use when launching pod in Kubernetes.
888
- gpu : int, optional, default None
889
- Number of GPUs required for this step. A value of zero implies that
890
- the scheduled node should not have GPUs.
891
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
892
- The vendor of the GPUs to be used for this step.
893
- tolerations : List[Dict[str,str]], default []
894
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
895
- Kubernetes tolerations to use when launching pod in Kubernetes.
896
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
897
- Kubernetes labels to use when launching pod in Kubernetes.
898
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
899
- Kubernetes annotations to use when launching pod in Kubernetes.
900
- use_tmpfs : bool, default False
901
- This enables an explicit tmpfs mount for this step.
902
- tmpfs_tempdir : bool, default True
903
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
904
- tmpfs_size : int, optional, default: None
905
- The value for the size (in MiB) of the tmpfs mount for this step.
906
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
907
- memory allocated for this step.
908
- tmpfs_path : str, optional, default /metaflow_temp
909
- Path to tmpfs mount for this step.
910
- persistent_volume_claims : Dict[str, str], optional, default None
911
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
912
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
913
- shared_memory: int, optional
914
- Shared memory size (in MiB) required for this step
915
- port: int, optional
916
- Port number to specify in the Kubernetes job object
917
- compute_pool : str, optional, default None
918
- Compute pool to be used for for this step.
919
- If not specified, any accessible compute pool within the perimeter is used.
920
- hostname_resolution_timeout: int, default 10 * 60
921
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
922
- Only applicable when @parallel is used.
923
- qos: str, default: Burstable
924
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
803
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
804
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
805
+ 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`.
806
+ 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
807
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
808
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
925
809
 
926
- security_context: Dict[str, Any], optional, default None
927
- Container security context. Applies to the task container. Allows the following keys:
928
- - privileged: bool, optional, default None
929
- - allow_privilege_escalation: bool, optional, default None
930
- - run_as_user: int, optional, default None
931
- - run_as_group: int, optional, default None
932
- - run_as_non_root: bool, optional, default None
933
- """
934
- ...
935
-
936
- @typing.overload
937
- def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
938
- """
939
- A simple decorator that demonstrates using CardDecoratorInjector
940
- to inject a card and render simple markdown content.
941
- """
942
- ...
943
-
944
- @typing.overload
945
- def test_append_card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
946
- ...
947
-
948
- def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
949
- """
950
- A simple decorator that demonstrates using CardDecoratorInjector
951
- to inject a card and render simple markdown content.
810
+ temp_dir_root : str, default: None
811
+ The root directory under which `current.model.loaded` will store loaded models
952
812
  """
953
813
  ...
954
814
 
955
- def coreweave_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]]]:
815
+ 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]]]:
956
816
  """
957
- `@coreweave_s3_proxy` is a CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
958
- It exists to make it easier for users to know that this decorator should only be used with
959
- a Neo Cloud like CoreWeave. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
960
-
961
-
962
- Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
963
- for S3 read and write requests.
964
-
965
- This decorator requires an integration in the Outerbounds platform that
966
- points to an external bucket. It affects S3 operations performed via
967
- Metaflow's `get_aws_client` and `S3` within a `@step`.
968
-
969
- Read operations
970
- ---------------
971
- All read operations pass through the proxy. If an object does not already
972
- exist in the external bucket, it is cached there. For example, if code reads
973
- from buckets `FOO` and `BAR` using the `S3` interface, objects from both
974
- buckets are cached in the external bucket.
975
-
976
- During task execution, all S3‑related read requests are routed through the
977
- proxy:
978
- - If the object is present in the external object store, the proxy
979
- streams it directly from there without accessing the requested origin
980
- bucket.
981
- - If the object is not present in the external storage, the proxy
982
- fetches it from the requested bucket, caches it in the external
983
- storage, and streams the response from the origin bucket.
984
-
985
- Warning
986
- -------
987
- All READ operations (e.g., GetObject, HeadObject) pass through the external
988
- bucket regardless of the bucket specified in user code. Even
989
- `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
990
- external bucket cache.
991
-
992
- Write operations
993
- ----------------
994
- Write behavior is controlled by the `write_mode` parameter, which determines
995
- whether writes also persist objects in the cache.
996
-
997
- `write_mode` values:
998
- - `origin-and-cache`: objects are written both to the cache and to their
999
- intended origin bucket.
1000
- - `origin`: objects are written only to their intended origin bucket.
817
+ Specifies that this step should execute on DGX cloud.
1001
818
 
1002
819
 
1003
820
  Parameters
1004
821
  ----------
1005
- integration_name : str, optional
1006
- [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
1007
- that holds the configuration for the external, S3‑compatible object
1008
- storage bucket. If not specified, the only available S3 proxy
1009
- integration in the namespace is used (fails if multiple exist).
1010
- write_mode : str, optional
1011
- Controls whether writes also go to the external bucket.
1012
- - `origin` (default)
1013
- - `origin-and-cache`
1014
- debug : bool, optional
1015
- Enables debug logging for proxy operations.
822
+ gpu : int
823
+ Number of GPUs to use.
824
+ gpu_type : str
825
+ Type of Nvidia GPU to use.
1016
826
  """
1017
827
  ...
1018
828
 
1019
829
  @typing.overload
1020
- 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]]]:
830
+ 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]]]:
1021
831
  """
1022
- Specifies a timeout for your step.
1023
-
1024
- This decorator is useful if this step may hang indefinitely.
1025
-
1026
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1027
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1028
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
832
+ Specifies the Conda environment for the step.
1029
833
 
1030
- Note that all the values specified in parameters are added together so if you specify
1031
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
834
+ Information in this decorator will augment any
835
+ attributes set in the `@conda_base` flow-level decorator. Hence,
836
+ you can use `@conda_base` to set packages required by all
837
+ steps and use `@conda` to specify step-specific overrides.
1032
838
 
1033
839
 
1034
840
  Parameters
1035
841
  ----------
1036
- seconds : int, default 0
1037
- Number of seconds to wait prior to timing out.
1038
- minutes : int, default 0
1039
- Number of minutes to wait prior to timing out.
1040
- hours : int, default 0
1041
- Number of hours to wait prior to timing out.
842
+ packages : Dict[str, str], default {}
843
+ Packages to use for this step. The key is the name of the package
844
+ and the value is the version to use.
845
+ libraries : Dict[str, str], default {}
846
+ Supported for backward compatibility. When used with packages, packages will take precedence.
847
+ python : str, optional, default None
848
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
849
+ that the version used will correspond to the version of the Python interpreter used to start the run.
850
+ disabled : bool, default False
851
+ If set to True, disables @conda.
1042
852
  """
1043
853
  ...
1044
854
 
1045
855
  @typing.overload
1046
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
856
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1047
857
  ...
1048
858
 
1049
859
  @typing.overload
1050
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
860
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1051
861
  ...
1052
862
 
1053
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
863
+ 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):
1054
864
  """
1055
- Specifies a timeout for your step.
1056
-
1057
- This decorator is useful if this step may hang indefinitely.
1058
-
1059
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1060
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1061
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
865
+ Specifies the Conda environment for the step.
1062
866
 
1063
- Note that all the values specified in parameters are added together so if you specify
1064
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
867
+ Information in this decorator will augment any
868
+ attributes set in the `@conda_base` flow-level decorator. Hence,
869
+ you can use `@conda_base` to set packages required by all
870
+ steps and use `@conda` to specify step-specific overrides.
1065
871
 
1066
872
 
1067
873
  Parameters
1068
874
  ----------
1069
- seconds : int, default 0
1070
- Number of seconds to wait prior to timing out.
1071
- minutes : int, default 0
1072
- Number of minutes to wait prior to timing out.
1073
- hours : int, default 0
1074
- Number of hours to wait prior to timing out.
875
+ packages : Dict[str, str], default {}
876
+ Packages to use for this step. The key is the name of the package
877
+ and the value is the version to use.
878
+ libraries : Dict[str, str], default {}
879
+ Supported for backward compatibility. When used with packages, packages will take precedence.
880
+ python : str, optional, default None
881
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
882
+ that the version used will correspond to the version of the Python interpreter used to start the run.
883
+ disabled : bool, default False
884
+ If set to True, disables @conda.
1075
885
  """
1076
886
  ...
1077
887
 
1078
888
  @typing.overload
1079
- 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]]]:
889
+ 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]]]:
1080
890
  """
1081
- Creates a human-readable report, a Metaflow Card, after this step completes.
891
+ Enables checkpointing for a step.
1082
892
 
1083
- Note that you may add multiple `@card` decorators in a step with different parameters.
893
+ > Examples
894
+
895
+ - Saving Checkpoints
896
+
897
+ ```python
898
+ @checkpoint
899
+ @step
900
+ def train(self):
901
+ model = create_model(self.parameters, checkpoint_path = None)
902
+ for i in range(self.epochs):
903
+ # some training logic
904
+ loss = model.train(self.dataset)
905
+ if i % 10 == 0:
906
+ model.save(
907
+ current.checkpoint.directory,
908
+ )
909
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
910
+ # and returns a reference dictionary to the checkpoint saved in the datastore
911
+ self.latest_checkpoint = current.checkpoint.save(
912
+ name="epoch_checkpoint",
913
+ metadata={
914
+ "epoch": i,
915
+ "loss": loss,
916
+ }
917
+ )
918
+ ```
919
+
920
+ - Using Loaded Checkpoints
921
+
922
+ ```python
923
+ @retry(times=3)
924
+ @checkpoint
925
+ @step
926
+ def train(self):
927
+ # Assume that the task has restarted and the previous attempt of the task
928
+ # saved a checkpoint
929
+ checkpoint_path = None
930
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
931
+ print("Loaded checkpoint from the previous attempt")
932
+ checkpoint_path = current.checkpoint.directory
933
+
934
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
935
+ for i in range(self.epochs):
936
+ ...
937
+ ```
1084
938
 
1085
939
 
1086
940
  Parameters
1087
941
  ----------
1088
- type : str, default 'default'
1089
- Card type.
1090
- id : str, optional, default None
1091
- If multiple cards are present, use this id to identify this card.
1092
- options : Dict[str, Any], default {}
1093
- Options passed to the card. The contents depend on the card type.
1094
- timeout : int, default 45
1095
- Interrupt reporting if it takes more than this many seconds.
942
+ load_policy : str, default: "fresh"
943
+ The policy for loading the checkpoint. The following policies are supported:
944
+ - "eager": Loads the the latest available checkpoint within the namespace.
945
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
946
+ will be loaded at the start of the task.
947
+ - "none": Do not load any checkpoint
948
+ - "fresh": Loads the lastest checkpoint created within the running Task.
949
+ This mode helps loading checkpoints across various retry attempts of the same task.
950
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
951
+ created within the task will be loaded when the task is retries execution on failure.
952
+
953
+ temp_dir_root : str, default: None
954
+ The root directory under which `current.checkpoint.directory` will be created.
1096
955
  """
1097
956
  ...
1098
957
 
1099
958
  @typing.overload
1100
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
959
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1101
960
  ...
1102
961
 
1103
962
  @typing.overload
1104
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1105
- ...
1106
-
1107
- 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):
1108
- """
1109
- Creates a human-readable report, a Metaflow Card, after this step completes.
1110
-
1111
- Note that you may add multiple `@card` decorators in a step with different parameters.
1112
-
1113
-
1114
- Parameters
1115
- ----------
1116
- type : str, default 'default'
1117
- Card type.
1118
- id : str, optional, default None
1119
- If multiple cards are present, use this id to identify this card.
1120
- options : Dict[str, Any], default {}
1121
- Options passed to the card. The contents depend on the card type.
1122
- timeout : int, default 45
1123
- Interrupt reporting if it takes more than this many seconds.
1124
- """
963
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1125
964
  ...
1126
965
 
1127
- def nebius_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]]]:
966
+ 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):
1128
967
  """
1129
- `@nebius_s3_proxy` is a Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1130
- It exists to make it easier for users to know that this decorator should only be used with
1131
- a Neo Cloud like Nebius. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
1132
-
1133
-
1134
- Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
1135
- for S3 read and write requests.
968
+ Enables checkpointing for a step.
1136
969
 
1137
- This decorator requires an integration in the Outerbounds platform that
1138
- points to an external bucket. It affects S3 operations performed via
1139
- Metaflow's `get_aws_client` and `S3` within a `@step`.
970
+ > Examples
1140
971
 
1141
- Read operations
1142
- ---------------
1143
- All read operations pass through the proxy. If an object does not already
1144
- exist in the external bucket, it is cached there. For example, if code reads
1145
- from buckets `FOO` and `BAR` using the `S3` interface, objects from both
1146
- buckets are cached in the external bucket.
972
+ - Saving Checkpoints
1147
973
 
1148
- During task execution, all S3‑related read requests are routed through the
1149
- proxy:
1150
- - If the object is present in the external object store, the proxy
1151
- streams it directly from there without accessing the requested origin
1152
- bucket.
1153
- - If the object is not present in the external storage, the proxy
1154
- fetches it from the requested bucket, caches it in the external
1155
- storage, and streams the response from the origin bucket.
974
+ ```python
975
+ @checkpoint
976
+ @step
977
+ def train(self):
978
+ model = create_model(self.parameters, checkpoint_path = None)
979
+ for i in range(self.epochs):
980
+ # some training logic
981
+ loss = model.train(self.dataset)
982
+ if i % 10 == 0:
983
+ model.save(
984
+ current.checkpoint.directory,
985
+ )
986
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
987
+ # and returns a reference dictionary to the checkpoint saved in the datastore
988
+ self.latest_checkpoint = current.checkpoint.save(
989
+ name="epoch_checkpoint",
990
+ metadata={
991
+ "epoch": i,
992
+ "loss": loss,
993
+ }
994
+ )
995
+ ```
1156
996
 
1157
- Warning
1158
- -------
1159
- All READ operations (e.g., GetObject, HeadObject) pass through the external
1160
- bucket regardless of the bucket specified in user code. Even
1161
- `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
1162
- external bucket cache.
997
+ - Using Loaded Checkpoints
1163
998
 
1164
- Write operations
1165
- ----------------
1166
- Write behavior is controlled by the `write_mode` parameter, which determines
1167
- whether writes also persist objects in the cache.
999
+ ```python
1000
+ @retry(times=3)
1001
+ @checkpoint
1002
+ @step
1003
+ def train(self):
1004
+ # Assume that the task has restarted and the previous attempt of the task
1005
+ # saved a checkpoint
1006
+ checkpoint_path = None
1007
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1008
+ print("Loaded checkpoint from the previous attempt")
1009
+ checkpoint_path = current.checkpoint.directory
1168
1010
 
1169
- `write_mode` values:
1170
- - `origin-and-cache`: objects are written both to the cache and to their
1171
- intended origin bucket.
1172
- - `origin`: objects are written only to their intended origin bucket.
1011
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1012
+ for i in range(self.epochs):
1013
+ ...
1014
+ ```
1173
1015
 
1174
1016
 
1175
1017
  Parameters
1176
1018
  ----------
1177
- integration_name : str, optional
1178
- [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
1179
- that holds the configuration for the external, S3‑compatible object
1180
- storage bucket. If not specified, the only available S3 proxy
1181
- integration in the namespace is used (fails if multiple exist).
1182
- write_mode : str, optional
1183
- Controls whether writes also go to the external bucket.
1184
- - `origin` (default)
1185
- - `origin-and-cache`
1186
- debug : bool, optional
1187
- Enables debug logging for proxy operations.
1019
+ load_policy : str, default: "fresh"
1020
+ The policy for loading the checkpoint. The following policies are supported:
1021
+ - "eager": Loads the the latest available checkpoint within the namespace.
1022
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1023
+ will be loaded at the start of the task.
1024
+ - "none": Do not load any checkpoint
1025
+ - "fresh": Loads the lastest checkpoint created within the running Task.
1026
+ This mode helps loading checkpoints across various retry attempts of the same task.
1027
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1028
+ created within the task will be loaded when the task is retries execution on failure.
1029
+
1030
+ temp_dir_root : str, default: None
1031
+ The root directory under which `current.checkpoint.directory` will be created.
1188
1032
  """
1189
1033
  ...
1190
1034
 
@@ -1239,36 +1083,6 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
1239
1083
  """
1240
1084
  ...
1241
1085
 
1242
- 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]]]:
1243
- """
1244
- Specifies that this step should execute on DGX cloud.
1245
-
1246
-
1247
- Parameters
1248
- ----------
1249
- gpu : int
1250
- Number of GPUs to use.
1251
- gpu_type : str
1252
- Type of Nvidia GPU to use.
1253
- queue_timeout : int
1254
- Time to keep the job in NVCF's queue.
1255
- """
1256
- ...
1257
-
1258
- 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]]]:
1259
- """
1260
- Specifies that this step should execute on DGX cloud.
1261
-
1262
-
1263
- Parameters
1264
- ----------
1265
- gpu : int
1266
- Number of GPUs to use.
1267
- gpu_type : str
1268
- Type of Nvidia GPU to use.
1269
- """
1270
- ...
1271
-
1272
1086
  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]]]:
1273
1087
  """
1274
1088
  This decorator is used to run Ollama APIs as Metaflow task sidecars.
@@ -1312,6 +1126,25 @@ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy:
1312
1126
  """
1313
1127
  ...
1314
1128
 
1129
+ @typing.overload
1130
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1131
+ """
1132
+ Decorator prototype for all step decorators. This function gets specialized
1133
+ and imported for all decorators types by _import_plugin_decorators().
1134
+ """
1135
+ ...
1136
+
1137
+ @typing.overload
1138
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1139
+ ...
1140
+
1141
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1142
+ """
1143
+ Decorator prototype for all step decorators. This function gets specialized
1144
+ and imported for all decorators types by _import_plugin_decorators().
1145
+ """
1146
+ ...
1147
+
1315
1148
  @typing.overload
1316
1149
  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]]]:
1317
1150
  """
@@ -1363,157 +1196,428 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
1363
1196
  """
1364
1197
  ...
1365
1198
 
1366
- @typing.overload
1367
- 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]]]:
1199
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1368
1200
  """
1369
- Enables checkpointing for a step.
1370
-
1371
- > Examples
1201
+ Specifies that this step should execute on Kubernetes.
1372
1202
 
1373
- - Saving Checkpoints
1374
1203
 
1375
- ```python
1376
- @checkpoint
1204
+ Parameters
1205
+ ----------
1206
+ cpu : int, default 1
1207
+ Number of CPUs required for this step. If `@resources` is
1208
+ also present, the maximum value from all decorators is used.
1209
+ memory : int, default 4096
1210
+ Memory size (in MB) required for this step. If
1211
+ `@resources` is also present, the maximum value from all decorators is
1212
+ used.
1213
+ disk : int, default 10240
1214
+ Disk size (in MB) required for this step. If
1215
+ `@resources` is also present, the maximum value from all decorators is
1216
+ used.
1217
+ image : str, optional, default None
1218
+ Docker image to use when launching on Kubernetes. If not specified, and
1219
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
1220
+ not, a default Docker image mapping to the current version of Python is used.
1221
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
1222
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
1223
+ image_pull_secrets: List[str], default []
1224
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
1225
+ Kubernetes image pull secrets to use when pulling container images
1226
+ in Kubernetes.
1227
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
1228
+ Kubernetes service account to use when launching pod in Kubernetes.
1229
+ secrets : List[str], optional, default None
1230
+ Kubernetes secrets to use when launching pod in Kubernetes. These
1231
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
1232
+ in Metaflow configuration.
1233
+ node_selector: Union[Dict[str,str], str], optional, default None
1234
+ Kubernetes node selector(s) to apply to the pod running the task.
1235
+ Can be passed in as a comma separated string of values e.g.
1236
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
1237
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
1238
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
1239
+ Kubernetes namespace to use when launching pod in Kubernetes.
1240
+ gpu : int, optional, default None
1241
+ Number of GPUs required for this step. A value of zero implies that
1242
+ the scheduled node should not have GPUs.
1243
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
1244
+ The vendor of the GPUs to be used for this step.
1245
+ tolerations : List[Dict[str,str]], default []
1246
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
1247
+ Kubernetes tolerations to use when launching pod in Kubernetes.
1248
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
1249
+ Kubernetes labels to use when launching pod in Kubernetes.
1250
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
1251
+ Kubernetes annotations to use when launching pod in Kubernetes.
1252
+ use_tmpfs : bool, default False
1253
+ This enables an explicit tmpfs mount for this step.
1254
+ tmpfs_tempdir : bool, default True
1255
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
1256
+ tmpfs_size : int, optional, default: None
1257
+ The value for the size (in MiB) of the tmpfs mount for this step.
1258
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
1259
+ memory allocated for this step.
1260
+ tmpfs_path : str, optional, default /metaflow_temp
1261
+ Path to tmpfs mount for this step.
1262
+ persistent_volume_claims : Dict[str, str], optional, default None
1263
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
1264
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
1265
+ shared_memory: int, optional
1266
+ Shared memory size (in MiB) required for this step
1267
+ port: int, optional
1268
+ Port number to specify in the Kubernetes job object
1269
+ compute_pool : str, optional, default None
1270
+ Compute pool to be used for for this step.
1271
+ If not specified, any accessible compute pool within the perimeter is used.
1272
+ hostname_resolution_timeout: int, default 10 * 60
1273
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
1274
+ Only applicable when @parallel is used.
1275
+ qos: str, default: Burstable
1276
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
1277
+
1278
+ security_context: Dict[str, Any], optional, default None
1279
+ Container security context. Applies to the task container. Allows the following keys:
1280
+ - privileged: bool, optional, default None
1281
+ - allow_privilege_escalation: bool, optional, default None
1282
+ - run_as_user: int, optional, default None
1283
+ - run_as_group: int, optional, default None
1284
+ - run_as_non_root: bool, optional, default None
1285
+ """
1286
+ ...
1287
+
1288
+ @typing.overload
1289
+ 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]]]:
1290
+ """
1291
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1292
+
1293
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1294
+
1295
+
1296
+ Parameters
1297
+ ----------
1298
+ type : str, default 'default'
1299
+ Card type.
1300
+ id : str, optional, default None
1301
+ If multiple cards are present, use this id to identify this card.
1302
+ options : Dict[str, Any], default {}
1303
+ Options passed to the card. The contents depend on the card type.
1304
+ timeout : int, default 45
1305
+ Interrupt reporting if it takes more than this many seconds.
1306
+ """
1307
+ ...
1308
+
1309
+ @typing.overload
1310
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1311
+ ...
1312
+
1313
+ @typing.overload
1314
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1315
+ ...
1316
+
1317
+ 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):
1318
+ """
1319
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1320
+
1321
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1322
+
1323
+
1324
+ Parameters
1325
+ ----------
1326
+ type : str, default 'default'
1327
+ Card type.
1328
+ id : str, optional, default None
1329
+ If multiple cards are present, use this id to identify this card.
1330
+ options : Dict[str, Any], default {}
1331
+ Options passed to the card. The contents depend on the card type.
1332
+ timeout : int, default 45
1333
+ Interrupt reporting if it takes more than this many seconds.
1334
+ """
1335
+ ...
1336
+
1337
+ 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]]]:
1338
+ """
1339
+ Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
1340
+ for S3 read and write requests.
1341
+
1342
+ This decorator requires an integration in the Outerbounds platform that
1343
+ points to an external bucket. It affects S3 operations performed via
1344
+ Metaflow's `get_aws_client` and `S3` within a `@step`.
1345
+
1346
+ Read operations
1347
+ ---------------
1348
+ All read operations pass through the proxy. If an object does not already
1349
+ exist in the external bucket, it is cached there. For example, if code reads
1350
+ from buckets `FOO` and `BAR` using the `S3` interface, objects from both
1351
+ buckets are cached in the external bucket.
1352
+
1353
+ During task execution, all S3‑related read requests are routed through the
1354
+ proxy:
1355
+ - If the object is present in the external object store, the proxy
1356
+ streams it directly from there without accessing the requested origin
1357
+ bucket.
1358
+ - If the object is not present in the external storage, the proxy
1359
+ fetches it from the requested bucket, caches it in the external
1360
+ storage, and streams the response from the origin bucket.
1361
+
1362
+ Warning
1363
+ -------
1364
+ All READ operations (e.g., GetObject, HeadObject) pass through the external
1365
+ bucket regardless of the bucket specified in user code. Even
1366
+ `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
1367
+ external bucket cache.
1368
+
1369
+ Write operations
1370
+ ----------------
1371
+ Write behavior is controlled by the `write_mode` parameter, which determines
1372
+ whether writes also persist objects in the cache.
1373
+
1374
+ `write_mode` values:
1375
+ - `origin-and-cache`: objects are written both to the cache and to their
1376
+ intended origin bucket.
1377
+ - `origin`: objects are written only to their intended origin bucket.
1378
+
1379
+
1380
+ Parameters
1381
+ ----------
1382
+ integration_name : str, optional
1383
+ [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
1384
+ that holds the configuration for the external, S3‑compatible object
1385
+ storage bucket. If not specified, the only available S3 proxy
1386
+ integration in the namespace is used (fails if multiple exist).
1387
+ write_mode : str, optional
1388
+ Controls whether writes also go to the external bucket.
1389
+ - `origin` (default)
1390
+ - `origin-and-cache`
1391
+ debug : bool, optional
1392
+ Enables debug logging for proxy operations.
1393
+ """
1394
+ ...
1395
+
1396
+ 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]]]:
1397
+ """
1398
+ Decorator that helps cache, version, and store models/datasets from the Hugging Face Hub.
1399
+
1400
+ Examples
1401
+ --------
1402
+
1403
+ ```python
1404
+ # **Usage: creating references to models from the Hugging Face Hub that may be loaded in downstream steps**
1405
+ @huggingface_hub
1377
1406
  @step
1378
- def train(self):
1379
- model = create_model(self.parameters, checkpoint_path = None)
1380
- for i in range(self.epochs):
1381
- # some training logic
1382
- loss = model.train(self.dataset)
1383
- if i % 10 == 0:
1384
- model.save(
1385
- current.checkpoint.directory,
1386
- )
1387
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
1388
- # and returns a reference dictionary to the checkpoint saved in the datastore
1389
- self.latest_checkpoint = current.checkpoint.save(
1390
- name="epoch_checkpoint",
1391
- metadata={
1392
- "epoch": i,
1393
- "loss": loss,
1394
- }
1395
- )
1407
+ def pull_model_from_huggingface(self):
1408
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
1409
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
1410
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
1411
+ # value of the function is a reference to the model in the backend storage.
1412
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
1413
+
1414
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
1415
+ self.llama_model = current.huggingface_hub.snapshot_download(
1416
+ repo_id=self.model_id,
1417
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1418
+ )
1419
+ self.next(self.train)
1420
+
1421
+ # **Usage: explicitly loading models at runtime from the Hugging Face Hub or from cache (from Metaflow's datastore)**
1422
+ @huggingface_hub
1423
+ @step
1424
+ def run_training(self):
1425
+ # Temporary directory (auto-cleaned on exit)
1426
+ with current.huggingface_hub.load(
1427
+ repo_id="google-bert/bert-base-uncased",
1428
+ allow_patterns=["*.bin"],
1429
+ ) as local_path:
1430
+ # Use files under local_path
1431
+ train_model(local_path)
1432
+ ...
1433
+
1434
+ # **Usage: loading models directly from the Hugging Face Hub or from cache (from Metaflow's datastore)**
1435
+
1436
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1437
+ @step
1438
+ def pull_model_from_huggingface(self):
1439
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1440
+
1441
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora", "/my-lora-directory")])
1442
+ @step
1443
+ def finetune_model(self):
1444
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1445
+ # path_to_model will be /my-directory
1446
+
1447
+
1448
+ # Takes all the arguments passed to `snapshot_download`
1449
+ # except for `local_dir`
1450
+ @huggingface_hub(load=[
1451
+ {
1452
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1453
+ },
1454
+ {
1455
+ "repo_id": "myorg/mistral-lora",
1456
+ "repo_type": "model",
1457
+ },
1458
+ ])
1459
+ @step
1460
+ def finetune_model(self):
1461
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1462
+ # path_to_model will be /my-directory
1463
+ ```
1464
+
1465
+
1466
+ Parameters
1467
+ ----------
1468
+ temp_dir_root : str, optional
1469
+ The root directory that will hold the temporary directory where objects will be downloaded.
1470
+
1471
+ cache_scope : str, optional
1472
+ The scope of the cache. Can be `checkpoint` / `flow` / `global`.
1473
+ - `checkpoint` (default): All repos are stored like objects saved by `@checkpoint`.
1474
+ i.e., the cached path is derived from the namespace, flow, step, and Metaflow foreach iteration.
1475
+ 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).
1476
+
1477
+ - `flow`: All repos are cached under the flow, regardless of namespace.
1478
+ i.e., the cached path is derived solely from the flow name.
1479
+ 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.
1480
+
1481
+ - `global`: All repos are cached under a globally static path.
1482
+ i.e., the base path of the cache is static and all repos are stored under it.
1483
+ When to use this mode:
1484
+ - All repos from the Hugging Face Hub need to be shared by users across all flow executions.
1485
+ - Each caching scope comes with its own trade-offs:
1486
+ - `checkpoint`:
1487
+ - 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.
1488
+ - 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.
1489
+ - `flow`:
1490
+ - 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.
1491
+ - The blast radius of a bad checkpoint is limited to all runs of a particular flow.
1492
+ - It doesn't promote cache reuse across flows.
1493
+ - `global`:
1494
+ - Has no control over when caches are populated (can be written by any flow execution) but has the highest cache hit rate.
1495
+ - It promotes cache reuse across flows.
1496
+ - The blast radius of a bad checkpoint spans every flow that could be using a particular repo.
1497
+
1498
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1499
+ The list of repos (models/datasets) to load.
1500
+
1501
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1502
+
1503
+ - If repo (model/dataset) is not found in the datastore:
1504
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1505
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1506
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1507
+
1508
+ - If repo is found in the datastore:
1509
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
1510
+ """
1511
+ ...
1512
+
1513
+ @typing.overload
1514
+ 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]]:
1515
+ """
1516
+ Specifies the flow(s) that this flow depends on.
1517
+
1518
+ ```
1519
+ @trigger_on_finish(flow='FooFlow')
1396
1520
  ```
1521
+ or
1522
+ ```
1523
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1524
+ ```
1525
+ This decorator respects the @project decorator and triggers the flow
1526
+ when upstream runs within the same namespace complete successfully
1397
1527
 
1398
- - Using Loaded Checkpoints
1399
-
1400
- ```python
1401
- @retry(times=3)
1402
- @checkpoint
1403
- @step
1404
- def train(self):
1405
- # Assume that the task has restarted and the previous attempt of the task
1406
- # saved a checkpoint
1407
- checkpoint_path = None
1408
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1409
- print("Loaded checkpoint from the previous attempt")
1410
- checkpoint_path = current.checkpoint.directory
1528
+ Additionally, you can specify project aware upstream flow dependencies
1529
+ by specifying the fully qualified project_flow_name.
1530
+ ```
1531
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1532
+ ```
1533
+ or
1534
+ ```
1535
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1536
+ ```
1411
1537
 
1412
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1413
- for i in range(self.epochs):
1414
- ...
1538
+ You can also specify just the project or project branch (other values will be
1539
+ inferred from the current project or project branch):
1540
+ ```
1541
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1415
1542
  ```
1416
1543
 
1544
+ Note that `branch` is typically one of:
1545
+ - `prod`
1546
+ - `user.bob`
1547
+ - `test.my_experiment`
1548
+ - `prod.staging`
1549
+
1417
1550
 
1418
1551
  Parameters
1419
1552
  ----------
1420
- load_policy : str, default: "fresh"
1421
- The policy for loading the checkpoint. The following policies are supported:
1422
- - "eager": Loads the the latest available checkpoint within the namespace.
1423
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1424
- will be loaded at the start of the task.
1425
- - "none": Do not load any checkpoint
1426
- - "fresh": Loads the lastest checkpoint created within the running Task.
1427
- This mode helps loading checkpoints across various retry attempts of the same task.
1428
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1429
- created within the task will be loaded when the task is retries execution on failure.
1430
-
1431
- temp_dir_root : str, default: None
1432
- The root directory under which `current.checkpoint.directory` will be created.
1553
+ flow : Union[str, Dict[str, str]], optional, default None
1554
+ Upstream flow dependency for this flow.
1555
+ flows : List[Union[str, Dict[str, str]]], default []
1556
+ Upstream flow dependencies for this flow.
1557
+ options : Dict[str, Any], default {}
1558
+ Backend-specific configuration for tuning eventing behavior.
1433
1559
  """
1434
1560
  ...
1435
1561
 
1436
1562
  @typing.overload
1437
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1438
- ...
1439
-
1440
- @typing.overload
1441
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1563
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1442
1564
  ...
1443
1565
 
1444
- 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):
1566
+ 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] = {}):
1445
1567
  """
1446
- Enables checkpointing for a step.
1447
-
1448
- > Examples
1449
-
1450
- - Saving Checkpoints
1568
+ Specifies the flow(s) that this flow depends on.
1451
1569
 
1452
- ```python
1453
- @checkpoint
1454
- @step
1455
- def train(self):
1456
- model = create_model(self.parameters, checkpoint_path = None)
1457
- for i in range(self.epochs):
1458
- # some training logic
1459
- loss = model.train(self.dataset)
1460
- if i % 10 == 0:
1461
- model.save(
1462
- current.checkpoint.directory,
1463
- )
1464
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
1465
- # and returns a reference dictionary to the checkpoint saved in the datastore
1466
- self.latest_checkpoint = current.checkpoint.save(
1467
- name="epoch_checkpoint",
1468
- metadata={
1469
- "epoch": i,
1470
- "loss": loss,
1471
- }
1472
- )
1473
1570
  ```
1571
+ @trigger_on_finish(flow='FooFlow')
1572
+ ```
1573
+ or
1574
+ ```
1575
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1576
+ ```
1577
+ This decorator respects the @project decorator and triggers the flow
1578
+ when upstream runs within the same namespace complete successfully
1474
1579
 
1475
- - Using Loaded Checkpoints
1476
-
1477
- ```python
1478
- @retry(times=3)
1479
- @checkpoint
1480
- @step
1481
- def train(self):
1482
- # Assume that the task has restarted and the previous attempt of the task
1483
- # saved a checkpoint
1484
- checkpoint_path = None
1485
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1486
- print("Loaded checkpoint from the previous attempt")
1487
- checkpoint_path = current.checkpoint.directory
1580
+ Additionally, you can specify project aware upstream flow dependencies
1581
+ by specifying the fully qualified project_flow_name.
1582
+ ```
1583
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1584
+ ```
1585
+ or
1586
+ ```
1587
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1588
+ ```
1488
1589
 
1489
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1490
- for i in range(self.epochs):
1491
- ...
1590
+ You can also specify just the project or project branch (other values will be
1591
+ inferred from the current project or project branch):
1592
+ ```
1593
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1492
1594
  ```
1493
1595
 
1596
+ Note that `branch` is typically one of:
1597
+ - `prod`
1598
+ - `user.bob`
1599
+ - `test.my_experiment`
1600
+ - `prod.staging`
1601
+
1494
1602
 
1495
1603
  Parameters
1496
1604
  ----------
1497
- load_policy : str, default: "fresh"
1498
- The policy for loading the checkpoint. The following policies are supported:
1499
- - "eager": Loads the the latest available checkpoint within the namespace.
1500
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1501
- will be loaded at the start of the task.
1502
- - "none": Do not load any checkpoint
1503
- - "fresh": Loads the lastest checkpoint created within the running Task.
1504
- This mode helps loading checkpoints across various retry attempts of the same task.
1505
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1506
- created within the task will be loaded when the task is retries execution on failure.
1507
-
1508
- temp_dir_root : str, default: None
1509
- The root directory under which `current.checkpoint.directory` will be created.
1605
+ flow : Union[str, Dict[str, str]], optional, default None
1606
+ Upstream flow dependency for this flow.
1607
+ flows : List[Union[str, Dict[str, str]]], default []
1608
+ Upstream flow dependencies for this flow.
1609
+ options : Dict[str, Any], default {}
1610
+ Backend-specific configuration for tuning eventing behavior.
1510
1611
  """
1511
1612
  ...
1512
1613
 
1513
- 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]]:
1614
+ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1514
1615
  """
1515
- 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.
1516
- 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.
1616
+ The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1617
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1618
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1619
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1620
+ starts only after all sensors finish.
1517
1621
 
1518
1622
 
1519
1623
  Parameters
@@ -1535,21 +1639,18 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1535
1639
  Name of the sensor on Airflow
1536
1640
  description : str
1537
1641
  Description of sensor in the Airflow UI
1538
- external_dag_id : str
1539
- The dag_id that contains the task you want to wait for.
1540
- external_task_ids : List[str]
1541
- The list of task_ids that you want to wait for.
1542
- If None (default value) the sensor waits for the DAG. (Default: None)
1543
- allowed_states : List[str]
1544
- Iterable of allowed states, (Default: ['success'])
1545
- failed_states : List[str]
1546
- Iterable of failed or dis-allowed states. (Default: None)
1547
- execution_delta : datetime.timedelta
1548
- time difference with the previous execution to look at,
1549
- the default is the same logical date as the current task or DAG. (Default: None)
1550
- check_existence: bool
1551
- Set to True to check if the external task exists or check if
1552
- the DAG to wait for exists. (Default: True)
1642
+ bucket_key : Union[str, List[str]]
1643
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1644
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1645
+ bucket_name : str
1646
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1647
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1648
+ wildcard_match : bool
1649
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1650
+ aws_conn_id : str
1651
+ a reference to the s3 connection on Airflow. (Default: None)
1652
+ verify : bool
1653
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1553
1654
  """
1554
1655
  ...
1555
1656
 
@@ -1589,94 +1690,53 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1589
1690
  ...
1590
1691
 
1591
1692
  @typing.overload
1592
- 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]]:
1593
- """
1594
- Specifies the times when the flow should be run when running on a
1595
- production scheduler.
1596
-
1597
-
1598
- Parameters
1599
- ----------
1600
- hourly : bool, default False
1601
- Run the workflow hourly.
1602
- daily : bool, default True
1603
- Run the workflow daily.
1604
- weekly : bool, default False
1605
- Run the workflow weekly.
1606
- cron : str, optional, default None
1607
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1608
- specified by this expression.
1609
- timezone : str, optional, default None
1610
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1611
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1612
- """
1613
- ...
1614
-
1615
- @typing.overload
1616
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1617
- ...
1618
-
1619
- 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):
1693
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1620
1694
  """
1621
- Specifies the times when the flow should be run when running on a
1622
- production scheduler.
1623
-
1695
+ Specifies the Conda environment for all steps of the flow.
1624
1696
 
1625
- Parameters
1626
- ----------
1627
- hourly : bool, default False
1628
- Run the workflow hourly.
1629
- daily : bool, default True
1630
- Run the workflow daily.
1631
- weekly : bool, default False
1632
- Run the workflow weekly.
1633
- cron : str, optional, default None
1634
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1635
- specified by this expression.
1636
- timezone : str, optional, default None
1637
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1638
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1639
- """
1640
- ...
1641
-
1642
- @typing.overload
1643
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1644
- """
1645
- Specifies the PyPI packages for all steps of the flow.
1697
+ Use `@conda_base` to set common libraries required by all
1698
+ steps and use `@conda` to specify step-specific additions.
1646
1699
 
1647
- Use `@pypi_base` to set common packages required by all
1648
- steps and use `@pypi` to specify step-specific overrides.
1649
1700
 
1650
1701
  Parameters
1651
1702
  ----------
1652
- packages : Dict[str, str], default: {}
1703
+ packages : Dict[str, str], default {}
1653
1704
  Packages to use for this flow. The key is the name of the package
1654
1705
  and the value is the version to use.
1655
- python : str, optional, default: None
1706
+ libraries : Dict[str, str], default {}
1707
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1708
+ python : str, optional, default None
1656
1709
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1657
1710
  that the version used will correspond to the version of the Python interpreter used to start the run.
1711
+ disabled : bool, default False
1712
+ If set to True, disables Conda.
1658
1713
  """
1659
1714
  ...
1660
1715
 
1661
1716
  @typing.overload
1662
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1717
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1663
1718
  ...
1664
1719
 
1665
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1720
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1666
1721
  """
1667
- Specifies the PyPI packages for all steps of the flow.
1722
+ Specifies the Conda environment for all steps of the flow.
1723
+
1724
+ Use `@conda_base` to set common libraries required by all
1725
+ steps and use `@conda` to specify step-specific additions.
1668
1726
 
1669
- Use `@pypi_base` to set common packages required by all
1670
- steps and use `@pypi` to specify step-specific overrides.
1671
1727
 
1672
1728
  Parameters
1673
1729
  ----------
1674
- packages : Dict[str, str], default: {}
1730
+ packages : Dict[str, str], default {}
1675
1731
  Packages to use for this flow. The key is the name of the package
1676
1732
  and the value is the version to use.
1677
- python : str, optional, default: None
1733
+ libraries : Dict[str, str], default {}
1734
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1735
+ python : str, optional, default None
1678
1736
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1679
1737
  that the version used will correspond to the version of the Python interpreter used to start the run.
1738
+ disabled : bool, default False
1739
+ If set to True, disables Conda.
1680
1740
  """
1681
1741
  ...
1682
1742
 
@@ -1773,6 +1833,49 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1773
1833
  """
1774
1834
  ...
1775
1835
 
1836
+ 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]]:
1837
+ """
1838
+ 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.
1839
+ 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
+
1841
+
1842
+ Parameters
1843
+ ----------
1844
+ timeout : int
1845
+ Time, in seconds before the task times out and fails. (Default: 3600)
1846
+ poke_interval : int
1847
+ Time in seconds that the job should wait in between each try. (Default: 60)
1848
+ mode : str
1849
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1850
+ exponential_backoff : bool
1851
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1852
+ pool : str
1853
+ the slot pool this task should run in,
1854
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1855
+ soft_fail : bool
1856
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1857
+ name : str
1858
+ Name of the sensor on Airflow
1859
+ description : str
1860
+ Description of sensor in the Airflow UI
1861
+ external_dag_id : str
1862
+ The dag_id that contains the task you want to wait for.
1863
+ external_task_ids : List[str]
1864
+ The list of task_ids that you want to wait for.
1865
+ If None (default value) the sensor waits for the DAG. (Default: None)
1866
+ allowed_states : List[str]
1867
+ Iterable of allowed states, (Default: ['success'])
1868
+ failed_states : List[str]
1869
+ Iterable of failed or dis-allowed states. (Default: None)
1870
+ execution_delta : datetime.timedelta
1871
+ time difference with the previous execution to look at,
1872
+ the default is the same logical date as the current task or DAG. (Default: None)
1873
+ check_existence: bool
1874
+ Set to True to check if the external task exists or check if
1875
+ the DAG to wait for exists. (Default: True)
1876
+ """
1877
+ ...
1878
+
1776
1879
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1777
1880
  """
1778
1881
  Allows setting external datastores to save data for the
@@ -1888,197 +1991,94 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1888
1991
  ...
1889
1992
 
1890
1993
  @typing.overload
1891
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1994
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1892
1995
  """
1893
- Specifies the Conda environment for all steps of the flow.
1894
-
1895
- Use `@conda_base` to set common libraries required by all
1896
- steps and use `@conda` to specify step-specific additions.
1996
+ Specifies the PyPI packages for all steps of the flow.
1897
1997
 
1998
+ Use `@pypi_base` to set common packages required by all
1999
+ steps and use `@pypi` to specify step-specific overrides.
1898
2000
 
1899
2001
  Parameters
1900
2002
  ----------
1901
- packages : Dict[str, str], default {}
2003
+ packages : Dict[str, str], default: {}
1902
2004
  Packages to use for this flow. The key is the name of the package
1903
2005
  and the value is the version to use.
1904
- libraries : Dict[str, str], default {}
1905
- Supported for backward compatibility. When used with packages, packages will take precedence.
1906
- python : str, optional, default None
2006
+ python : str, optional, default: None
1907
2007
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1908
2008
  that the version used will correspond to the version of the Python interpreter used to start the run.
1909
- disabled : bool, default False
1910
- If set to True, disables Conda.
1911
2009
  """
1912
2010
  ...
1913
2011
 
1914
2012
  @typing.overload
1915
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
2013
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1916
2014
  ...
1917
2015
 
1918
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
2016
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1919
2017
  """
1920
- Specifies the Conda environment for all steps of the flow.
1921
-
1922
- Use `@conda_base` to set common libraries required by all
1923
- steps and use `@conda` to specify step-specific additions.
2018
+ Specifies the PyPI packages for all steps of the flow.
1924
2019
 
2020
+ Use `@pypi_base` to set common packages required by all
2021
+ steps and use `@pypi` to specify step-specific overrides.
1925
2022
 
1926
2023
  Parameters
1927
2024
  ----------
1928
- packages : Dict[str, str], default {}
2025
+ packages : Dict[str, str], default: {}
1929
2026
  Packages to use for this flow. The key is the name of the package
1930
2027
  and the value is the version to use.
1931
- libraries : Dict[str, str], default {}
1932
- Supported for backward compatibility. When used with packages, packages will take precedence.
1933
- python : str, optional, default None
2028
+ python : str, optional, default: None
1934
2029
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1935
2030
  that the version used will correspond to the version of the Python interpreter used to start the run.
1936
- disabled : bool, default False
1937
- If set to True, disables Conda.
1938
- """
1939
- ...
1940
-
1941
- def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, bucket_key: typing.Union[str, typing.List[str]], bucket_name: str, wildcard_match: bool, aws_conn_id: str, verify: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1942
- """
1943
- The `@airflow_s3_key_sensor` decorator attaches a Airflow [S3KeySensor](https://airflow.apache.org/docs/apache-airflow-providers-amazon/stable/_api/airflow/providers/amazon/aws/sensors/s3/index.html#airflow.providers.amazon.aws.sensors.s3.S3KeySensor)
1944
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1945
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1946
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1947
- starts only after all sensors finish.
1948
-
1949
-
1950
- Parameters
1951
- ----------
1952
- timeout : int
1953
- Time, in seconds before the task times out and fails. (Default: 3600)
1954
- poke_interval : int
1955
- Time in seconds that the job should wait in between each try. (Default: 60)
1956
- mode : str
1957
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1958
- exponential_backoff : bool
1959
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1960
- pool : str
1961
- the slot pool this task should run in,
1962
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1963
- soft_fail : bool
1964
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1965
- name : str
1966
- Name of the sensor on Airflow
1967
- description : str
1968
- Description of sensor in the Airflow UI
1969
- bucket_key : Union[str, List[str]]
1970
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1971
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1972
- bucket_name : str
1973
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1974
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1975
- wildcard_match : bool
1976
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1977
- aws_conn_id : str
1978
- a reference to the s3 connection on Airflow. (Default: None)
1979
- verify : bool
1980
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1981
2031
  """
1982
2032
  ...
1983
2033
 
1984
2034
  @typing.overload
1985
- 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]]:
2035
+ 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]]:
1986
2036
  """
1987
- Specifies the flow(s) that this flow depends on.
1988
-
1989
- ```
1990
- @trigger_on_finish(flow='FooFlow')
1991
- ```
1992
- or
1993
- ```
1994
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1995
- ```
1996
- This decorator respects the @project decorator and triggers the flow
1997
- when upstream runs within the same namespace complete successfully
1998
-
1999
- Additionally, you can specify project aware upstream flow dependencies
2000
- by specifying the fully qualified project_flow_name.
2001
- ```
2002
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
2003
- ```
2004
- or
2005
- ```
2006
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
2007
- ```
2008
-
2009
- You can also specify just the project or project branch (other values will be
2010
- inferred from the current project or project branch):
2011
- ```
2012
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
2013
- ```
2014
-
2015
- Note that `branch` is typically one of:
2016
- - `prod`
2017
- - `user.bob`
2018
- - `test.my_experiment`
2019
- - `prod.staging`
2037
+ Specifies the times when the flow should be run when running on a
2038
+ production scheduler.
2020
2039
 
2021
2040
 
2022
2041
  Parameters
2023
2042
  ----------
2024
- flow : Union[str, Dict[str, str]], optional, default None
2025
- Upstream flow dependency for this flow.
2026
- flows : List[Union[str, Dict[str, str]]], default []
2027
- Upstream flow dependencies for this flow.
2028
- options : Dict[str, Any], default {}
2029
- Backend-specific configuration for tuning eventing behavior.
2043
+ hourly : bool, default False
2044
+ Run the workflow hourly.
2045
+ daily : bool, default True
2046
+ Run the workflow daily.
2047
+ weekly : bool, default False
2048
+ Run the workflow weekly.
2049
+ cron : str, optional, default None
2050
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
2051
+ specified by this expression.
2052
+ timezone : str, optional, default None
2053
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
2054
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
2030
2055
  """
2031
2056
  ...
2032
2057
 
2033
2058
  @typing.overload
2034
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
2059
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
2035
2060
  ...
2036
2061
 
2037
- 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] = {}):
2062
+ 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):
2038
2063
  """
2039
- Specifies the flow(s) that this flow depends on.
2040
-
2041
- ```
2042
- @trigger_on_finish(flow='FooFlow')
2043
- ```
2044
- or
2045
- ```
2046
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
2047
- ```
2048
- This decorator respects the @project decorator and triggers the flow
2049
- when upstream runs within the same namespace complete successfully
2050
-
2051
- Additionally, you can specify project aware upstream flow dependencies
2052
- by specifying the fully qualified project_flow_name.
2053
- ```
2054
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
2055
- ```
2056
- or
2057
- ```
2058
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
2059
- ```
2060
-
2061
- You can also specify just the project or project branch (other values will be
2062
- inferred from the current project or project branch):
2063
- ```
2064
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
2065
- ```
2066
-
2067
- Note that `branch` is typically one of:
2068
- - `prod`
2069
- - `user.bob`
2070
- - `test.my_experiment`
2071
- - `prod.staging`
2064
+ Specifies the times when the flow should be run when running on a
2065
+ production scheduler.
2072
2066
 
2073
2067
 
2074
2068
  Parameters
2075
2069
  ----------
2076
- flow : Union[str, Dict[str, str]], optional, default None
2077
- Upstream flow dependency for this flow.
2078
- flows : List[Union[str, Dict[str, str]]], default []
2079
- Upstream flow dependencies for this flow.
2080
- options : Dict[str, Any], default {}
2081
- Backend-specific configuration for tuning eventing behavior.
2070
+ hourly : bool, default False
2071
+ Run the workflow hourly.
2072
+ daily : bool, default True
2073
+ Run the workflow daily.
2074
+ weekly : bool, default False
2075
+ Run the workflow weekly.
2076
+ cron : str, optional, default None
2077
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
2078
+ specified by this expression.
2079
+ timezone : str, optional, default None
2080
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
2081
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
2082
2082
  """
2083
2083
  ...
2084
2084