ob-metaflow-stubs 6.0.7.2__py2.py3-none-any.whl → 6.0.7.4__py2.py3-none-any.whl

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