ob-metaflow-stubs 6.0.10.20__py2.py3-none-any.whl → 6.0.10.21__py2.py3-none-any.whl

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