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