ob-metaflow-stubs 6.0.10.9__py2.py3-none-any.whl → 6.0.10.11__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 +856 -732
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +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 +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +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 +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +2 -2
  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 +1 -1
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +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 +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +3 -3
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +3 -3
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +3 -3
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  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 +2 -2
  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 +3 -3
  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 +3 -3
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +57 -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 +1 -1
  120. metaflow-stubs/packaging_sys/distribution_support.pyi +3 -3
  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 +13 -13
  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 +3 -3
  137. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  139. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  140. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  141. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  142. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  143. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  144. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  146. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  147. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  148. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  149. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  150. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  151. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  152. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  153. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  154. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  155. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +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 +2 -2
  161. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  162. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  163. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  164. metaflow-stubs/plugins/cards/__init__.pyi +5 -5
  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 +1 -1
  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 +4 -4
  184. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  185. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  186. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  187. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  188. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  189. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  190. metaflow-stubs/plugins/exit_hook/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +1 -1
  192. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  193. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  194. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  195. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  196. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  197. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  198. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  199. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  200. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  201. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +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 +2 -2
  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 +2 -2
  213. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  214. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  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 +1 -1
  223. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  224. metaflow-stubs/plugins/secrets/secrets_func.pyi +1 -1
  225. metaflow-stubs/plugins/secrets/secrets_spec.pyi +1 -1
  226. metaflow-stubs/plugins/secrets/utils.pyi +1 -1
  227. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  228. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  229. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  230. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  231. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  232. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  233. metaflow-stubs/plugins/uv/uv_environment.pyi +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 +4 -4
  238. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  239. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  240. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  241. metaflow-stubs/runner/nbrun.pyi +1 -1
  242. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  243. metaflow-stubs/runner/utils.pyi +3 -3
  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 +4 -4
  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 +2 -2
  255. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  256. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  257. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  258. {ob_metaflow_stubs-6.0.10.9.dist-info → ob_metaflow_stubs-6.0.10.11.dist-info}/METADATA +1 -1
  259. ob_metaflow_stubs-6.0.10.11.dist-info/RECORD +262 -0
  260. ob_metaflow_stubs-6.0.10.9.dist-info/RECORD +0 -262
  261. {ob_metaflow_stubs-6.0.10.9.dist-info → ob_metaflow_stubs-6.0.10.11.dist-info}/WHEEL +0 -0
  262. {ob_metaflow_stubs-6.0.10.9.dist-info → ob_metaflow_stubs-6.0.10.11.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.18.7.5+obcheckpoint(0.2.7);ob(v1) #
4
- # Generated on 2025-09-23T00:18:02.202668 #
4
+ # Generated on 2025-09-29T21:43:14.782191 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -40,17 +40,17 @@ from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
42
  from . import cards as cards
43
- from . import events as events
44
43
  from . import metaflow_git as metaflow_git
45
44
  from . import tuple_util as tuple_util
45
+ from . import events as events
46
46
  from . import runner as runner
47
47
  from . import plugins as plugins
48
48
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
49
49
  from . import includefile as includefile
50
50
  from .includefile import IncludeFile as IncludeFile
51
- from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
52
51
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
53
52
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
53
+ from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
54
54
  from . import client as client
55
55
  from .client.core import namespace as namespace
56
56
  from .client.core import get_namespace as get_namespace
@@ -77,6 +77,7 @@ from .mf_extensions.outerbounds.toplevel.s3_proxy import get_S3_with_s3_proxy as
77
77
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import set_s3_proxy_config as set_s3_proxy_config
78
78
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import clear_s3_proxy_config as clear_s3_proxy_config
79
79
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_s3_proxy_config as get_s3_proxy_config
80
+ from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_s3_proxy_config_from_env as get_s3_proxy_config_from_env
80
81
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import get_aws_client as get_aws_client
81
82
  from .mf_extensions.outerbounds.plugins.snowflake.snowflake import Snowflake as Snowflake
82
83
  from .mf_extensions.outerbounds.plugins.checkpoint_datastores.nebius import nebius_checkpoints as nebius_checkpoints
@@ -167,6 +168,61 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
167
168
  """
168
169
  ...
169
170
 
171
+ @typing.overload
172
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
173
+ """
174
+ Specifies the number of times the task corresponding
175
+ to a step needs to be retried.
176
+
177
+ This decorator is useful for handling transient errors, such as networking issues.
178
+ If your task contains operations that can't be retried safely, e.g. database updates,
179
+ it is advisable to annotate it with `@retry(times=0)`.
180
+
181
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
182
+ decorator will execute a no-op task after all retries have been exhausted,
183
+ ensuring that the flow execution can continue.
184
+
185
+
186
+ Parameters
187
+ ----------
188
+ times : int, default 3
189
+ Number of times to retry this task.
190
+ minutes_between_retries : int, default 2
191
+ Number of minutes between retries.
192
+ """
193
+ ...
194
+
195
+ @typing.overload
196
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
197
+ ...
198
+
199
+ @typing.overload
200
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
201
+ ...
202
+
203
+ 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):
204
+ """
205
+ Specifies the number of times the task corresponding
206
+ to a step needs to be retried.
207
+
208
+ This decorator is useful for handling transient errors, such as networking issues.
209
+ If your task contains operations that can't be retried safely, e.g. database updates,
210
+ it is advisable to annotate it with `@retry(times=0)`.
211
+
212
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
213
+ decorator will execute a no-op task after all retries have been exhausted,
214
+ ensuring that the flow execution can continue.
215
+
216
+
217
+ Parameters
218
+ ----------
219
+ times : int, default 3
220
+ Number of times to retry this task.
221
+ minutes_between_retries : int, default 2
222
+ Number of minutes between retries.
223
+ """
224
+ ...
225
+
170
226
  @typing.overload
171
227
  def test_append_card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
172
228
  """
@@ -186,46 +242,52 @@ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag]
186
242
  """
187
243
  ...
188
244
 
189
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
245
+ @typing.overload
246
+ 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]]]:
190
247
  """
191
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
248
+ Creates a human-readable report, a Metaflow Card, after this step completes.
192
249
 
193
- User code call
194
- --------------
195
- @ollama(
196
- models=[...],
197
- ...
198
- )
250
+ Note that you may add multiple `@card` decorators in a step with different parameters.
199
251
 
200
- Valid backend options
201
- ---------------------
202
- - 'local': Run as a separate process on the local task machine.
203
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
204
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
205
252
 
206
- Valid model options
207
- -------------------
208
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
253
+ Parameters
254
+ ----------
255
+ type : str, default 'default'
256
+ Card type.
257
+ id : str, optional, default None
258
+ If multiple cards are present, use this id to identify this card.
259
+ options : Dict[str, Any], default {}
260
+ Options passed to the card. The contents depend on the card type.
261
+ timeout : int, default 45
262
+ Interrupt reporting if it takes more than this many seconds.
263
+ """
264
+ ...
265
+
266
+ @typing.overload
267
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
268
+ ...
269
+
270
+ @typing.overload
271
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
272
+ ...
273
+
274
+ 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):
275
+ """
276
+ Creates a human-readable report, a Metaflow Card, after this step completes.
277
+
278
+ Note that you may add multiple `@card` decorators in a step with different parameters.
209
279
 
210
280
 
211
281
  Parameters
212
282
  ----------
213
- models: list[str]
214
- List of Ollama containers running models in sidecars.
215
- backend: str
216
- Determines where and how to run the Ollama process.
217
- force_pull: bool
218
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
219
- cache_update_policy: str
220
- Cache update policy: "auto", "force", or "never".
221
- force_cache_update: bool
222
- Simple override for "force" cache update policy.
223
- debug: bool
224
- Whether to turn on verbose debugging logs.
225
- circuit_breaker_config: dict
226
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
227
- timeout_config: dict
228
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
283
+ type : str, default 'default'
284
+ Card type.
285
+ id : str, optional, default None
286
+ If multiple cards are present, use this id to identify this card.
287
+ options : Dict[str, Any], default {}
288
+ Options passed to the card. The contents depend on the card type.
289
+ timeout : int, default 45
290
+ Interrupt reporting if it takes more than this many seconds.
229
291
  """
230
292
  ...
231
293
 
@@ -279,233 +341,417 @@ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card
279
341
  """
280
342
  ...
281
343
 
282
- @typing.overload
283
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
284
- """
285
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
286
- It exists to make it easier for users to know that this decorator should only be used with
287
- a Neo Cloud like CoreWeave.
288
- """
289
- ...
290
-
291
- @typing.overload
292
- def coreweave_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
293
- ...
294
-
295
- def coreweave_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
296
- """
297
- CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
298
- It exists to make it easier for users to know that this decorator should only be used with
299
- a Neo Cloud like CoreWeave.
300
- """
301
- ...
302
-
303
- @typing.overload
304
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
305
- """
306
- Decorator prototype for all step decorators. This function gets specialized
307
- and imported for all decorators types by _import_plugin_decorators().
308
- """
309
- ...
310
-
311
- @typing.overload
312
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
313
- ...
314
-
315
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
316
- """
317
- Decorator prototype for all step decorators. This function gets specialized
318
- and imported for all decorators types by _import_plugin_decorators().
319
- """
320
- ...
321
-
322
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
344
+ 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]]]:
323
345
  """
324
- Specifies that this step should execute on DGX cloud.
346
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
347
+
348
+ User code call
349
+ --------------
350
+ @ollama(
351
+ models=[...],
352
+ ...
353
+ )
354
+
355
+ Valid backend options
356
+ ---------------------
357
+ - 'local': Run as a separate process on the local task machine.
358
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
359
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
360
+
361
+ Valid model options
362
+ -------------------
363
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
325
364
 
326
365
 
327
366
  Parameters
328
367
  ----------
329
- gpu : int
330
- Number of GPUs to use.
331
- gpu_type : str
332
- Type of Nvidia GPU to use.
368
+ models: list[str]
369
+ List of Ollama containers running models in sidecars.
370
+ backend: str
371
+ Determines where and how to run the Ollama process.
372
+ force_pull: bool
373
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
374
+ cache_update_policy: str
375
+ Cache update policy: "auto", "force", or "never".
376
+ force_cache_update: bool
377
+ Simple override for "force" cache update policy.
378
+ debug: bool
379
+ Whether to turn on verbose debugging logs.
380
+ circuit_breaker_config: dict
381
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
382
+ timeout_config: dict
383
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
333
384
  """
334
385
  ...
335
386
 
336
387
  @typing.overload
337
- def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
388
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
338
389
  """
339
- Creates a human-readable report, a Metaflow Card, after this step completes.
390
+ Specifies the PyPI packages for the step.
340
391
 
341
- Note that you may add multiple `@card` decorators in a step with different parameters.
392
+ Information in this decorator will augment any
393
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
394
+ you can use `@pypi_base` to set packages required by all
395
+ steps and use `@pypi` to specify step-specific overrides.
342
396
 
343
397
 
344
398
  Parameters
345
399
  ----------
346
- type : str, default 'default'
347
- Card type.
348
- id : str, optional, default None
349
- If multiple cards are present, use this id to identify this card.
350
- options : Dict[str, Any], default {}
351
- Options passed to the card. The contents depend on the card type.
352
- timeout : int, default 45
353
- Interrupt reporting if it takes more than this many seconds.
400
+ packages : Dict[str, str], default: {}
401
+ Packages to use for this step. The key is the name of the package
402
+ and the value is the version to use.
403
+ python : str, optional, default: None
404
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
405
+ that the version used will correspond to the version of the Python interpreter used to start the run.
354
406
  """
355
407
  ...
356
408
 
357
409
  @typing.overload
358
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
410
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
359
411
  ...
360
412
 
361
413
  @typing.overload
362
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
414
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
363
415
  ...
364
416
 
365
- def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
417
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
366
418
  """
367
- Creates a human-readable report, a Metaflow Card, after this step completes.
419
+ Specifies the PyPI packages for the step.
368
420
 
369
- Note that you may add multiple `@card` decorators in a step with different parameters.
421
+ Information in this decorator will augment any
422
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
423
+ you can use `@pypi_base` to set packages required by all
424
+ steps and use `@pypi` to specify step-specific overrides.
370
425
 
371
426
 
372
427
  Parameters
373
428
  ----------
374
- type : str, default 'default'
375
- Card type.
376
- id : str, optional, default None
377
- If multiple cards are present, use this id to identify this card.
378
- options : Dict[str, Any], default {}
379
- Options passed to the card. The contents depend on the card type.
380
- timeout : int, default 45
381
- Interrupt reporting if it takes more than this many seconds.
429
+ packages : Dict[str, str], default: {}
430
+ Packages to use for this step. The key is the name of the package
431
+ and the value is the version to use.
432
+ python : str, optional, default: None
433
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
434
+ that the version used will correspond to the version of the Python interpreter used to start the run.
382
435
  """
383
436
  ...
384
437
 
385
- @typing.overload
386
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
438
+ 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]]]:
387
439
  """
388
- Specifies environment variables to be set prior to the execution of a step.
440
+ Specifies that this step should execute on Kubernetes.
389
441
 
390
442
 
391
443
  Parameters
392
444
  ----------
393
- vars : Dict[str, str], default {}
394
- Dictionary of environment variables to set.
445
+ cpu : int, default 1
446
+ Number of CPUs required for this step. If `@resources` is
447
+ also present, the maximum value from all decorators is used.
448
+ memory : int, default 4096
449
+ Memory size (in MB) required for this step. If
450
+ `@resources` is also present, the maximum value from all decorators is
451
+ used.
452
+ disk : int, default 10240
453
+ Disk size (in MB) required for this step. If
454
+ `@resources` is also present, the maximum value from all decorators is
455
+ used.
456
+ image : str, optional, default None
457
+ Docker image to use when launching on Kubernetes. If not specified, and
458
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
459
+ not, a default Docker image mapping to the current version of Python is used.
460
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
461
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
462
+ image_pull_secrets: List[str], default []
463
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
464
+ Kubernetes image pull secrets to use when pulling container images
465
+ in Kubernetes.
466
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
467
+ Kubernetes service account to use when launching pod in Kubernetes.
468
+ secrets : List[str], optional, default None
469
+ Kubernetes secrets to use when launching pod in Kubernetes. These
470
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
471
+ in Metaflow configuration.
472
+ node_selector: Union[Dict[str,str], str], optional, default None
473
+ Kubernetes node selector(s) to apply to the pod running the task.
474
+ Can be passed in as a comma separated string of values e.g.
475
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
476
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
477
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
478
+ Kubernetes namespace to use when launching pod in Kubernetes.
479
+ gpu : int, optional, default None
480
+ Number of GPUs required for this step. A value of zero implies that
481
+ the scheduled node should not have GPUs.
482
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
483
+ The vendor of the GPUs to be used for this step.
484
+ tolerations : List[Dict[str,str]], default []
485
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
486
+ Kubernetes tolerations to use when launching pod in Kubernetes.
487
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
488
+ Kubernetes labels to use when launching pod in Kubernetes.
489
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
490
+ Kubernetes annotations to use when launching pod in Kubernetes.
491
+ use_tmpfs : bool, default False
492
+ This enables an explicit tmpfs mount for this step.
493
+ tmpfs_tempdir : bool, default True
494
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
495
+ tmpfs_size : int, optional, default: None
496
+ The value for the size (in MiB) of the tmpfs mount for this step.
497
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
498
+ memory allocated for this step.
499
+ tmpfs_path : str, optional, default /metaflow_temp
500
+ Path to tmpfs mount for this step.
501
+ persistent_volume_claims : Dict[str, str], optional, default None
502
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
503
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
504
+ shared_memory: int, optional
505
+ Shared memory size (in MiB) required for this step
506
+ port: int, optional
507
+ Port number to specify in the Kubernetes job object
508
+ compute_pool : str, optional, default None
509
+ Compute pool to be used for for this step.
510
+ If not specified, any accessible compute pool within the perimeter is used.
511
+ hostname_resolution_timeout: int, default 10 * 60
512
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
513
+ Only applicable when @parallel is used.
514
+ qos: str, default: Burstable
515
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
516
+
517
+ security_context: Dict[str, Any], optional, default None
518
+ Container security context. Applies to the task container. Allows the following keys:
519
+ - privileged: bool, optional, default None
520
+ - allow_privilege_escalation: bool, optional, default None
521
+ - run_as_user: int, optional, default None
522
+ - run_as_group: int, optional, default None
523
+ - run_as_non_root: bool, optional, default None
395
524
  """
396
525
  ...
397
526
 
398
527
  @typing.overload
399
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
528
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], 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
+ """
400
533
  ...
401
534
 
402
535
  @typing.overload
403
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
536
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
404
537
  ...
405
538
 
406
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
539
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
407
540
  """
408
- Specifies environment variables to be set prior to the execution of a step.
409
-
410
-
411
- Parameters
412
- ----------
413
- vars : Dict[str, str], default {}
414
- Dictionary of environment variables to set.
541
+ Decorator prototype for all step decorators. This function gets specialized
542
+ and imported for all decorators types by _import_plugin_decorators().
415
543
  """
416
544
  ...
417
545
 
418
546
  @typing.overload
419
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
547
+ 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]]]:
420
548
  """
421
- Specifies the number of times the task corresponding
422
- to a step needs to be retried.
549
+ Enables checkpointing for a step.
423
550
 
424
- This decorator is useful for handling transient errors, such as networking issues.
425
- If your task contains operations that can't be retried safely, e.g. database updates,
426
- it is advisable to annotate it with `@retry(times=0)`.
551
+ > Examples
427
552
 
428
- This can be used in conjunction with the `@catch` decorator. The `@catch`
429
- decorator will execute a no-op task after all retries have been exhausted,
430
- ensuring that the flow execution can continue.
553
+ - Saving Checkpoints
554
+
555
+ ```python
556
+ @checkpoint
557
+ @step
558
+ def train(self):
559
+ model = create_model(self.parameters, checkpoint_path = None)
560
+ for i in range(self.epochs):
561
+ # some training logic
562
+ loss = model.train(self.dataset)
563
+ if i % 10 == 0:
564
+ model.save(
565
+ current.checkpoint.directory,
566
+ )
567
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
568
+ # and returns a reference dictionary to the checkpoint saved in the datastore
569
+ self.latest_checkpoint = current.checkpoint.save(
570
+ name="epoch_checkpoint",
571
+ metadata={
572
+ "epoch": i,
573
+ "loss": loss,
574
+ }
575
+ )
576
+ ```
577
+
578
+ - Using Loaded Checkpoints
579
+
580
+ ```python
581
+ @retry(times=3)
582
+ @checkpoint
583
+ @step
584
+ def train(self):
585
+ # Assume that the task has restarted and the previous attempt of the task
586
+ # saved a checkpoint
587
+ checkpoint_path = None
588
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
589
+ print("Loaded checkpoint from the previous attempt")
590
+ checkpoint_path = current.checkpoint.directory
591
+
592
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
593
+ for i in range(self.epochs):
594
+ ...
595
+ ```
431
596
 
432
597
 
433
598
  Parameters
434
599
  ----------
435
- times : int, default 3
436
- Number of times to retry this task.
437
- minutes_between_retries : int, default 2
438
- Number of minutes between retries.
600
+ load_policy : str, default: "fresh"
601
+ The policy for loading the checkpoint. The following policies are supported:
602
+ - "eager": Loads the the latest available checkpoint within the namespace.
603
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
604
+ will be loaded at the start of the task.
605
+ - "none": Do not load any checkpoint
606
+ - "fresh": Loads the lastest checkpoint created within the running Task.
607
+ This mode helps loading checkpoints across various retry attempts of the same task.
608
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
609
+ created within the task will be loaded when the task is retries execution on failure.
610
+
611
+ temp_dir_root : str, default: None
612
+ The root directory under which `current.checkpoint.directory` will be created.
439
613
  """
440
614
  ...
441
615
 
442
616
  @typing.overload
443
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
617
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
444
618
  ...
445
619
 
446
620
  @typing.overload
447
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
621
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
448
622
  ...
449
623
 
450
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
624
+ 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):
451
625
  """
452
- Specifies the number of times the task corresponding
453
- to a step needs to be retried.
626
+ Enables checkpointing for a step.
454
627
 
455
- This decorator is useful for handling transient errors, such as networking issues.
456
- If your task contains operations that can't be retried safely, e.g. database updates,
457
- it is advisable to annotate it with `@retry(times=0)`.
628
+ > Examples
458
629
 
459
- This can be used in conjunction with the `@catch` decorator. The `@catch`
460
- decorator will execute a no-op task after all retries have been exhausted,
461
- ensuring that the flow execution can continue.
630
+ - Saving Checkpoints
631
+
632
+ ```python
633
+ @checkpoint
634
+ @step
635
+ def train(self):
636
+ model = create_model(self.parameters, checkpoint_path = None)
637
+ for i in range(self.epochs):
638
+ # some training logic
639
+ loss = model.train(self.dataset)
640
+ if i % 10 == 0:
641
+ model.save(
642
+ current.checkpoint.directory,
643
+ )
644
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
645
+ # and returns a reference dictionary to the checkpoint saved in the datastore
646
+ self.latest_checkpoint = current.checkpoint.save(
647
+ name="epoch_checkpoint",
648
+ metadata={
649
+ "epoch": i,
650
+ "loss": loss,
651
+ }
652
+ )
653
+ ```
654
+
655
+ - Using Loaded Checkpoints
656
+
657
+ ```python
658
+ @retry(times=3)
659
+ @checkpoint
660
+ @step
661
+ def train(self):
662
+ # Assume that the task has restarted and the previous attempt of the task
663
+ # saved a checkpoint
664
+ checkpoint_path = None
665
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
666
+ print("Loaded checkpoint from the previous attempt")
667
+ checkpoint_path = current.checkpoint.directory
668
+
669
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
670
+ for i in range(self.epochs):
671
+ ...
672
+ ```
462
673
 
463
674
 
464
675
  Parameters
465
676
  ----------
466
- times : int, default 3
467
- Number of times to retry this task.
468
- minutes_between_retries : int, default 2
469
- Number of minutes between retries.
677
+ load_policy : str, default: "fresh"
678
+ The policy for loading the checkpoint. The following policies are supported:
679
+ - "eager": Loads the the latest available checkpoint within the namespace.
680
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
681
+ will be loaded at the start of the task.
682
+ - "none": Do not load any checkpoint
683
+ - "fresh": Loads the lastest checkpoint created within the running Task.
684
+ This mode helps loading checkpoints across various retry attempts of the same task.
685
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
686
+ created within the task will be loaded when the task is retries execution on failure.
687
+
688
+ temp_dir_root : str, default: None
689
+ The root directory under which `current.checkpoint.directory` will be created.
690
+ """
691
+ ...
692
+
693
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
694
+ """
695
+ Specifies that this step should execute on DGX cloud.
696
+
697
+
698
+ Parameters
699
+ ----------
700
+ gpu : int
701
+ Number of GPUs to use.
702
+ gpu_type : str
703
+ Type of Nvidia GPU to use.
470
704
  """
471
705
  ...
472
706
 
473
707
  @typing.overload
474
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
708
+ 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]]]:
475
709
  """
476
- Specifies secrets to be retrieved and injected as environment variables prior to
477
- the execution of a step.
710
+ Specifies that the step will success under all circumstances.
711
+
712
+ The decorator will create an optional artifact, specified by `var`, which
713
+ contains the exception raised. You can use it to detect the presence
714
+ of errors, indicating that all happy-path artifacts produced by the step
715
+ are missing.
478
716
 
479
717
 
480
718
  Parameters
481
719
  ----------
482
- sources : List[Union[str, Dict[str, Any]]], default: []
483
- List of secret specs, defining how the secrets are to be retrieved
484
- role : str, optional, default: None
485
- Role to use for fetching secrets
720
+ var : str, optional, default None
721
+ Name of the artifact in which to store the caught exception.
722
+ If not specified, the exception is not stored.
723
+ print_exception : bool, default True
724
+ Determines whether or not the exception is printed to
725
+ stdout when caught.
486
726
  """
487
727
  ...
488
728
 
489
729
  @typing.overload
490
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
730
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
491
731
  ...
492
732
 
493
733
  @typing.overload
494
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
734
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
495
735
  ...
496
736
 
497
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None):
737
+ 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):
498
738
  """
499
- Specifies secrets to be retrieved and injected as environment variables prior to
500
- the execution of a step.
739
+ Specifies that the step will success under all circumstances.
740
+
741
+ The decorator will create an optional artifact, specified by `var`, which
742
+ contains the exception raised. You can use it to detect the presence
743
+ of errors, indicating that all happy-path artifacts produced by the step
744
+ are missing.
501
745
 
502
746
 
503
747
  Parameters
504
748
  ----------
505
- sources : List[Union[str, Dict[str, Any]]], default: []
506
- List of secret specs, defining how the secrets are to be retrieved
507
- role : str, optional, default: None
508
- Role to use for fetching secrets
749
+ var : str, optional, default None
750
+ Name of the artifact in which to store the caught exception.
751
+ If not specified, the exception is not stored.
752
+ print_exception : bool, default True
753
+ Determines whether or not the exception is printed to
754
+ stdout when caught.
509
755
  """
510
756
  ...
511
757
 
@@ -569,58 +815,46 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
569
815
  ...
570
816
 
571
817
  @typing.overload
572
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
818
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
573
819
  """
574
- Specifies that the step will success under all circumstances.
575
-
576
- The decorator will create an optional artifact, specified by `var`, which
577
- contains the exception raised. You can use it to detect the presence
578
- of errors, indicating that all happy-path artifacts produced by the step
579
- are missing.
820
+ Specifies secrets to be retrieved and injected as environment variables prior to
821
+ the execution of a step.
580
822
 
581
823
 
582
824
  Parameters
583
825
  ----------
584
- var : str, optional, default None
585
- Name of the artifact in which to store the caught exception.
586
- If not specified, the exception is not stored.
587
- print_exception : bool, default True
588
- Determines whether or not the exception is printed to
589
- stdout when caught.
826
+ sources : List[Union[str, Dict[str, Any]]], default: []
827
+ List of secret specs, defining how the secrets are to be retrieved
828
+ role : str, optional, default: None
829
+ Role to use for fetching secrets
590
830
  """
591
831
  ...
592
832
 
593
833
  @typing.overload
594
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
834
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
595
835
  ...
596
836
 
597
837
  @typing.overload
598
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
838
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
599
839
  ...
600
840
 
601
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
841
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None):
602
842
  """
603
- Specifies that the step will success under all circumstances.
604
-
605
- The decorator will create an optional artifact, specified by `var`, which
606
- contains the exception raised. You can use it to detect the presence
607
- of errors, indicating that all happy-path artifacts produced by the step
608
- are missing.
843
+ Specifies secrets to be retrieved and injected as environment variables prior to
844
+ the execution of a step.
609
845
 
610
846
 
611
847
  Parameters
612
848
  ----------
613
- var : str, optional, default None
614
- Name of the artifact in which to store the caught exception.
615
- If not specified, the exception is not stored.
616
- print_exception : bool, default True
617
- Determines whether or not the exception is printed to
618
- stdout when caught.
849
+ sources : List[Union[str, Dict[str, Any]]], default: []
850
+ List of secret specs, defining how the secrets are to be retrieved
851
+ role : str, optional, default: None
852
+ Role to use for fetching secrets
619
853
  """
620
854
  ...
621
855
 
622
856
  @typing.overload
623
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
857
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
624
858
  """
625
859
  Decorator prototype for all step decorators. This function gets specialized
626
860
  and imported for all decorators types by _import_plugin_decorators().
@@ -628,64 +862,151 @@ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.
628
862
  ...
629
863
 
630
864
  @typing.overload
631
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
865
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
632
866
  ...
633
867
 
634
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
868
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
635
869
  """
636
870
  Decorator prototype for all step decorators. This function gets specialized
637
871
  and imported for all decorators types by _import_plugin_decorators().
638
872
  """
639
873
  ...
640
874
 
875
+ 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]]]:
876
+ """
877
+ Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
878
+ for S3 read and write requests.
879
+
880
+ This decorator requires an integration in the Outerbounds platform that
881
+ points to an external bucket. It affects S3 operations performed via
882
+ Metaflow's `get_aws_client` and `S3` within a `@step`.
883
+
884
+ Read operations
885
+ ---------------
886
+ All read operations pass through the proxy. If an object does not already
887
+ exist in the external bucket, it is cached there. For example, if code reads
888
+ from buckets `FOO` and `BAR` using the `S3` interface, objects from both
889
+ buckets are cached in the external bucket.
890
+
891
+ During task execution, all S3‑related read requests are routed through the
892
+ proxy:
893
+ - If the object is present in the external object store, the proxy
894
+ streams it directly from there without accessing the requested origin
895
+ bucket.
896
+ - If the object is not present in the external storage, the proxy
897
+ fetches it from the requested bucket, caches it in the external
898
+ storage, and streams the response from the origin bucket.
899
+
900
+ Warning
901
+ -------
902
+ All READ operations (e.g., GetObject, HeadObject) pass through the external
903
+ bucket regardless of the bucket specified in user code. Even
904
+ `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
905
+ external bucket cache.
906
+
907
+ Write operations
908
+ ----------------
909
+ Write behavior is controlled by the `write_mode` parameter, which determines
910
+ whether writes also persist objects in the cache.
911
+
912
+ `write_mode` values:
913
+ - `origin-and-cache`: objects are written both to the cache and to their
914
+ intended origin bucket.
915
+ - `origin`: objects are written only to their intended origin bucket.
916
+
917
+
918
+ Parameters
919
+ ----------
920
+ integration_name : str, optional
921
+ [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
922
+ that holds the configuration for the external, S3‑compatible object
923
+ storage bucket. If not specified, the only available S3 proxy
924
+ integration in the namespace is used (fails if multiple exist).
925
+ write_mode : str, optional
926
+ Controls whether writes also go to the external bucket.
927
+ - `origin` (default)
928
+ - `origin-and-cache`
929
+ debug : bool, optional
930
+ Enables debug logging for proxy operations.
931
+ """
932
+ ...
933
+
641
934
  @typing.overload
642
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
935
+ 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]]]:
643
936
  """
644
- Specifies the PyPI packages for the step.
937
+ Specifies the resources needed when executing this step.
645
938
 
646
- Information in this decorator will augment any
647
- attributes set in the `@pyi_base` flow-level decorator. Hence,
648
- you can use `@pypi_base` to set packages required by all
649
- steps and use `@pypi` to specify step-specific overrides.
939
+ Use `@resources` to specify the resource requirements
940
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
941
+
942
+ You can choose the compute layer on the command line by executing e.g.
943
+ ```
944
+ python myflow.py run --with batch
945
+ ```
946
+ or
947
+ ```
948
+ python myflow.py run --with kubernetes
949
+ ```
950
+ which executes the flow on the desired system using the
951
+ requirements specified in `@resources`.
650
952
 
651
953
 
652
954
  Parameters
653
955
  ----------
654
- packages : Dict[str, str], default: {}
655
- Packages to use for this step. The key is the name of the package
656
- and the value is the version to use.
657
- python : str, optional, default: None
658
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
659
- that the version used will correspond to the version of the Python interpreter used to start the run.
956
+ cpu : int, default 1
957
+ Number of CPUs required for this step.
958
+ gpu : int, optional, default None
959
+ Number of GPUs required for this step.
960
+ disk : int, optional, default None
961
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
962
+ memory : int, default 4096
963
+ Memory size (in MB) required for this step.
964
+ shared_memory : int, optional, default None
965
+ The value for the size (in MiB) of the /dev/shm volume for this step.
966
+ This parameter maps to the `--shm-size` option in Docker.
660
967
  """
661
968
  ...
662
969
 
663
970
  @typing.overload
664
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
971
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
665
972
  ...
666
973
 
667
974
  @typing.overload
668
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
975
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
669
976
  ...
670
977
 
671
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
978
+ 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):
672
979
  """
673
- Specifies the PyPI packages for the step.
980
+ Specifies the resources needed when executing this step.
674
981
 
675
- Information in this decorator will augment any
676
- attributes set in the `@pyi_base` flow-level decorator. Hence,
677
- you can use `@pypi_base` to set packages required by all
678
- steps and use `@pypi` to specify step-specific overrides.
982
+ Use `@resources` to specify the resource requirements
983
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
984
+
985
+ You can choose the compute layer on the command line by executing e.g.
986
+ ```
987
+ python myflow.py run --with batch
988
+ ```
989
+ or
990
+ ```
991
+ python myflow.py run --with kubernetes
992
+ ```
993
+ which executes the flow on the desired system using the
994
+ requirements specified in `@resources`.
679
995
 
680
996
 
681
997
  Parameters
682
998
  ----------
683
- packages : Dict[str, str], default: {}
684
- Packages to use for this step. The key is the name of the package
685
- and the value is the version to use.
686
- python : str, optional, default: None
687
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
688
- that the version used will correspond to the version of the Python interpreter used to start the run.
999
+ cpu : int, default 1
1000
+ Number of CPUs required for this step.
1001
+ gpu : int, optional, default None
1002
+ Number of GPUs required for this step.
1003
+ disk : int, optional, default None
1004
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1005
+ memory : int, default 4096
1006
+ Memory size (in MB) required for this step.
1007
+ shared_memory : int, optional, default None
1008
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1009
+ This parameter maps to the `--shm-size` option in Docker.
689
1010
  """
690
1011
  ...
691
1012
 
@@ -819,208 +1140,158 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
819
1140
  ...
820
1141
 
821
1142
  @typing.overload
822
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
823
- """
824
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
825
- It exists to make it easier for users to know that this decorator should only be used with
826
- a Neo Cloud like Nebius.
827
- """
828
- ...
829
-
830
- @typing.overload
831
- def nebius_s3_proxy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
832
- ...
833
-
834
- def nebius_s3_proxy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1143
+ 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]]]:
835
1144
  """
836
- Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
837
- It exists to make it easier for users to know that this decorator should only be used with
838
- a Neo Cloud like Nebius.
1145
+ Specifies environment variables to be set prior to the execution of a step.
1146
+
1147
+
1148
+ Parameters
1149
+ ----------
1150
+ vars : Dict[str, str], default {}
1151
+ Dictionary of environment variables to set.
839
1152
  """
840
1153
  ...
841
1154
 
842
1155
  @typing.overload
843
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
844
- """
845
- Internal decorator to support Fast bakery
846
- """
1156
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
847
1157
  ...
848
1158
 
849
1159
  @typing.overload
850
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1160
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
851
1161
  ...
852
1162
 
853
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1163
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
854
1164
  """
855
- Internal decorator to support Fast bakery
1165
+ Specifies environment variables to be set prior to the execution of a step.
1166
+
1167
+
1168
+ Parameters
1169
+ ----------
1170
+ vars : Dict[str, str], default {}
1171
+ Dictionary of environment variables to set.
856
1172
  """
857
1173
  ...
858
1174
 
859
- 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]]]:
1175
+ def coreweave_s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
860
1176
  """
861
- Specifies that this step should execute on Kubernetes.
1177
+ `@coreweave_s3_proxy` is a CoreWeave-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1178
+ It exists to make it easier for users to know that this decorator should only be used with
1179
+ a Neo Cloud like CoreWeave. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
1180
+
1181
+
1182
+ Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
1183
+ for S3 read and write requests.
1184
+
1185
+ This decorator requires an integration in the Outerbounds platform that
1186
+ points to an external bucket. It affects S3 operations performed via
1187
+ Metaflow's `get_aws_client` and `S3` within a `@step`.
1188
+
1189
+ Read operations
1190
+ ---------------
1191
+ All read operations pass through the proxy. If an object does not already
1192
+ exist in the external bucket, it is cached there. For example, if code reads
1193
+ from buckets `FOO` and `BAR` using the `S3` interface, objects from both
1194
+ buckets are cached in the external bucket.
1195
+
1196
+ During task execution, all S3‑related read requests are routed through the
1197
+ proxy:
1198
+ - If the object is present in the external object store, the proxy
1199
+ streams it directly from there without accessing the requested origin
1200
+ bucket.
1201
+ - If the object is not present in the external storage, the proxy
1202
+ fetches it from the requested bucket, caches it in the external
1203
+ storage, and streams the response from the origin bucket.
1204
+
1205
+ Warning
1206
+ -------
1207
+ All READ operations (e.g., GetObject, HeadObject) pass through the external
1208
+ bucket regardless of the bucket specified in user code. Even
1209
+ `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
1210
+ external bucket cache.
1211
+
1212
+ Write operations
1213
+ ----------------
1214
+ Write behavior is controlled by the `write_mode` parameter, which determines
1215
+ whether writes also persist objects in the cache.
1216
+
1217
+ `write_mode` values:
1218
+ - `origin-and-cache`: objects are written both to the cache and to their
1219
+ intended origin bucket.
1220
+ - `origin`: objects are written only to their intended origin bucket.
862
1221
 
863
1222
 
864
1223
  Parameters
865
1224
  ----------
866
- cpu : int, default 1
867
- Number of CPUs required for this step. If `@resources` is
868
- also present, the maximum value from all decorators is used.
869
- memory : int, default 4096
870
- Memory size (in MB) required for this step. If
871
- `@resources` is also present, the maximum value from all decorators is
872
- used.
873
- disk : int, default 10240
874
- Disk size (in MB) required for this step. If
875
- `@resources` is also present, the maximum value from all decorators is
876
- used.
877
- image : str, optional, default None
878
- Docker image to use when launching on Kubernetes. If not specified, and
879
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
880
- not, a default Docker image mapping to the current version of Python is used.
881
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
882
- If given, the imagePullPolicy to be applied to the Docker image of the step.
883
- image_pull_secrets: List[str], default []
884
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
885
- Kubernetes image pull secrets to use when pulling container images
886
- in Kubernetes.
887
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
888
- Kubernetes service account to use when launching pod in Kubernetes.
889
- secrets : List[str], optional, default None
890
- Kubernetes secrets to use when launching pod in Kubernetes. These
891
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
892
- in Metaflow configuration.
893
- node_selector: Union[Dict[str,str], str], optional, default None
894
- Kubernetes node selector(s) to apply to the pod running the task.
895
- Can be passed in as a comma separated string of values e.g.
896
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
897
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
898
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
899
- Kubernetes namespace to use when launching pod in Kubernetes.
900
- gpu : int, optional, default None
901
- Number of GPUs required for this step. A value of zero implies that
902
- the scheduled node should not have GPUs.
903
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
904
- The vendor of the GPUs to be used for this step.
905
- tolerations : List[Dict[str,str]], default []
906
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
907
- Kubernetes tolerations to use when launching pod in Kubernetes.
908
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
909
- Kubernetes labels to use when launching pod in Kubernetes.
910
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
911
- Kubernetes annotations to use when launching pod in Kubernetes.
912
- use_tmpfs : bool, default False
913
- This enables an explicit tmpfs mount for this step.
914
- tmpfs_tempdir : bool, default True
915
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
916
- tmpfs_size : int, optional, default: None
917
- The value for the size (in MiB) of the tmpfs mount for this step.
918
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
919
- memory allocated for this step.
920
- tmpfs_path : str, optional, default /metaflow_temp
921
- Path to tmpfs mount for this step.
922
- persistent_volume_claims : Dict[str, str], optional, default None
923
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
924
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
925
- shared_memory: int, optional
926
- Shared memory size (in MiB) required for this step
927
- port: int, optional
928
- Port number to specify in the Kubernetes job object
929
- compute_pool : str, optional, default None
930
- Compute pool to be used for for this step.
931
- If not specified, any accessible compute pool within the perimeter is used.
932
- hostname_resolution_timeout: int, default 10 * 60
933
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
934
- Only applicable when @parallel is used.
935
- qos: str, default: Burstable
936
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
937
-
938
- security_context: Dict[str, Any], optional, default None
939
- Container security context. Applies to the task container. Allows the following keys:
940
- - privileged: bool, optional, default None
941
- - allow_privilege_escalation: bool, optional, default None
942
- - run_as_user: int, optional, default None
943
- - run_as_group: int, optional, default None
944
- - run_as_non_root: bool, optional, default None
1225
+ integration_name : str, optional
1226
+ [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
1227
+ that holds the configuration for the external, S3‑compatible object
1228
+ storage bucket. If not specified, the only available S3 proxy
1229
+ integration in the namespace is used (fails if multiple exist).
1230
+ write_mode : str, optional
1231
+ Controls whether writes also go to the external bucket.
1232
+ - `origin` (default)
1233
+ - `origin-and-cache`
1234
+ debug : bool, optional
1235
+ Enables debug logging for proxy operations.
945
1236
  """
946
1237
  ...
947
1238
 
948
1239
  @typing.overload
949
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1240
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
950
1241
  """
951
- Specifies the resources needed when executing this step.
1242
+ Specifies a timeout for your step.
952
1243
 
953
- Use `@resources` to specify the resource requirements
954
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1244
+ This decorator is useful if this step may hang indefinitely.
955
1245
 
956
- You can choose the compute layer on the command line by executing e.g.
957
- ```
958
- python myflow.py run --with batch
959
- ```
960
- or
961
- ```
962
- python myflow.py run --with kubernetes
963
- ```
964
- which executes the flow on the desired system using the
965
- requirements specified in `@resources`.
1246
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1247
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1248
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1249
+
1250
+ Note that all the values specified in parameters are added together so if you specify
1251
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
966
1252
 
967
1253
 
968
1254
  Parameters
969
1255
  ----------
970
- cpu : int, default 1
971
- Number of CPUs required for this step.
972
- gpu : int, optional, default None
973
- Number of GPUs required for this step.
974
- disk : int, optional, default None
975
- Disk size (in MB) required for this step. Only applies on Kubernetes.
976
- memory : int, default 4096
977
- Memory size (in MB) required for this step.
978
- shared_memory : int, optional, default None
979
- The value for the size (in MiB) of the /dev/shm volume for this step.
980
- This parameter maps to the `--shm-size` option in Docker.
1256
+ seconds : int, default 0
1257
+ Number of seconds to wait prior to timing out.
1258
+ minutes : int, default 0
1259
+ Number of minutes to wait prior to timing out.
1260
+ hours : int, default 0
1261
+ Number of hours to wait prior to timing out.
981
1262
  """
982
1263
  ...
983
1264
 
984
1265
  @typing.overload
985
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1266
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
986
1267
  ...
987
1268
 
988
1269
  @typing.overload
989
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1270
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
990
1271
  ...
991
1272
 
992
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
1273
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
993
1274
  """
994
- Specifies the resources needed when executing this step.
1275
+ Specifies a timeout for your step.
995
1276
 
996
- Use `@resources` to specify the resource requirements
997
- independently of the specific compute layer (`@batch`, `@kubernetes`).
1277
+ This decorator is useful if this step may hang indefinitely.
998
1278
 
999
- You can choose the compute layer on the command line by executing e.g.
1000
- ```
1001
- python myflow.py run --with batch
1002
- ```
1003
- or
1004
- ```
1005
- python myflow.py run --with kubernetes
1006
- ```
1007
- which executes the flow on the desired system using the
1008
- requirements specified in `@resources`.
1279
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1280
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1281
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1282
+
1283
+ Note that all the values specified in parameters are added together so if you specify
1284
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1009
1285
 
1010
1286
 
1011
1287
  Parameters
1012
1288
  ----------
1013
- cpu : int, default 1
1014
- Number of CPUs required for this step.
1015
- gpu : int, optional, default None
1016
- Number of GPUs required for this step.
1017
- disk : int, optional, default None
1018
- Disk size (in MB) required for this step. Only applies on Kubernetes.
1019
- memory : int, default 4096
1020
- Memory size (in MB) required for this step.
1021
- shared_memory : int, optional, default None
1022
- The value for the size (in MiB) of the /dev/shm volume for this step.
1023
- This parameter maps to the `--shm-size` option in Docker.
1289
+ seconds : int, default 0
1290
+ Number of seconds to wait prior to timing out.
1291
+ minutes : int, default 0
1292
+ Number of minutes to wait prior to timing out.
1293
+ hours : int, default 0
1294
+ Number of hours to wait prior to timing out.
1024
1295
  """
1025
1296
  ...
1026
1297
 
@@ -1141,252 +1412,132 @@ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, cache_scope:
1141
1412
  """
1142
1413
  ...
1143
1414
 
1144
- @typing.overload
1145
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1415
+ 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]]]:
1146
1416
  """
1147
- Specifies a timeout for your step.
1148
-
1149
- This decorator is useful if this step may hang indefinitely.
1150
-
1151
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1152
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1153
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1154
-
1155
- Note that all the values specified in parameters are added together so if you specify
1156
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1417
+ Specifies that this step should execute on DGX cloud.
1157
1418
 
1158
1419
 
1159
1420
  Parameters
1160
1421
  ----------
1161
- seconds : int, default 0
1162
- Number of seconds to wait prior to timing out.
1163
- minutes : int, default 0
1164
- Number of minutes to wait prior to timing out.
1165
- hours : int, default 0
1166
- Number of hours to wait prior to timing out.
1422
+ gpu : int
1423
+ Number of GPUs to use.
1424
+ gpu_type : str
1425
+ Type of Nvidia GPU to use.
1426
+ queue_timeout : int
1427
+ Time to keep the job in NVCF's queue.
1167
1428
  """
1168
1429
  ...
1169
1430
 
1170
1431
  @typing.overload
1171
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1432
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1433
+ """
1434
+ Internal decorator to support Fast bakery
1435
+ """
1172
1436
  ...
1173
1437
 
1174
1438
  @typing.overload
1175
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1439
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1176
1440
  ...
1177
1441
 
1178
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
1442
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1179
1443
  """
1180
- Specifies a timeout for your step.
1181
-
1182
- This decorator is useful if this step may hang indefinitely.
1183
-
1184
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1185
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1186
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1187
-
1188
- Note that all the values specified in parameters are added together so if you specify
1189
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1190
-
1191
-
1192
- Parameters
1193
- ----------
1194
- seconds : int, default 0
1195
- Number of seconds to wait prior to timing out.
1196
- minutes : int, default 0
1197
- Number of minutes to wait prior to timing out.
1198
- hours : int, default 0
1199
- Number of hours to wait prior to timing out.
1444
+ Internal decorator to support Fast bakery
1200
1445
  """
1201
1446
  ...
1202
1447
 
1203
- @typing.overload
1204
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1448
+ def nebius_s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1205
1449
  """
1206
- Enables checkpointing for a step.
1207
-
1208
- > Examples
1209
-
1210
- - Saving Checkpoints
1211
-
1212
- ```python
1213
- @checkpoint
1214
- @step
1215
- def train(self):
1216
- model = create_model(self.parameters, checkpoint_path = None)
1217
- for i in range(self.epochs):
1218
- # some training logic
1219
- loss = model.train(self.dataset)
1220
- if i % 10 == 0:
1221
- model.save(
1222
- current.checkpoint.directory,
1223
- )
1224
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
1225
- # and returns a reference dictionary to the checkpoint saved in the datastore
1226
- self.latest_checkpoint = current.checkpoint.save(
1227
- name="epoch_checkpoint",
1228
- metadata={
1229
- "epoch": i,
1230
- "loss": loss,
1231
- }
1232
- )
1233
- ```
1234
-
1235
- - Using Loaded Checkpoints
1236
-
1237
- ```python
1238
- @retry(times=3)
1239
- @checkpoint
1240
- @step
1241
- def train(self):
1242
- # Assume that the task has restarted and the previous attempt of the task
1243
- # saved a checkpoint
1244
- checkpoint_path = None
1245
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1246
- print("Loaded checkpoint from the previous attempt")
1247
- checkpoint_path = current.checkpoint.directory
1248
-
1249
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1250
- for i in range(self.epochs):
1251
- ...
1252
- ```
1253
-
1450
+ `@nebius_s3_proxy` is a Nebius-specific S3 Proxy decorator for routing S3 requests through a local proxy service.
1451
+ It exists to make it easier for users to know that this decorator should only be used with
1452
+ a Neo Cloud like Nebius. The underlying mechanics of the decorator is the same as the `@s3_proxy`:
1254
1453
 
1255
- Parameters
1256
- ----------
1257
- load_policy : str, default: "fresh"
1258
- The policy for loading the checkpoint. The following policies are supported:
1259
- - "eager": Loads the the latest available checkpoint within the namespace.
1260
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1261
- will be loaded at the start of the task.
1262
- - "none": Do not load any checkpoint
1263
- - "fresh": Loads the lastest checkpoint created within the running Task.
1264
- This mode helps loading checkpoints across various retry attempts of the same task.
1265
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1266
- created within the task will be loaded when the task is retries execution on failure.
1267
1454
 
1268
- temp_dir_root : str, default: None
1269
- The root directory under which `current.checkpoint.directory` will be created.
1270
- """
1271
- ...
1272
-
1273
- @typing.overload
1274
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1275
- ...
1276
-
1277
- @typing.overload
1278
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1279
- ...
1280
-
1281
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
1282
- """
1283
- Enables checkpointing for a step.
1455
+ Set up an S3 proxy that caches objects in an external, S3‑compatible bucket
1456
+ for S3 read and write requests.
1284
1457
 
1285
- > Examples
1458
+ This decorator requires an integration in the Outerbounds platform that
1459
+ points to an external bucket. It affects S3 operations performed via
1460
+ Metaflow's `get_aws_client` and `S3` within a `@step`.
1286
1461
 
1287
- - Saving Checkpoints
1462
+ Read operations
1463
+ ---------------
1464
+ All read operations pass through the proxy. If an object does not already
1465
+ exist in the external bucket, it is cached there. For example, if code reads
1466
+ from buckets `FOO` and `BAR` using the `S3` interface, objects from both
1467
+ buckets are cached in the external bucket.
1288
1468
 
1289
- ```python
1290
- @checkpoint
1291
- @step
1292
- def train(self):
1293
- model = create_model(self.parameters, checkpoint_path = None)
1294
- for i in range(self.epochs):
1295
- # some training logic
1296
- loss = model.train(self.dataset)
1297
- if i % 10 == 0:
1298
- model.save(
1299
- current.checkpoint.directory,
1300
- )
1301
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
1302
- # and returns a reference dictionary to the checkpoint saved in the datastore
1303
- self.latest_checkpoint = current.checkpoint.save(
1304
- name="epoch_checkpoint",
1305
- metadata={
1306
- "epoch": i,
1307
- "loss": loss,
1308
- }
1309
- )
1310
- ```
1469
+ During task execution, all S3‑related read requests are routed through the
1470
+ proxy:
1471
+ - If the object is present in the external object store, the proxy
1472
+ streams it directly from there without accessing the requested origin
1473
+ bucket.
1474
+ - If the object is not present in the external storage, the proxy
1475
+ fetches it from the requested bucket, caches it in the external
1476
+ storage, and streams the response from the origin bucket.
1311
1477
 
1312
- - Using Loaded Checkpoints
1478
+ Warning
1479
+ -------
1480
+ All READ operations (e.g., GetObject, HeadObject) pass through the external
1481
+ bucket regardless of the bucket specified in user code. Even
1482
+ `S3(run=self)` and `S3(s3root="mybucketfoo")` requests go through the
1483
+ external bucket cache.
1313
1484
 
1314
- ```python
1315
- @retry(times=3)
1316
- @checkpoint
1317
- @step
1318
- def train(self):
1319
- # Assume that the task has restarted and the previous attempt of the task
1320
- # saved a checkpoint
1321
- checkpoint_path = None
1322
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
1323
- print("Loaded checkpoint from the previous attempt")
1324
- checkpoint_path = current.checkpoint.directory
1485
+ Write operations
1486
+ ----------------
1487
+ Write behavior is controlled by the `write_mode` parameter, which determines
1488
+ whether writes also persist objects in the cache.
1325
1489
 
1326
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
1327
- for i in range(self.epochs):
1328
- ...
1329
- ```
1490
+ `write_mode` values:
1491
+ - `origin-and-cache`: objects are written both to the cache and to their
1492
+ intended origin bucket.
1493
+ - `origin`: objects are written only to their intended origin bucket.
1330
1494
 
1331
1495
 
1332
1496
  Parameters
1333
1497
  ----------
1334
- load_policy : str, default: "fresh"
1335
- The policy for loading the checkpoint. The following policies are supported:
1336
- - "eager": Loads the the latest available checkpoint within the namespace.
1337
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
1338
- will be loaded at the start of the task.
1339
- - "none": Do not load any checkpoint
1340
- - "fresh": Loads the lastest checkpoint created within the running Task.
1341
- This mode helps loading checkpoints across various retry attempts of the same task.
1342
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
1343
- created within the task will be loaded when the task is retries execution on failure.
1344
-
1345
- temp_dir_root : str, default: None
1346
- The root directory under which `current.checkpoint.directory` will be created.
1498
+ integration_name : str, optional
1499
+ [Outerbounds integration name](https://docs.outerbounds.com/outerbounds/configuring-secrets/#integrations-view)
1500
+ that holds the configuration for the external, S3‑compatible object
1501
+ storage bucket. If not specified, the only available S3 proxy
1502
+ integration in the namespace is used (fails if multiple exist).
1503
+ write_mode : str, optional
1504
+ Controls whether writes also go to the external bucket.
1505
+ - `origin` (default)
1506
+ - `origin-and-cache`
1507
+ debug : bool, optional
1508
+ Enables debug logging for proxy operations.
1347
1509
  """
1348
1510
  ...
1349
1511
 
1350
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1512
+ @typing.overload
1513
+ 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]]:
1351
1514
  """
1352
- Specifies that this step should execute on DGX cloud.
1515
+ Specifies the times when the flow should be run when running on a
1516
+ production scheduler.
1353
1517
 
1354
1518
 
1355
1519
  Parameters
1356
1520
  ----------
1357
- gpu : int
1358
- Number of GPUs to use.
1359
- gpu_type : str
1360
- Type of Nvidia GPU to use.
1361
- queue_timeout : int
1362
- Time to keep the job in NVCF's queue.
1521
+ hourly : bool, default False
1522
+ Run the workflow hourly.
1523
+ daily : bool, default True
1524
+ Run the workflow daily.
1525
+ weekly : bool, default False
1526
+ Run the workflow weekly.
1527
+ cron : str, optional, default None
1528
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1529
+ specified by this expression.
1530
+ timezone : str, optional, default None
1531
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1532
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1363
1533
  """
1364
1534
  ...
1365
1535
 
1366
- def s3_proxy(*, integration_name: typing.Optional[str] = None, write_mode: typing.Optional[str] = None, debug: typing.Optional[bool] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1367
- """
1368
- S3 Proxy decorator for routing S3 requests through a local proxy service.
1369
-
1370
-
1371
- Parameters
1372
- ----------
1373
- integration_name : str, optional
1374
- Name of the S3 proxy integration. If not specified, will use the only
1375
- available S3 proxy integration in the namespace (fails if multiple exist).
1376
- write_mode : str, optional
1377
- The desired behavior during write operations to target (origin) S3 bucket.
1378
- allowed options are:
1379
- "origin-and-cache" -> write to both the target S3 bucket and local object
1380
- storage
1381
- "origin" -> only write to the target S3 bucket
1382
- "cache" -> only write to the object storage service used for caching
1383
- debug : bool, optional
1384
- Enable debug logging for proxy operations.
1385
- """
1536
+ @typing.overload
1537
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1386
1538
  ...
1387
1539
 
1388
- @typing.overload
1389
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1540
+ 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):
1390
1541
  """
1391
1542
  Specifies the times when the flow should be run when running on a
1392
1543
  production scheduler.
@@ -1410,72 +1561,78 @@ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False,
1410
1561
  ...
1411
1562
 
1412
1563
  @typing.overload
1413
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1564
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1565
+ """
1566
+ Specifies the PyPI packages for all steps of the flow.
1567
+
1568
+ Use `@pypi_base` to set common packages required by all
1569
+ steps and use `@pypi` to specify step-specific overrides.
1570
+
1571
+ Parameters
1572
+ ----------
1573
+ packages : Dict[str, str], default: {}
1574
+ Packages to use for this flow. The key is the name of the package
1575
+ and the value is the version to use.
1576
+ python : str, optional, default: None
1577
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1578
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1579
+ """
1414
1580
  ...
1415
1581
 
1416
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1582
+ @typing.overload
1583
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1584
+ ...
1585
+
1586
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1417
1587
  """
1418
- Specifies the times when the flow should be run when running on a
1419
- production scheduler.
1588
+ Specifies the PyPI packages for all steps of the flow.
1420
1589
 
1590
+ Use `@pypi_base` to set common packages required by all
1591
+ steps and use `@pypi` to specify step-specific overrides.
1421
1592
 
1422
1593
  Parameters
1423
1594
  ----------
1424
- hourly : bool, default False
1425
- Run the workflow hourly.
1426
- daily : bool, default True
1427
- Run the workflow daily.
1428
- weekly : bool, default False
1429
- Run the workflow weekly.
1430
- cron : str, optional, default None
1431
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1432
- specified by this expression.
1433
- timezone : str, optional, default None
1434
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1435
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1595
+ packages : Dict[str, str], default: {}
1596
+ Packages to use for this flow. The key is the name of the package
1597
+ and the value is the version to use.
1598
+ python : str, optional, default: None
1599
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1600
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1436
1601
  """
1437
1602
  ...
1438
1603
 
1439
- 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]]:
1604
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1440
1605
  """
1441
- 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)
1442
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1443
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1444
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1445
- starts only after all sensors finish.
1606
+ Specifies what flows belong to the same project.
1607
+
1608
+ A project-specific namespace is created for all flows that
1609
+ use the same `@project(name)`.
1446
1610
 
1447
1611
 
1448
1612
  Parameters
1449
1613
  ----------
1450
- timeout : int
1451
- Time, in seconds before the task times out and fails. (Default: 3600)
1452
- poke_interval : int
1453
- Time in seconds that the job should wait in between each try. (Default: 60)
1454
- mode : str
1455
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1456
- exponential_backoff : bool
1457
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1458
- pool : str
1459
- the slot pool this task should run in,
1460
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1461
- soft_fail : bool
1462
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1463
1614
  name : str
1464
- Name of the sensor on Airflow
1465
- description : str
1466
- Description of sensor in the Airflow UI
1467
- bucket_key : Union[str, List[str]]
1468
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1469
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1470
- bucket_name : str
1471
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1472
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1473
- wildcard_match : bool
1474
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1475
- aws_conn_id : str
1476
- a reference to the s3 connection on Airflow. (Default: None)
1477
- verify : bool
1478
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1615
+ Project name. Make sure that the name is unique amongst all
1616
+ projects that use the same production scheduler. The name may
1617
+ contain only lowercase alphanumeric characters and underscores.
1618
+
1619
+ branch : Optional[str], default None
1620
+ The branch to use. If not specified, the branch is set to
1621
+ `user.<username>` unless `production` is set to `True`. This can
1622
+ also be set on the command line using `--branch` as a top-level option.
1623
+ It is an error to specify `branch` in the decorator and on the command line.
1624
+
1625
+ production : bool, default False
1626
+ Whether or not the branch is the production branch. This can also be set on the
1627
+ command line using `--production` as a top-level option. It is an error to specify
1628
+ `production` in the decorator and on the command line.
1629
+ The project branch name will be:
1630
+ - if `branch` is specified:
1631
+ - if `production` is True: `prod.<branch>`
1632
+ - if `production` is False: `test.<branch>`
1633
+ - if `branch` is not specified:
1634
+ - if `production` is True: `prod`
1635
+ - if `production` is False: `user.<username>`
1479
1636
  """
1480
1637
  ...
1481
1638
 
@@ -1644,44 +1801,89 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1644
1801
  """
1645
1802
  ...
1646
1803
 
1647
- @typing.overload
1648
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1804
+ 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]]:
1649
1805
  """
1650
- Specifies the PyPI packages for all steps of the flow.
1806
+ 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)
1807
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1808
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1809
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1810
+ starts only after all sensors finish.
1651
1811
 
1652
- Use `@pypi_base` to set common packages required by all
1653
- steps and use `@pypi` to specify step-specific overrides.
1654
1812
 
1655
1813
  Parameters
1656
1814
  ----------
1657
- packages : Dict[str, str], default: {}
1658
- Packages to use for this flow. The key is the name of the package
1659
- and the value is the version to use.
1660
- python : str, optional, default: None
1661
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1662
- that the version used will correspond to the version of the Python interpreter used to start the run.
1815
+ timeout : int
1816
+ Time, in seconds before the task times out and fails. (Default: 3600)
1817
+ poke_interval : int
1818
+ Time in seconds that the job should wait in between each try. (Default: 60)
1819
+ mode : str
1820
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1821
+ exponential_backoff : bool
1822
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1823
+ pool : str
1824
+ the slot pool this task should run in,
1825
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1826
+ soft_fail : bool
1827
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1828
+ name : str
1829
+ Name of the sensor on Airflow
1830
+ description : str
1831
+ Description of sensor in the Airflow UI
1832
+ bucket_key : Union[str, List[str]]
1833
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1834
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1835
+ bucket_name : str
1836
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1837
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1838
+ wildcard_match : bool
1839
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1840
+ aws_conn_id : str
1841
+ a reference to the s3 connection on Airflow. (Default: None)
1842
+ verify : bool
1843
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1663
1844
  """
1664
1845
  ...
1665
1846
 
1666
- @typing.overload
1667
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1668
- ...
1669
-
1670
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1847
+ 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]]:
1671
1848
  """
1672
- Specifies the PyPI packages for all steps of the flow.
1849
+ 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.
1850
+ 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.
1673
1851
 
1674
- Use `@pypi_base` to set common packages required by all
1675
- steps and use `@pypi` to specify step-specific overrides.
1676
1852
 
1677
1853
  Parameters
1678
1854
  ----------
1679
- packages : Dict[str, str], default: {}
1680
- Packages to use for this flow. The key is the name of the package
1681
- and the value is the version to use.
1682
- python : str, optional, default: None
1683
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1684
- that the version used will correspond to the version of the Python interpreter used to start the run.
1855
+ timeout : int
1856
+ Time, in seconds before the task times out and fails. (Default: 3600)
1857
+ poke_interval : int
1858
+ Time in seconds that the job should wait in between each try. (Default: 60)
1859
+ mode : str
1860
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1861
+ exponential_backoff : bool
1862
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1863
+ pool : str
1864
+ the slot pool this task should run in,
1865
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1866
+ soft_fail : bool
1867
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1868
+ name : str
1869
+ Name of the sensor on Airflow
1870
+ description : str
1871
+ Description of sensor in the Airflow UI
1872
+ external_dag_id : str
1873
+ The dag_id that contains the task you want to wait for.
1874
+ external_task_ids : List[str]
1875
+ The list of task_ids that you want to wait for.
1876
+ If None (default value) the sensor waits for the DAG. (Default: None)
1877
+ allowed_states : List[str]
1878
+ Iterable of allowed states, (Default: ['success'])
1879
+ failed_states : List[str]
1880
+ Iterable of failed or dis-allowed states. (Default: None)
1881
+ execution_delta : datetime.timedelta
1882
+ time difference with the previous execution to look at,
1883
+ the default is the same logical date as the current task or DAG. (Default: None)
1884
+ check_existence: bool
1885
+ Set to True to check if the external task exists or check if
1886
+ the DAG to wait for exists. (Default: True)
1685
1887
  """
1686
1888
  ...
1687
1889
 
@@ -1879,83 +2081,5 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1879
2081
  """
1880
2082
  ...
1881
2083
 
1882
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1883
- """
1884
- Specifies what flows belong to the same project.
1885
-
1886
- A project-specific namespace is created for all flows that
1887
- use the same `@project(name)`.
1888
-
1889
-
1890
- Parameters
1891
- ----------
1892
- name : str
1893
- Project name. Make sure that the name is unique amongst all
1894
- projects that use the same production scheduler. The name may
1895
- contain only lowercase alphanumeric characters and underscores.
1896
-
1897
- branch : Optional[str], default None
1898
- The branch to use. If not specified, the branch is set to
1899
- `user.<username>` unless `production` is set to `True`. This can
1900
- also be set on the command line using `--branch` as a top-level option.
1901
- It is an error to specify `branch` in the decorator and on the command line.
1902
-
1903
- production : bool, default False
1904
- Whether or not the branch is the production branch. This can also be set on the
1905
- command line using `--production` as a top-level option. It is an error to specify
1906
- `production` in the decorator and on the command line.
1907
- The project branch name will be:
1908
- - if `branch` is specified:
1909
- - if `production` is True: `prod.<branch>`
1910
- - if `production` is False: `test.<branch>`
1911
- - if `branch` is not specified:
1912
- - if `production` is True: `prod`
1913
- - if `production` is False: `user.<username>`
1914
- """
1915
- ...
1916
-
1917
- def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str, exponential_backoff: bool, pool: str, soft_fail: bool, name: str, description: str, external_dag_id: str, external_task_ids: typing.List[str], allowed_states: typing.List[str], failed_states: typing.List[str], execution_delta: "datetime.timedelta", check_existence: bool) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1918
- """
1919
- The `@airflow_external_task_sensor` decorator attaches a Airflow [ExternalTaskSensor](https://airflow.apache.org/docs/apache-airflow/stable/_api/airflow/sensors/external_task/index.html#airflow.sensors.external_task.ExternalTaskSensor) before the start step of the flow.
1920
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1921
-
1922
-
1923
- Parameters
1924
- ----------
1925
- timeout : int
1926
- Time, in seconds before the task times out and fails. (Default: 3600)
1927
- poke_interval : int
1928
- Time in seconds that the job should wait in between each try. (Default: 60)
1929
- mode : str
1930
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1931
- exponential_backoff : bool
1932
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1933
- pool : str
1934
- the slot pool this task should run in,
1935
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1936
- soft_fail : bool
1937
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1938
- name : str
1939
- Name of the sensor on Airflow
1940
- description : str
1941
- Description of sensor in the Airflow UI
1942
- external_dag_id : str
1943
- The dag_id that contains the task you want to wait for.
1944
- external_task_ids : List[str]
1945
- The list of task_ids that you want to wait for.
1946
- If None (default value) the sensor waits for the DAG. (Default: None)
1947
- allowed_states : List[str]
1948
- Iterable of allowed states, (Default: ['success'])
1949
- failed_states : List[str]
1950
- Iterable of failed or dis-allowed states. (Default: None)
1951
- execution_delta : datetime.timedelta
1952
- time difference with the previous execution to look at,
1953
- the default is the same logical date as the current task or DAG. (Default: None)
1954
- check_existence: bool
1955
- Set to True to check if the external task exists or check if
1956
- the DAG to wait for exists. (Default: True)
1957
- """
1958
- ...
1959
-
1960
2084
  pkg_name: str
1961
2085