ob-metaflow-stubs 6.0.7.1__py2.py3-none-any.whl → 6.0.7.3__py2.py3-none-any.whl

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

Potentially problematic release.


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

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