ob-metaflow-stubs 6.0.8.3__py2.py3-none-any.whl → 6.0.9.1__py2.py3-none-any.whl

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

Potentially problematic release.


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

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