ob-metaflow-stubs 6.0.4.2__py2.py3-none-any.whl → 6.0.4.4rc0__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (260) hide show
  1. metaflow-stubs/__init__.pyi +863 -863
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +6 -6
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/meta_files.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +34 -34
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +4 -4
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +2 -2
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +5 -5
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +5 -5
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +3 -3
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +3 -3
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +4 -4
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/__init__.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_state_machine.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/__init__.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/_vendor/spinner/spinners.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_cli.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/app_config.pyi +2 -2
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/capsule.pyi +2 -2
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/click_importer.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/__init__.pyi +2 -2
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/code_package/code_packager.pyi +2 -2
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/__init__.pyi +2 -2
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/cli_generator.pyi +2 -2
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/config_utils.pyi +3 -3
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/schema_export.pyi +2 -2
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/typed_configs.pyi +3 -3
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/config/unified_config.pyi +3 -3
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/dependencies.pyi +3 -3
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/deployer.pyi +4 -4
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/experimental/__init__.pyi +2 -2
  87. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/perimeters.pyi +2 -2
  88. metaflow-stubs/mf_extensions/outerbounds/plugins/apps/core/utils.pyi +3 -3
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +3 -3
  91. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  92. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +3 -3
  93. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  94. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +2 -2
  95. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +3 -3
  96. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +3 -3
  97. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +2 -2
  98. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +3 -3
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +2 -2
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  110. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  111. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  112. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  115. metaflow-stubs/multicore_utils.pyi +2 -2
  116. metaflow-stubs/ob_internal.pyi +2 -2
  117. metaflow-stubs/packaging_sys/__init__.pyi +6 -6
  118. metaflow-stubs/packaging_sys/backend.pyi +3 -3
  119. metaflow-stubs/packaging_sys/distribution_support.pyi +4 -4
  120. metaflow-stubs/packaging_sys/tar_backend.pyi +5 -5
  121. metaflow-stubs/packaging_sys/utils.pyi +2 -2
  122. metaflow-stubs/packaging_sys/v1.pyi +2 -2
  123. metaflow-stubs/parameters.pyi +3 -3
  124. metaflow-stubs/plugins/__init__.pyi +14 -14
  125. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  126. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  127. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  128. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  129. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  130. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  131. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  132. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  133. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  134. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  135. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  136. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  137. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  138. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  139. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  140. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  141. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  142. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  143. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  144. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  145. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  146. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  147. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  148. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  149. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  150. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  151. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  152. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  153. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  154. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  155. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  156. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  157. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  158. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  159. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  160. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  161. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  162. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  163. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  164. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  165. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  166. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  167. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  169. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  170. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  171. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  172. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  173. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  174. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  175. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  176. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  177. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  178. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  179. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  180. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  181. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  182. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  183. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  184. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  185. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  186. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  187. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  188. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  189. metaflow-stubs/plugins/exit_hook/__init__.pyi +2 -2
  190. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +2 -2
  191. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  192. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  193. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  194. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  195. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  196. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  197. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  198. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  199. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  200. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  201. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  202. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  203. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  204. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  205. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  206. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  207. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  208. metaflow-stubs/plugins/perimeters.pyi +2 -2
  209. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  210. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  211. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  212. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  213. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  214. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  215. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  216. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  217. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  218. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  219. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  220. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  221. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  222. metaflow-stubs/plugins/secrets/secrets_func.pyi +2 -2
  223. metaflow-stubs/plugins/secrets/secrets_spec.pyi +2 -2
  224. metaflow-stubs/plugins/secrets/utils.pyi +2 -2
  225. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  226. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  227. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  228. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  229. metaflow-stubs/plugins/torchtune/__init__.pyi +2 -2
  230. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  231. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  232. metaflow-stubs/profilers/__init__.pyi +2 -2
  233. metaflow-stubs/pylint_wrapper.pyi +2 -2
  234. metaflow-stubs/runner/__init__.pyi +2 -2
  235. metaflow-stubs/runner/deployer.pyi +28 -28
  236. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  237. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  238. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  239. metaflow-stubs/runner/nbrun.pyi +2 -2
  240. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  241. metaflow-stubs/runner/utils.pyi +2 -2
  242. metaflow-stubs/system/__init__.pyi +2 -2
  243. metaflow-stubs/system/system_logger.pyi +2 -2
  244. metaflow-stubs/system/system_monitor.pyi +2 -2
  245. metaflow-stubs/tagging_util.pyi +2 -2
  246. metaflow-stubs/tuple_util.pyi +2 -2
  247. metaflow-stubs/user_configs/__init__.pyi +2 -2
  248. metaflow-stubs/user_configs/config_options.pyi +2 -2
  249. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  250. metaflow-stubs/user_decorators/__init__.pyi +2 -2
  251. metaflow-stubs/user_decorators/common.pyi +2 -2
  252. metaflow-stubs/user_decorators/mutable_flow.pyi +6 -6
  253. metaflow-stubs/user_decorators/mutable_step.pyi +5 -5
  254. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  255. metaflow-stubs/user_decorators/user_step_decorator.pyi +5 -5
  256. {ob_metaflow_stubs-6.0.4.2.dist-info → ob_metaflow_stubs-6.0.4.4rc0.dist-info}/METADATA +1 -1
  257. ob_metaflow_stubs-6.0.4.4rc0.dist-info/RECORD +260 -0
  258. ob_metaflow_stubs-6.0.4.2.dist-info/RECORD +0 -260
  259. {ob_metaflow_stubs-6.0.4.2.dist-info → ob_metaflow_stubs-6.0.4.4rc0.dist-info}/WHEEL +0 -0
  260. {ob_metaflow_stubs-6.0.4.2.dist-info → ob_metaflow_stubs-6.0.4.4rc0.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.16.0.1+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-07-15T03:12:46.861592 #
3
+ # MF version: 2.16.1.1+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-07-15T21:03:18.366830 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -39,9 +39,9 @@ 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
44
+ from . import tuple_util as tuple_util
45
45
  from . import events as events
46
46
  from . import runner as runner
47
47
  from . import plugins as plugins
@@ -162,452 +162,460 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
162
162
  """
163
163
  ...
164
164
 
165
- @typing.overload
166
- def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
165
+ 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]]]:
167
166
  """
168
- Specifies a timeout for your step.
167
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
169
168
 
170
- This decorator is useful if this step may hang indefinitely.
169
+ User code call
170
+ --------------
171
+ @ollama(
172
+ models=[...],
173
+ ...
174
+ )
171
175
 
172
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
173
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
174
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
176
+ Valid backend options
177
+ ---------------------
178
+ - 'local': Run as a separate process on the local task machine.
179
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
180
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
175
181
 
176
- Note that all the values specified in parameters are added together so if you specify
177
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
182
+ Valid model options
183
+ -------------------
184
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
178
185
 
179
186
 
180
187
  Parameters
181
188
  ----------
182
- seconds : int, default 0
183
- Number of seconds to wait prior to timing out.
184
- minutes : int, default 0
185
- Number of minutes to wait prior to timing out.
186
- hours : int, default 0
187
- Number of hours to wait prior to timing out.
189
+ models: list[str]
190
+ List of Ollama containers running models in sidecars.
191
+ backend: str
192
+ Determines where and how to run the Ollama process.
193
+ force_pull: bool
194
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
195
+ cache_update_policy: str
196
+ Cache update policy: "auto", "force", or "never".
197
+ force_cache_update: bool
198
+ Simple override for "force" cache update policy.
199
+ debug: bool
200
+ Whether to turn on verbose debugging logs.
201
+ circuit_breaker_config: dict
202
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
203
+ timeout_config: dict
204
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
188
205
  """
189
206
  ...
190
207
 
191
208
  @typing.overload
192
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
209
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
210
+ """
211
+ Specifies the number of times the task corresponding
212
+ to a step needs to be retried.
213
+
214
+ This decorator is useful for handling transient errors, such as networking issues.
215
+ If your task contains operations that can't be retried safely, e.g. database updates,
216
+ it is advisable to annotate it with `@retry(times=0)`.
217
+
218
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
219
+ decorator will execute a no-op task after all retries have been exhausted,
220
+ ensuring that the flow execution can continue.
221
+
222
+
223
+ Parameters
224
+ ----------
225
+ times : int, default 3
226
+ Number of times to retry this task.
227
+ minutes_between_retries : int, default 2
228
+ Number of minutes between retries.
229
+ """
193
230
  ...
194
231
 
195
232
  @typing.overload
196
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
233
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
197
234
  ...
198
235
 
199
- def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
236
+ @typing.overload
237
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
238
+ ...
239
+
240
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
200
241
  """
201
- Specifies a timeout for your step.
242
+ Specifies the number of times the task corresponding
243
+ to a step needs to be retried.
202
244
 
203
- This decorator is useful if this step may hang indefinitely.
245
+ This decorator is useful for handling transient errors, such as networking issues.
246
+ If your task contains operations that can't be retried safely, e.g. database updates,
247
+ it is advisable to annotate it with `@retry(times=0)`.
204
248
 
205
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
206
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
207
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
249
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
250
+ decorator will execute a no-op task after all retries have been exhausted,
251
+ ensuring that the flow execution can continue.
208
252
 
209
- Note that all the values specified in parameters are added together so if you specify
210
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
253
+
254
+ Parameters
255
+ ----------
256
+ times : int, default 3
257
+ Number of times to retry this task.
258
+ minutes_between_retries : int, default 2
259
+ Number of minutes between retries.
260
+ """
261
+ ...
262
+
263
+ 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]]]:
264
+ """
265
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
266
+
267
+ User code call
268
+ --------------
269
+ @vllm(
270
+ model="...",
271
+ ...
272
+ )
273
+
274
+ Valid backend options
275
+ ---------------------
276
+ - 'local': Run as a separate process on the local task machine.
277
+
278
+ Valid model options
279
+ -------------------
280
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
281
+
282
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
283
+ If you need multiple models, you must create multiple @vllm decorators.
211
284
 
212
285
 
213
286
  Parameters
214
287
  ----------
215
- seconds : int, default 0
216
- Number of seconds to wait prior to timing out.
217
- minutes : int, default 0
218
- Number of minutes to wait prior to timing out.
219
- hours : int, default 0
220
- Number of hours to wait prior to timing out.
288
+ model: str
289
+ HuggingFace model identifier to be served by vLLM.
290
+ backend: str
291
+ Determines where and how to run the vLLM process.
292
+ openai_api_server: bool
293
+ Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
294
+ Default is False (uses native engine).
295
+ Set to True for backward compatibility with existing code.
296
+ debug: bool
297
+ Whether to turn on verbose debugging logs.
298
+ card_refresh_interval: int
299
+ Interval in seconds for refreshing the vLLM status card.
300
+ Only used when openai_api_server=True.
301
+ max_retries: int
302
+ Maximum number of retries checking for vLLM server startup.
303
+ Only used when openai_api_server=True.
304
+ retry_alert_frequency: int
305
+ Frequency of alert logs for vLLM server startup retries.
306
+ Only used when openai_api_server=True.
307
+ engine_args : dict
308
+ Additional keyword arguments to pass to the vLLM engine.
309
+ For example, `tensor_parallel_size=2`.
221
310
  """
222
311
  ...
223
312
 
224
313
  @typing.overload
225
- 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]]]:
314
+ 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]]]:
226
315
  """
227
- Specifies the resources needed when executing this step.
228
-
229
- Use `@resources` to specify the resource requirements
230
- independently of the specific compute layer (`@batch`, `@kubernetes`).
316
+ Specifies the Conda environment for the step.
231
317
 
232
- You can choose the compute layer on the command line by executing e.g.
233
- ```
234
- python myflow.py run --with batch
235
- ```
236
- or
237
- ```
238
- python myflow.py run --with kubernetes
239
- ```
240
- which executes the flow on the desired system using the
241
- requirements specified in `@resources`.
318
+ Information in this decorator will augment any
319
+ attributes set in the `@conda_base` flow-level decorator. Hence,
320
+ you can use `@conda_base` to set packages required by all
321
+ steps and use `@conda` to specify step-specific overrides.
242
322
 
243
323
 
244
324
  Parameters
245
325
  ----------
246
- cpu : int, default 1
247
- Number of CPUs required for this step.
248
- gpu : int, optional, default None
249
- Number of GPUs required for this step.
250
- disk : int, optional, default None
251
- Disk size (in MB) required for this step. Only applies on Kubernetes.
252
- memory : int, default 4096
253
- Memory size (in MB) required for this step.
254
- shared_memory : int, optional, default None
255
- The value for the size (in MiB) of the /dev/shm volume for this step.
256
- This parameter maps to the `--shm-size` option in Docker.
326
+ packages : Dict[str, str], default {}
327
+ Packages to use for this step. The key is the name of the package
328
+ and the value is the version to use.
329
+ libraries : Dict[str, str], default {}
330
+ Supported for backward compatibility. When used with packages, packages will take precedence.
331
+ python : str, optional, default None
332
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
333
+ that the version used will correspond to the version of the Python interpreter used to start the run.
334
+ disabled : bool, default False
335
+ If set to True, disables @conda.
257
336
  """
258
337
  ...
259
338
 
260
339
  @typing.overload
261
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
340
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
262
341
  ...
263
342
 
264
343
  @typing.overload
265
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
344
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
266
345
  ...
267
346
 
268
- 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):
347
+ 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):
269
348
  """
270
- Specifies the resources needed when executing this step.
271
-
272
- Use `@resources` to specify the resource requirements
273
- independently of the specific compute layer (`@batch`, `@kubernetes`).
349
+ Specifies the Conda environment for the step.
274
350
 
275
- You can choose the compute layer on the command line by executing e.g.
276
- ```
277
- python myflow.py run --with batch
278
- ```
279
- or
280
- ```
281
- python myflow.py run --with kubernetes
282
- ```
283
- which executes the flow on the desired system using the
284
- requirements specified in `@resources`.
351
+ Information in this decorator will augment any
352
+ attributes set in the `@conda_base` flow-level decorator. Hence,
353
+ you can use `@conda_base` to set packages required by all
354
+ steps and use `@conda` to specify step-specific overrides.
285
355
 
286
356
 
287
357
  Parameters
288
358
  ----------
289
- cpu : int, default 1
290
- Number of CPUs required for this step.
291
- gpu : int, optional, default None
292
- Number of GPUs required for this step.
293
- disk : int, optional, default None
294
- Disk size (in MB) required for this step. Only applies on Kubernetes.
295
- memory : int, default 4096
296
- Memory size (in MB) required for this step.
297
- shared_memory : int, optional, default None
298
- The value for the size (in MiB) of the /dev/shm volume for this step.
299
- This parameter maps to the `--shm-size` option in Docker.
359
+ packages : Dict[str, str], default {}
360
+ Packages to use for this step. The key is the name of the package
361
+ and the value is the version to use.
362
+ libraries : Dict[str, str], default {}
363
+ Supported for backward compatibility. When used with packages, packages will take precedence.
364
+ python : str, optional, default None
365
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
366
+ that the version used will correspond to the version of the Python interpreter used to start the run.
367
+ disabled : bool, default False
368
+ If set to True, disables @conda.
300
369
  """
301
370
  ...
302
371
 
303
372
  @typing.overload
304
- 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]]]:
373
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
305
374
  """
306
- Enables checkpointing for a step.
307
-
308
- > Examples
309
-
310
- - Saving Checkpoints
311
-
312
- ```python
313
- @checkpoint
314
- @step
315
- def train(self):
316
- model = create_model(self.parameters, checkpoint_path = None)
317
- for i in range(self.epochs):
318
- # some training logic
319
- loss = model.train(self.dataset)
320
- if i % 10 == 0:
321
- model.save(
322
- current.checkpoint.directory,
323
- )
324
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
325
- # and returns a reference dictionary to the checkpoint saved in the datastore
326
- self.latest_checkpoint = current.checkpoint.save(
327
- name="epoch_checkpoint",
328
- metadata={
329
- "epoch": i,
330
- "loss": loss,
331
- }
332
- )
333
- ```
334
-
335
- - Using Loaded Checkpoints
375
+ Decorator prototype for all step decorators. This function gets specialized
376
+ and imported for all decorators types by _import_plugin_decorators().
377
+ """
378
+ ...
379
+
380
+ @typing.overload
381
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
382
+ ...
383
+
384
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
385
+ """
386
+ Decorator prototype for all step decorators. This function gets specialized
387
+ and imported for all decorators types by _import_plugin_decorators().
388
+ """
389
+ ...
390
+
391
+ @typing.overload
392
+ 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]]]:
393
+ """
394
+ Enables loading / saving of models within a step.
336
395
 
396
+ > Examples
397
+ - Saving Models
337
398
  ```python
338
- @retry(times=3)
339
- @checkpoint
399
+ @model
340
400
  @step
341
401
  def train(self):
342
- # Assume that the task has restarted and the previous attempt of the task
343
- # saved a checkpoint
344
- checkpoint_path = None
345
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
346
- print("Loaded checkpoint from the previous attempt")
347
- checkpoint_path = current.checkpoint.directory
402
+ # current.model.save returns a dictionary reference to the model saved
403
+ self.my_model = current.model.save(
404
+ path_to_my_model,
405
+ label="my_model",
406
+ metadata={
407
+ "epochs": 10,
408
+ "batch-size": 32,
409
+ "learning-rate": 0.001,
410
+ }
411
+ )
412
+ self.next(self.test)
348
413
 
349
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
350
- for i in range(self.epochs):
351
- ...
414
+ @model(load="my_model")
415
+ @step
416
+ def test(self):
417
+ # `current.model.loaded` returns a dictionary of the loaded models
418
+ # where the key is the name of the artifact and the value is the path to the model
419
+ print(os.listdir(current.model.loaded["my_model"]))
420
+ self.next(self.end)
421
+ ```
422
+
423
+ - Loading models
424
+ ```python
425
+ @step
426
+ def train(self):
427
+ # current.model.load returns the path to the model loaded
428
+ checkpoint_path = current.model.load(
429
+ self.checkpoint_key,
430
+ )
431
+ model_path = current.model.load(
432
+ self.model,
433
+ )
434
+ self.next(self.test)
352
435
  ```
353
436
 
354
437
 
355
438
  Parameters
356
439
  ----------
357
- load_policy : str, default: "fresh"
358
- The policy for loading the checkpoint. The following policies are supported:
359
- - "eager": Loads the the latest available checkpoint within the namespace.
360
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
361
- will be loaded at the start of the task.
362
- - "none": Do not load any checkpoint
363
- - "fresh": Loads the lastest checkpoint created within the running Task.
364
- This mode helps loading checkpoints across various retry attempts of the same task.
365
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
366
- created within the task will be loaded when the task is retries execution on failure.
440
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
441
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
442
+ 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`.
443
+ 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
444
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
445
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
367
446
 
368
447
  temp_dir_root : str, default: None
369
- The root directory under which `current.checkpoint.directory` will be created.
448
+ The root directory under which `current.model.loaded` will store loaded models
370
449
  """
371
450
  ...
372
451
 
373
452
  @typing.overload
374
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
453
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
375
454
  ...
376
455
 
377
456
  @typing.overload
378
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
457
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
379
458
  ...
380
459
 
381
- 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):
460
+ 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):
382
461
  """
383
- Enables checkpointing for a step.
462
+ Enables loading / saving of models within a step.
384
463
 
385
464
  > Examples
386
-
387
- - Saving Checkpoints
388
-
465
+ - Saving Models
389
466
  ```python
390
- @checkpoint
467
+ @model
391
468
  @step
392
469
  def train(self):
393
- model = create_model(self.parameters, checkpoint_path = None)
394
- for i in range(self.epochs):
395
- # some training logic
396
- loss = model.train(self.dataset)
397
- if i % 10 == 0:
398
- model.save(
399
- current.checkpoint.directory,
400
- )
401
- # saves the contents of the `current.checkpoint.directory` as a checkpoint
402
- # and returns a reference dictionary to the checkpoint saved in the datastore
403
- self.latest_checkpoint = current.checkpoint.save(
404
- name="epoch_checkpoint",
405
- metadata={
406
- "epoch": i,
407
- "loss": loss,
408
- }
409
- )
410
- ```
470
+ # current.model.save returns a dictionary reference to the model saved
471
+ self.my_model = current.model.save(
472
+ path_to_my_model,
473
+ label="my_model",
474
+ metadata={
475
+ "epochs": 10,
476
+ "batch-size": 32,
477
+ "learning-rate": 0.001,
478
+ }
479
+ )
480
+ self.next(self.test)
411
481
 
412
- - Using Loaded Checkpoints
482
+ @model(load="my_model")
483
+ @step
484
+ def test(self):
485
+ # `current.model.loaded` returns a dictionary of the loaded models
486
+ # where the key is the name of the artifact and the value is the path to the model
487
+ print(os.listdir(current.model.loaded["my_model"]))
488
+ self.next(self.end)
489
+ ```
413
490
 
491
+ - Loading models
414
492
  ```python
415
- @retry(times=3)
416
- @checkpoint
417
493
  @step
418
494
  def train(self):
419
- # Assume that the task has restarted and the previous attempt of the task
420
- # saved a checkpoint
421
- checkpoint_path = None
422
- if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
423
- print("Loaded checkpoint from the previous attempt")
424
- checkpoint_path = current.checkpoint.directory
425
-
426
- model = create_model(self.parameters, checkpoint_path = checkpoint_path)
427
- for i in range(self.epochs):
428
- ...
495
+ # current.model.load returns the path to the model loaded
496
+ checkpoint_path = current.model.load(
497
+ self.checkpoint_key,
498
+ )
499
+ model_path = current.model.load(
500
+ self.model,
501
+ )
502
+ self.next(self.test)
429
503
  ```
430
504
 
431
505
 
432
506
  Parameters
433
507
  ----------
434
- load_policy : str, default: "fresh"
435
- The policy for loading the checkpoint. The following policies are supported:
436
- - "eager": Loads the the latest available checkpoint within the namespace.
437
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
438
- will be loaded at the start of the task.
439
- - "none": Do not load any checkpoint
440
- - "fresh": Loads the lastest checkpoint created within the running Task.
441
- This mode helps loading checkpoints across various retry attempts of the same task.
442
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
443
- created within the task will be loaded when the task is retries execution on failure.
508
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
509
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
510
+ 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`.
511
+ 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
512
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
513
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
444
514
 
445
515
  temp_dir_root : str, default: None
446
- The root directory under which `current.checkpoint.directory` will be created.
516
+ The root directory under which `current.model.loaded` will store loaded models
447
517
  """
448
518
  ...
449
519
 
450
- 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]]]:
520
+ @typing.overload
521
+ 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]]]:
451
522
  """
452
- Specifies that this step should execute on DGX cloud.
523
+ Specifies that the step will success under all circumstances.
524
+
525
+ The decorator will create an optional artifact, specified by `var`, which
526
+ contains the exception raised. You can use it to detect the presence
527
+ of errors, indicating that all happy-path artifacts produced by the step
528
+ are missing.
453
529
 
454
530
 
455
531
  Parameters
456
532
  ----------
457
- gpu : int
458
- Number of GPUs to use.
459
- gpu_type : str
460
- Type of Nvidia GPU to use.
461
- queue_timeout : int
462
- Time to keep the job in NVCF's queue.
533
+ var : str, optional, default None
534
+ Name of the artifact in which to store the caught exception.
535
+ If not specified, the exception is not stored.
536
+ print_exception : bool, default True
537
+ Determines whether or not the exception is printed to
538
+ stdout when caught.
463
539
  """
464
540
  ...
465
541
 
466
542
  @typing.overload
467
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
468
- """
469
- Internal decorator to support Fast bakery
470
- """
543
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
471
544
  ...
472
545
 
473
546
  @typing.overload
474
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
547
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
475
548
  ...
476
549
 
477
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
550
+ 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):
478
551
  """
479
- Internal decorator to support Fast bakery
552
+ Specifies that the step will success under all circumstances.
553
+
554
+ The decorator will create an optional artifact, specified by `var`, which
555
+ contains the exception raised. You can use it to detect the presence
556
+ of errors, indicating that all happy-path artifacts produced by the step
557
+ are missing.
558
+
559
+
560
+ Parameters
561
+ ----------
562
+ var : str, optional, default None
563
+ Name of the artifact in which to store the caught exception.
564
+ If not specified, the exception is not stored.
565
+ print_exception : bool, default True
566
+ Determines whether or not the exception is printed to
567
+ stdout when caught.
480
568
  """
481
569
  ...
482
570
 
483
571
  @typing.overload
484
- 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]]]:
572
+ 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]]]:
485
573
  """
486
- Specifies the Conda environment for the step.
574
+ Specifies the PyPI packages for the step.
487
575
 
488
576
  Information in this decorator will augment any
489
- attributes set in the `@conda_base` flow-level decorator. Hence,
490
- you can use `@conda_base` to set packages required by all
491
- steps and use `@conda` to specify step-specific overrides.
577
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
578
+ you can use `@pypi_base` to set packages required by all
579
+ steps and use `@pypi` to specify step-specific overrides.
492
580
 
493
581
 
494
582
  Parameters
495
583
  ----------
496
- packages : Dict[str, str], default {}
584
+ packages : Dict[str, str], default: {}
497
585
  Packages to use for this step. The key is the name of the package
498
586
  and the value is the version to use.
499
- libraries : Dict[str, str], default {}
500
- Supported for backward compatibility. When used with packages, packages will take precedence.
501
- python : str, optional, default None
587
+ python : str, optional, default: None
502
588
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
503
589
  that the version used will correspond to the version of the Python interpreter used to start the run.
504
- disabled : bool, default False
505
- If set to True, disables @conda.
506
590
  """
507
591
  ...
508
592
 
509
593
  @typing.overload
510
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
594
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
511
595
  ...
512
596
 
513
597
  @typing.overload
514
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
598
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
515
599
  ...
516
600
 
517
- 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):
601
+ 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):
518
602
  """
519
- Specifies the Conda environment for the step.
603
+ Specifies the PyPI packages for the step.
520
604
 
521
605
  Information in this decorator will augment any
522
- attributes set in the `@conda_base` flow-level decorator. Hence,
523
- you can use `@conda_base` to set packages required by all
524
- steps and use `@conda` to specify step-specific overrides.
606
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
607
+ you can use `@pypi_base` to set packages required by all
608
+ steps and use `@pypi` to specify step-specific overrides.
525
609
 
526
610
 
527
611
  Parameters
528
612
  ----------
529
- packages : Dict[str, str], default {}
613
+ packages : Dict[str, str], default: {}
530
614
  Packages to use for this step. The key is the name of the package
531
615
  and the value is the version to use.
532
- libraries : Dict[str, str], default {}
533
- Supported for backward compatibility. When used with packages, packages will take precedence.
534
- python : str, optional, default None
616
+ python : str, optional, default: None
535
617
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
536
618
  that the version used will correspond to the version of the Python interpreter used to start the run.
537
- disabled : bool, default False
538
- If set to True, disables @conda.
539
- """
540
- ...
541
-
542
- @typing.overload
543
- 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]]]:
544
- """
545
- Specifies secrets to be retrieved and injected as environment variables prior to
546
- the execution of a step.
547
-
548
-
549
- Parameters
550
- ----------
551
- sources : List[Union[str, Dict[str, Any]]], default: []
552
- List of secret specs, defining how the secrets are to be retrieved
553
- role : str, optional, default: None
554
- Role to use for fetching secrets
555
- """
556
- ...
557
-
558
- @typing.overload
559
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
560
- ...
561
-
562
- @typing.overload
563
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
564
- ...
565
-
566
- 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):
567
- """
568
- Specifies secrets to be retrieved and injected as environment variables prior to
569
- the execution of a step.
570
-
571
-
572
- Parameters
573
- ----------
574
- sources : List[Union[str, Dict[str, Any]]], default: []
575
- List of secret specs, defining how the secrets are to be retrieved
576
- role : str, optional, default: None
577
- Role to use for fetching secrets
578
- """
579
- ...
580
-
581
- @typing.overload
582
- 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]]]:
583
- """
584
- Specifies environment variables to be set prior to the execution of a step.
585
-
586
-
587
- Parameters
588
- ----------
589
- vars : Dict[str, str], default {}
590
- Dictionary of environment variables to set.
591
- """
592
- ...
593
-
594
- @typing.overload
595
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
596
- ...
597
-
598
- @typing.overload
599
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
600
- ...
601
-
602
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
603
- """
604
- Specifies environment variables to be set prior to the execution of a step.
605
-
606
-
607
- Parameters
608
- ----------
609
- vars : Dict[str, str], default {}
610
- Dictionary of environment variables to set.
611
619
  """
612
620
  ...
613
621
 
@@ -701,7 +709,7 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
701
709
  ...
702
710
 
703
711
  @typing.overload
704
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
712
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
705
713
  """
706
714
  Decorator prototype for all step decorators. This function gets specialized
707
715
  and imported for all decorators types by _import_plugin_decorators().
@@ -709,554 +717,698 @@ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.
709
717
  ...
710
718
 
711
719
  @typing.overload
712
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
720
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
713
721
  ...
714
722
 
715
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
723
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
716
724
  """
717
725
  Decorator prototype for all step decorators. This function gets specialized
718
726
  and imported for all decorators types by _import_plugin_decorators().
719
727
  """
720
728
  ...
721
729
 
722
- 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]]]:
730
+ 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]]]:
723
731
  """
724
- Decorator that helps cache, version and store models/datasets from huggingface hub.
732
+ Specifies that this step should execute on DGX cloud.
725
733
 
726
- > Examples
727
734
 
728
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
729
- ```python
730
- @huggingface_hub
731
- @step
732
- def pull_model_from_huggingface(self):
733
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
734
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
735
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
736
- # value of the function is a reference to the model in the backend storage.
737
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
735
+ Parameters
736
+ ----------
737
+ gpu : int
738
+ Number of GPUs to use.
739
+ gpu_type : str
740
+ Type of Nvidia GPU to use.
741
+ queue_timeout : int
742
+ Time to keep the job in NVCF's queue.
743
+ """
744
+ ...
745
+
746
+ @typing.overload
747
+ 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]]]:
748
+ """
749
+ Specifies environment variables to be set prior to the execution of a step.
738
750
 
739
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
740
- self.llama_model = current.huggingface_hub.snapshot_download(
741
- repo_id=self.model_id,
742
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
743
- )
744
- self.next(self.train)
745
- ```
746
751
 
747
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
748
- ```python
749
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
750
- @step
751
- def pull_model_from_huggingface(self):
752
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
753
- ```
752
+ Parameters
753
+ ----------
754
+ vars : Dict[str, str], default {}
755
+ Dictionary of environment variables to set.
756
+ """
757
+ ...
758
+
759
+ @typing.overload
760
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
761
+ ...
762
+
763
+ @typing.overload
764
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
765
+ ...
766
+
767
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
768
+ """
769
+ Specifies environment variables to be set prior to the execution of a step.
770
+
771
+
772
+ Parameters
773
+ ----------
774
+ vars : Dict[str, str], default {}
775
+ Dictionary of environment variables to set.
776
+ """
777
+ ...
778
+
779
+ @typing.overload
780
+ 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]]]:
781
+ """
782
+ Enables checkpointing for a step.
783
+
784
+ > Examples
785
+
786
+ - Saving Checkpoints
754
787
 
755
788
  ```python
756
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
757
- @step
758
- def finetune_model(self):
759
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
760
- # path_to_model will be /my-directory
789
+ @checkpoint
790
+ @step
791
+ def train(self):
792
+ model = create_model(self.parameters, checkpoint_path = None)
793
+ for i in range(self.epochs):
794
+ # some training logic
795
+ loss = model.train(self.dataset)
796
+ if i % 10 == 0:
797
+ model.save(
798
+ current.checkpoint.directory,
799
+ )
800
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
801
+ # and returns a reference dictionary to the checkpoint saved in the datastore
802
+ self.latest_checkpoint = current.checkpoint.save(
803
+ name="epoch_checkpoint",
804
+ metadata={
805
+ "epoch": i,
806
+ "loss": loss,
807
+ }
808
+ )
761
809
  ```
762
810
 
811
+ - Using Loaded Checkpoints
812
+
763
813
  ```python
764
- # Takes all the arguments passed to `snapshot_download`
765
- # except for `local_dir`
766
- @huggingface_hub(load=[
767
- {
768
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
769
- },
770
- {
771
- "repo_id": "myorg/mistral-lora",
772
- "repo_type": "model",
773
- },
774
- ])
775
- @step
776
- def finetune_model(self):
777
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
778
- # path_to_model will be /my-directory
814
+ @retry(times=3)
815
+ @checkpoint
816
+ @step
817
+ def train(self):
818
+ # Assume that the task has restarted and the previous attempt of the task
819
+ # saved a checkpoint
820
+ checkpoint_path = None
821
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
822
+ print("Loaded checkpoint from the previous attempt")
823
+ checkpoint_path = current.checkpoint.directory
824
+
825
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
826
+ for i in range(self.epochs):
827
+ ...
779
828
  ```
780
829
 
781
830
 
782
831
  Parameters
783
832
  ----------
784
- temp_dir_root : str, optional
785
- The root directory that will hold the temporary directory where objects will be downloaded.
786
-
787
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
788
- The list of repos (models/datasets) to load.
789
-
790
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
791
-
792
- - If repo (model/dataset) is not found in the datastore:
793
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
794
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
795
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
833
+ load_policy : str, default: "fresh"
834
+ The policy for loading the checkpoint. The following policies are supported:
835
+ - "eager": Loads the the latest available checkpoint within the namespace.
836
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
837
+ will be loaded at the start of the task.
838
+ - "none": Do not load any checkpoint
839
+ - "fresh": Loads the lastest checkpoint created within the running Task.
840
+ This mode helps loading checkpoints across various retry attempts of the same task.
841
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
842
+ created within the task will be loaded when the task is retries execution on failure.
796
843
 
797
- - If repo is found in the datastore:
798
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
844
+ temp_dir_root : str, default: None
845
+ The root directory under which `current.checkpoint.directory` will be created.
799
846
  """
800
847
  ...
801
848
 
802
849
  @typing.overload
803
- 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]]]:
850
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
851
+ ...
852
+
853
+ @typing.overload
854
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
855
+ ...
856
+
857
+ 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):
804
858
  """
805
- Enables loading / saving of models within a step.
859
+ Enables checkpointing for a step.
806
860
 
807
861
  > Examples
808
- - Saving Models
862
+
863
+ - Saving Checkpoints
864
+
809
865
  ```python
810
- @model
866
+ @checkpoint
811
867
  @step
812
868
  def train(self):
813
- # current.model.save returns a dictionary reference to the model saved
814
- self.my_model = current.model.save(
815
- path_to_my_model,
816
- label="my_model",
817
- metadata={
818
- "epochs": 10,
819
- "batch-size": 32,
820
- "learning-rate": 0.001,
821
- }
822
- )
823
- self.next(self.test)
824
-
825
- @model(load="my_model")
826
- @step
827
- def test(self):
828
- # `current.model.loaded` returns a dictionary of the loaded models
829
- # where the key is the name of the artifact and the value is the path to the model
830
- print(os.listdir(current.model.loaded["my_model"]))
831
- self.next(self.end)
869
+ model = create_model(self.parameters, checkpoint_path = None)
870
+ for i in range(self.epochs):
871
+ # some training logic
872
+ loss = model.train(self.dataset)
873
+ if i % 10 == 0:
874
+ model.save(
875
+ current.checkpoint.directory,
876
+ )
877
+ # saves the contents of the `current.checkpoint.directory` as a checkpoint
878
+ # and returns a reference dictionary to the checkpoint saved in the datastore
879
+ self.latest_checkpoint = current.checkpoint.save(
880
+ name="epoch_checkpoint",
881
+ metadata={
882
+ "epoch": i,
883
+ "loss": loss,
884
+ }
885
+ )
832
886
  ```
833
887
 
834
- - Loading models
888
+ - Using Loaded Checkpoints
889
+
835
890
  ```python
891
+ @retry(times=3)
892
+ @checkpoint
836
893
  @step
837
894
  def train(self):
838
- # current.model.load returns the path to the model loaded
839
- checkpoint_path = current.model.load(
840
- self.checkpoint_key,
841
- )
842
- model_path = current.model.load(
843
- self.model,
844
- )
845
- self.next(self.test)
895
+ # Assume that the task has restarted and the previous attempt of the task
896
+ # saved a checkpoint
897
+ checkpoint_path = None
898
+ if current.checkpoint.is_loaded: # Check if a checkpoint is loaded
899
+ print("Loaded checkpoint from the previous attempt")
900
+ checkpoint_path = current.checkpoint.directory
901
+
902
+ model = create_model(self.parameters, checkpoint_path = checkpoint_path)
903
+ for i in range(self.epochs):
904
+ ...
846
905
  ```
847
906
 
848
907
 
849
908
  Parameters
850
909
  ----------
851
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
852
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
853
- 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`.
854
- 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
855
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
856
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
910
+ load_policy : str, default: "fresh"
911
+ The policy for loading the checkpoint. The following policies are supported:
912
+ - "eager": Loads the the latest available checkpoint within the namespace.
913
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
914
+ will be loaded at the start of the task.
915
+ - "none": Do not load any checkpoint
916
+ - "fresh": Loads the lastest checkpoint created within the running Task.
917
+ This mode helps loading checkpoints across various retry attempts of the same task.
918
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
919
+ created within the task will be loaded when the task is retries execution on failure.
857
920
 
858
921
  temp_dir_root : str, default: None
859
- The root directory under which `current.model.loaded` will store loaded models
922
+ The root directory under which `current.checkpoint.directory` will be created.
860
923
  """
861
924
  ...
862
925
 
863
926
  @typing.overload
864
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
927
+ 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]]]:
928
+ """
929
+ Specifies secrets to be retrieved and injected as environment variables prior to
930
+ the execution of a step.
931
+
932
+
933
+ Parameters
934
+ ----------
935
+ sources : List[Union[str, Dict[str, Any]]], default: []
936
+ List of secret specs, defining how the secrets are to be retrieved
937
+ role : str, optional, default: None
938
+ Role to use for fetching secrets
939
+ """
865
940
  ...
866
941
 
867
942
  @typing.overload
868
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
943
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
869
944
  ...
870
945
 
871
- 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):
946
+ @typing.overload
947
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
948
+ ...
949
+
950
+ 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):
872
951
  """
873
- Enables loading / saving of models within a step.
952
+ Specifies secrets to be retrieved and injected as environment variables prior to
953
+ the execution of a step.
954
+
955
+
956
+ Parameters
957
+ ----------
958
+ sources : List[Union[str, Dict[str, Any]]], default: []
959
+ List of secret specs, defining how the secrets are to be retrieved
960
+ role : str, optional, default: None
961
+ Role to use for fetching secrets
962
+ """
963
+ ...
964
+
965
+ 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]]]:
966
+ """
967
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
874
968
 
875
969
  > Examples
876
- - Saving Models
970
+
971
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
877
972
  ```python
878
- @model
879
- @step
880
- def train(self):
881
- # current.model.save returns a dictionary reference to the model saved
882
- self.my_model = current.model.save(
883
- path_to_my_model,
884
- label="my_model",
885
- metadata={
886
- "epochs": 10,
887
- "batch-size": 32,
888
- "learning-rate": 0.001,
889
- }
890
- )
891
- self.next(self.test)
973
+ @huggingface_hub
974
+ @step
975
+ def pull_model_from_huggingface(self):
976
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
977
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
978
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
979
+ # value of the function is a reference to the model in the backend storage.
980
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
892
981
 
893
- @model(load="my_model")
894
- @step
895
- def test(self):
896
- # `current.model.loaded` returns a dictionary of the loaded models
897
- # where the key is the name of the artifact and the value is the path to the model
898
- print(os.listdir(current.model.loaded["my_model"]))
899
- self.next(self.end)
982
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
983
+ self.llama_model = current.huggingface_hub.snapshot_download(
984
+ repo_id=self.model_id,
985
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
986
+ )
987
+ self.next(self.train)
988
+ ```
989
+
990
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
991
+ ```python
992
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
993
+ @step
994
+ def pull_model_from_huggingface(self):
995
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
900
996
  ```
901
997
 
902
- - Loading models
903
998
  ```python
904
- @step
905
- def train(self):
906
- # current.model.load returns the path to the model loaded
907
- checkpoint_path = current.model.load(
908
- self.checkpoint_key,
909
- )
910
- model_path = current.model.load(
911
- self.model,
912
- )
913
- self.next(self.test)
999
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1000
+ @step
1001
+ def finetune_model(self):
1002
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1003
+ # path_to_model will be /my-directory
1004
+ ```
1005
+
1006
+ ```python
1007
+ # Takes all the arguments passed to `snapshot_download`
1008
+ # except for `local_dir`
1009
+ @huggingface_hub(load=[
1010
+ {
1011
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1012
+ },
1013
+ {
1014
+ "repo_id": "myorg/mistral-lora",
1015
+ "repo_type": "model",
1016
+ },
1017
+ ])
1018
+ @step
1019
+ def finetune_model(self):
1020
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1021
+ # path_to_model will be /my-directory
914
1022
  ```
915
1023
 
916
1024
 
917
1025
  Parameters
918
1026
  ----------
919
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
920
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
921
- 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`.
922
- 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
923
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
924
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
1027
+ temp_dir_root : str, optional
1028
+ The root directory that will hold the temporary directory where objects will be downloaded.
925
1029
 
926
- temp_dir_root : str, default: None
927
- The root directory under which `current.model.loaded` will store loaded models
1030
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1031
+ The list of repos (models/datasets) to load.
1032
+
1033
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1034
+
1035
+ - If repo (model/dataset) is not found in the datastore:
1036
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1037
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1038
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1039
+
1040
+ - If repo is found in the datastore:
1041
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
928
1042
  """
929
1043
  ...
930
1044
 
931
- @typing.overload
932
- 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]]]:
1045
+ 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]]]:
933
1046
  """
934
- Specifies that the step will success under all circumstances.
935
-
936
- The decorator will create an optional artifact, specified by `var`, which
937
- contains the exception raised. You can use it to detect the presence
938
- of errors, indicating that all happy-path artifacts produced by the step
939
- are missing.
1047
+ Specifies that this step should execute on DGX cloud.
940
1048
 
941
1049
 
942
1050
  Parameters
943
1051
  ----------
944
- var : str, optional, default None
945
- Name of the artifact in which to store the caught exception.
946
- If not specified, the exception is not stored.
947
- print_exception : bool, default True
948
- Determines whether or not the exception is printed to
949
- stdout when caught.
1052
+ gpu : int
1053
+ Number of GPUs to use.
1054
+ gpu_type : str
1055
+ Type of Nvidia GPU to use.
950
1056
  """
951
1057
  ...
952
1058
 
953
1059
  @typing.overload
954
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
955
- ...
956
-
957
- @typing.overload
958
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
959
- ...
960
-
961
- 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):
1060
+ def timeout(*, seconds: int = 0, minutes: int = 0, hours: int = 0) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
962
1061
  """
963
- Specifies that the step will success under all circumstances.
1062
+ Specifies a timeout for your step.
964
1063
 
965
- The decorator will create an optional artifact, specified by `var`, which
966
- contains the exception raised. You can use it to detect the presence
967
- of errors, indicating that all happy-path artifacts produced by the step
968
- are missing.
1064
+ This decorator is useful if this step may hang indefinitely.
969
1065
 
1066
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1067
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1068
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
970
1069
 
971
- Parameters
972
- ----------
973
- var : str, optional, default None
974
- Name of the artifact in which to store the caught exception.
975
- If not specified, the exception is not stored.
976
- print_exception : bool, default True
977
- Determines whether or not the exception is printed to
978
- stdout when caught.
979
- """
980
- ...
981
-
982
- 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]]]:
983
- """
984
- Specifies that this step should execute on DGX cloud.
1070
+ Note that all the values specified in parameters are added together so if you specify
1071
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
985
1072
 
986
1073
 
987
1074
  Parameters
988
1075
  ----------
989
- gpu : int
990
- Number of GPUs to use.
991
- gpu_type : str
992
- Type of Nvidia GPU to use.
1076
+ seconds : int, default 0
1077
+ Number of seconds to wait prior to timing out.
1078
+ minutes : int, default 0
1079
+ Number of minutes to wait prior to timing out.
1080
+ hours : int, default 0
1081
+ Number of hours to wait prior to timing out.
993
1082
  """
994
1083
  ...
995
1084
 
996
- 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]]]:
1085
+ @typing.overload
1086
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1087
+ ...
1088
+
1089
+ @typing.overload
1090
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1091
+ ...
1092
+
1093
+ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, seconds: int = 0, minutes: int = 0, hours: int = 0):
997
1094
  """
998
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
1095
+ Specifies a timeout for your step.
999
1096
 
1000
- User code call
1001
- --------------
1002
- @ollama(
1003
- models=[...],
1004
- ...
1005
- )
1097
+ This decorator is useful if this step may hang indefinitely.
1006
1098
 
1007
- Valid backend options
1008
- ---------------------
1009
- - 'local': Run as a separate process on the local task machine.
1010
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
1011
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
1099
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1100
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1101
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1012
1102
 
1013
- Valid model options
1014
- -------------------
1015
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
1103
+ Note that all the values specified in parameters are added together so if you specify
1104
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1016
1105
 
1017
1106
 
1018
1107
  Parameters
1019
1108
  ----------
1020
- models: list[str]
1021
- List of Ollama containers running models in sidecars.
1022
- backend: str
1023
- Determines where and how to run the Ollama process.
1024
- force_pull: bool
1025
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
1026
- cache_update_policy: str
1027
- Cache update policy: "auto", "force", or "never".
1028
- force_cache_update: bool
1029
- Simple override for "force" cache update policy.
1030
- debug: bool
1031
- Whether to turn on verbose debugging logs.
1032
- circuit_breaker_config: dict
1033
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
1034
- timeout_config: dict
1035
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
1109
+ seconds : int, default 0
1110
+ Number of seconds to wait prior to timing out.
1111
+ minutes : int, default 0
1112
+ Number of minutes to wait prior to timing out.
1113
+ hours : int, default 0
1114
+ Number of hours to wait prior to timing out.
1036
1115
  """
1037
1116
  ...
1038
1117
 
1039
1118
  @typing.overload
1040
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1119
+ 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]]]:
1041
1120
  """
1042
- Specifies the number of times the task corresponding
1043
- to a step needs to be retried.
1044
-
1045
- This decorator is useful for handling transient errors, such as networking issues.
1046
- If your task contains operations that can't be retried safely, e.g. database updates,
1047
- it is advisable to annotate it with `@retry(times=0)`.
1121
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1048
1122
 
1049
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1050
- decorator will execute a no-op task after all retries have been exhausted,
1051
- ensuring that the flow execution can continue.
1123
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1052
1124
 
1053
1125
 
1054
1126
  Parameters
1055
1127
  ----------
1056
- times : int, default 3
1057
- Number of times to retry this task.
1058
- minutes_between_retries : int, default 2
1059
- Number of minutes between retries.
1128
+ type : str, default 'default'
1129
+ Card type.
1130
+ id : str, optional, default None
1131
+ If multiple cards are present, use this id to identify this card.
1132
+ options : Dict[str, Any], default {}
1133
+ Options passed to the card. The contents depend on the card type.
1134
+ timeout : int, default 45
1135
+ Interrupt reporting if it takes more than this many seconds.
1060
1136
  """
1061
1137
  ...
1062
1138
 
1063
1139
  @typing.overload
1064
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1140
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1065
1141
  ...
1066
1142
 
1067
1143
  @typing.overload
1068
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1144
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1069
1145
  ...
1070
1146
 
1071
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
1147
+ 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):
1072
1148
  """
1073
- Specifies the number of times the task corresponding
1074
- to a step needs to be retried.
1075
-
1076
- This decorator is useful for handling transient errors, such as networking issues.
1077
- If your task contains operations that can't be retried safely, e.g. database updates,
1078
- it is advisable to annotate it with `@retry(times=0)`.
1149
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1079
1150
 
1080
- This can be used in conjunction with the `@catch` decorator. The `@catch`
1081
- decorator will execute a no-op task after all retries have been exhausted,
1082
- ensuring that the flow execution can continue.
1151
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1083
1152
 
1084
1153
 
1085
1154
  Parameters
1086
1155
  ----------
1087
- times : int, default 3
1088
- Number of times to retry this task.
1089
- minutes_between_retries : int, default 2
1090
- Number of minutes between retries.
1156
+ type : str, default 'default'
1157
+ Card type.
1158
+ id : str, optional, default None
1159
+ If multiple cards are present, use this id to identify this card.
1160
+ options : Dict[str, Any], default {}
1161
+ Options passed to the card. The contents depend on the card type.
1162
+ timeout : int, default 45
1163
+ Interrupt reporting if it takes more than this many seconds.
1091
1164
  """
1092
1165
  ...
1093
1166
 
1094
1167
  @typing.overload
1095
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1168
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1096
1169
  """
1097
- Decorator prototype for all step decorators. This function gets specialized
1098
- and imported for all decorators types by _import_plugin_decorators().
1170
+ Internal decorator to support Fast bakery
1099
1171
  """
1100
1172
  ...
1101
1173
 
1102
1174
  @typing.overload
1103
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1175
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1104
1176
  ...
1105
1177
 
1106
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1178
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
1107
1179
  """
1108
- Decorator prototype for all step decorators. This function gets specialized
1109
- and imported for all decorators types by _import_plugin_decorators().
1180
+ Internal decorator to support Fast bakery
1110
1181
  """
1111
1182
  ...
1112
1183
 
1113
1184
  @typing.overload
1114
- 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]]]:
1185
+ 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]]]:
1115
1186
  """
1116
- Specifies the PyPI packages for the step.
1187
+ Specifies the resources needed when executing this step.
1117
1188
 
1118
- Information in this decorator will augment any
1119
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1120
- you can use `@pypi_base` to set packages required by all
1121
- steps and use `@pypi` to specify step-specific overrides.
1189
+ Use `@resources` to specify the resource requirements
1190
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1191
+
1192
+ You can choose the compute layer on the command line by executing e.g.
1193
+ ```
1194
+ python myflow.py run --with batch
1195
+ ```
1196
+ or
1197
+ ```
1198
+ python myflow.py run --with kubernetes
1199
+ ```
1200
+ which executes the flow on the desired system using the
1201
+ requirements specified in `@resources`.
1122
1202
 
1123
1203
 
1124
1204
  Parameters
1125
1205
  ----------
1126
- packages : Dict[str, str], default: {}
1127
- Packages to use for this step. The key is the name of the package
1128
- and the value is the version to use.
1129
- python : str, optional, default: None
1130
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1131
- that the version used will correspond to the version of the Python interpreter used to start the run.
1206
+ cpu : int, default 1
1207
+ Number of CPUs required for this step.
1208
+ gpu : int, optional, default None
1209
+ Number of GPUs required for this step.
1210
+ disk : int, optional, default None
1211
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1212
+ memory : int, default 4096
1213
+ Memory size (in MB) required for this step.
1214
+ shared_memory : int, optional, default None
1215
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1216
+ This parameter maps to the `--shm-size` option in Docker.
1132
1217
  """
1133
1218
  ...
1134
1219
 
1135
1220
  @typing.overload
1136
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1221
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1137
1222
  ...
1138
1223
 
1139
1224
  @typing.overload
1140
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1225
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1141
1226
  ...
1142
1227
 
1143
- 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):
1228
+ 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):
1144
1229
  """
1145
- Specifies the PyPI packages for the step.
1230
+ Specifies the resources needed when executing this step.
1146
1231
 
1147
- Information in this decorator will augment any
1148
- attributes set in the `@pyi_base` flow-level decorator. Hence,
1149
- you can use `@pypi_base` to set packages required by all
1150
- steps and use `@pypi` to specify step-specific overrides.
1232
+ Use `@resources` to specify the resource requirements
1233
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
1234
+
1235
+ You can choose the compute layer on the command line by executing e.g.
1236
+ ```
1237
+ python myflow.py run --with batch
1238
+ ```
1239
+ or
1240
+ ```
1241
+ python myflow.py run --with kubernetes
1242
+ ```
1243
+ which executes the flow on the desired system using the
1244
+ requirements specified in `@resources`.
1151
1245
 
1152
1246
 
1153
1247
  Parameters
1154
1248
  ----------
1155
- packages : Dict[str, str], default: {}
1156
- Packages to use for this step. The key is the name of the package
1157
- and the value is the version to use.
1158
- python : str, optional, default: None
1159
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1160
- that the version used will correspond to the version of the Python interpreter used to start the run.
1249
+ cpu : int, default 1
1250
+ Number of CPUs required for this step.
1251
+ gpu : int, optional, default None
1252
+ Number of GPUs required for this step.
1253
+ disk : int, optional, default None
1254
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
1255
+ memory : int, default 4096
1256
+ Memory size (in MB) required for this step.
1257
+ shared_memory : int, optional, default None
1258
+ The value for the size (in MiB) of the /dev/shm volume for this step.
1259
+ This parameter maps to the `--shm-size` option in Docker.
1161
1260
  """
1162
1261
  ...
1163
1262
 
1164
- 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]]]:
1263
+ @typing.overload
1264
+ 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]]:
1165
1265
  """
1166
- This decorator is used to run vllm APIs as Metaflow task sidecars.
1266
+ Specifies the flow(s) that this flow depends on.
1167
1267
 
1168
- User code call
1169
- --------------
1170
- @vllm(
1171
- model="...",
1172
- ...
1173
- )
1268
+ ```
1269
+ @trigger_on_finish(flow='FooFlow')
1270
+ ```
1271
+ or
1272
+ ```
1273
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1274
+ ```
1275
+ This decorator respects the @project decorator and triggers the flow
1276
+ when upstream runs within the same namespace complete successfully
1174
1277
 
1175
- Valid backend options
1176
- ---------------------
1177
- - 'local': Run as a separate process on the local task machine.
1278
+ Additionally, you can specify project aware upstream flow dependencies
1279
+ by specifying the fully qualified project_flow_name.
1280
+ ```
1281
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1282
+ ```
1283
+ or
1284
+ ```
1285
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1286
+ ```
1178
1287
 
1179
- Valid model options
1180
- -------------------
1181
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
1288
+ You can also specify just the project or project branch (other values will be
1289
+ inferred from the current project or project branch):
1290
+ ```
1291
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1292
+ ```
1182
1293
 
1183
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
1184
- If you need multiple models, you must create multiple @vllm decorators.
1294
+ Note that `branch` is typically one of:
1295
+ - `prod`
1296
+ - `user.bob`
1297
+ - `test.my_experiment`
1298
+ - `prod.staging`
1185
1299
 
1186
1300
 
1187
1301
  Parameters
1188
1302
  ----------
1189
- model: str
1190
- HuggingFace model identifier to be served by vLLM.
1191
- backend: str
1192
- Determines where and how to run the vLLM process.
1193
- openai_api_server: bool
1194
- Whether to use OpenAI-compatible API server mode (subprocess) instead of native engine.
1195
- Default is False (uses native engine).
1196
- Set to True for backward compatibility with existing code.
1197
- debug: bool
1198
- Whether to turn on verbose debugging logs.
1199
- card_refresh_interval: int
1200
- Interval in seconds for refreshing the vLLM status card.
1201
- Only used when openai_api_server=True.
1202
- max_retries: int
1203
- Maximum number of retries checking for vLLM server startup.
1204
- Only used when openai_api_server=True.
1205
- retry_alert_frequency: int
1206
- Frequency of alert logs for vLLM server startup retries.
1207
- Only used when openai_api_server=True.
1208
- engine_args : dict
1209
- Additional keyword arguments to pass to the vLLM engine.
1210
- For example, `tensor_parallel_size=2`.
1303
+ flow : Union[str, Dict[str, str]], optional, default None
1304
+ Upstream flow dependency for this flow.
1305
+ flows : List[Union[str, Dict[str, str]]], default []
1306
+ Upstream flow dependencies for this flow.
1307
+ options : Dict[str, Any], default {}
1308
+ Backend-specific configuration for tuning eventing behavior.
1211
1309
  """
1212
1310
  ...
1213
1311
 
1214
1312
  @typing.overload
1215
- 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]]]:
1313
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1314
+ ...
1315
+
1316
+ 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] = {}):
1216
1317
  """
1217
- Creates a human-readable report, a Metaflow Card, after this step completes.
1318
+ Specifies the flow(s) that this flow depends on.
1218
1319
 
1219
- Note that you may add multiple `@card` decorators in a step with different parameters.
1320
+ ```
1321
+ @trigger_on_finish(flow='FooFlow')
1322
+ ```
1323
+ or
1324
+ ```
1325
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1326
+ ```
1327
+ This decorator respects the @project decorator and triggers the flow
1328
+ when upstream runs within the same namespace complete successfully
1329
+
1330
+ Additionally, you can specify project aware upstream flow dependencies
1331
+ by specifying the fully qualified project_flow_name.
1332
+ ```
1333
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1334
+ ```
1335
+ or
1336
+ ```
1337
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1338
+ ```
1339
+
1340
+ You can also specify just the project or project branch (other values will be
1341
+ inferred from the current project or project branch):
1342
+ ```
1343
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1344
+ ```
1345
+
1346
+ Note that `branch` is typically one of:
1347
+ - `prod`
1348
+ - `user.bob`
1349
+ - `test.my_experiment`
1350
+ - `prod.staging`
1220
1351
 
1221
1352
 
1222
1353
  Parameters
1223
1354
  ----------
1224
- type : str, default 'default'
1225
- Card type.
1226
- id : str, optional, default None
1227
- If multiple cards are present, use this id to identify this card.
1355
+ flow : Union[str, Dict[str, str]], optional, default None
1356
+ Upstream flow dependency for this flow.
1357
+ flows : List[Union[str, Dict[str, str]]], default []
1358
+ Upstream flow dependencies for this flow.
1228
1359
  options : Dict[str, Any], default {}
1229
- Options passed to the card. The contents depend on the card type.
1230
- timeout : int, default 45
1231
- Interrupt reporting if it takes more than this many seconds.
1360
+ Backend-specific configuration for tuning eventing behavior.
1232
1361
  """
1233
1362
  ...
1234
1363
 
1235
1364
  @typing.overload
1236
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1365
+ 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]]:
1366
+ """
1367
+ Specifies the times when the flow should be run when running on a
1368
+ production scheduler.
1369
+
1370
+
1371
+ Parameters
1372
+ ----------
1373
+ hourly : bool, default False
1374
+ Run the workflow hourly.
1375
+ daily : bool, default True
1376
+ Run the workflow daily.
1377
+ weekly : bool, default False
1378
+ Run the workflow weekly.
1379
+ cron : str, optional, default None
1380
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1381
+ specified by this expression.
1382
+ timezone : str, optional, default None
1383
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1384
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1385
+ """
1237
1386
  ...
1238
1387
 
1239
1388
  @typing.overload
1240
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1389
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1241
1390
  ...
1242
1391
 
1243
- 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):
1392
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1244
1393
  """
1245
- Creates a human-readable report, a Metaflow Card, after this step completes.
1246
-
1247
- Note that you may add multiple `@card` decorators in a step with different parameters.
1394
+ Specifies the times when the flow should be run when running on a
1395
+ production scheduler.
1248
1396
 
1249
1397
 
1250
1398
  Parameters
1251
1399
  ----------
1252
- type : str, default 'default'
1253
- Card type.
1254
- id : str, optional, default None
1255
- If multiple cards are present, use this id to identify this card.
1256
- options : Dict[str, Any], default {}
1257
- Options passed to the card. The contents depend on the card type.
1258
- timeout : int, default 45
1259
- Interrupt reporting if it takes more than this many seconds.
1400
+ hourly : bool, default False
1401
+ Run the workflow hourly.
1402
+ daily : bool, default True
1403
+ Run the workflow daily.
1404
+ weekly : bool, default False
1405
+ Run the workflow weekly.
1406
+ cron : str, optional, default None
1407
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1408
+ specified by this expression.
1409
+ timezone : str, optional, default None
1410
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1411
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1260
1412
  """
1261
1413
  ...
1262
1414
 
@@ -1418,53 +1570,78 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1418
1570
  ...
1419
1571
 
1420
1572
  @typing.overload
1421
- 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]]:
1573
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1574
+ """
1575
+ Specifies the PyPI packages for all steps of the flow.
1576
+
1577
+ Use `@pypi_base` to set common packages required by all
1578
+ steps and use `@pypi` to specify step-specific overrides.
1579
+
1580
+ Parameters
1581
+ ----------
1582
+ packages : Dict[str, str], default: {}
1583
+ Packages to use for this flow. The key is the name of the package
1584
+ and the value is the version to use.
1585
+ python : str, optional, default: None
1586
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1587
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1588
+ """
1589
+ ...
1590
+
1591
+ @typing.overload
1592
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1593
+ ...
1594
+
1595
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1422
1596
  """
1423
- Specifies the times when the flow should be run when running on a
1424
- production scheduler.
1597
+ Specifies the PyPI packages for all steps of the flow.
1425
1598
 
1599
+ Use `@pypi_base` to set common packages required by all
1600
+ steps and use `@pypi` to specify step-specific overrides.
1426
1601
 
1427
1602
  Parameters
1428
1603
  ----------
1429
- hourly : bool, default False
1430
- Run the workflow hourly.
1431
- daily : bool, default True
1432
- Run the workflow daily.
1433
- weekly : bool, default False
1434
- Run the workflow weekly.
1435
- cron : str, optional, default None
1436
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1437
- specified by this expression.
1438
- timezone : str, optional, default None
1439
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1440
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1604
+ packages : Dict[str, str], default: {}
1605
+ Packages to use for this flow. The key is the name of the package
1606
+ and the value is the version to use.
1607
+ python : str, optional, default: None
1608
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1609
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1441
1610
  """
1442
1611
  ...
1443
1612
 
1444
- @typing.overload
1445
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1446
- ...
1447
-
1448
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1613
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1449
1614
  """
1450
- Specifies the times when the flow should be run when running on a
1451
- production scheduler.
1615
+ Specifies what flows belong to the same project.
1616
+
1617
+ A project-specific namespace is created for all flows that
1618
+ use the same `@project(name)`.
1452
1619
 
1453
1620
 
1454
1621
  Parameters
1455
1622
  ----------
1456
- hourly : bool, default False
1457
- Run the workflow hourly.
1458
- daily : bool, default True
1459
- Run the workflow daily.
1460
- weekly : bool, default False
1461
- Run the workflow weekly.
1462
- cron : str, optional, default None
1463
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1464
- specified by this expression.
1465
- timezone : str, optional, default None
1466
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1467
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1623
+ name : str
1624
+ Project name. Make sure that the name is unique amongst all
1625
+ projects that use the same production scheduler. The name may
1626
+ contain only lowercase alphanumeric characters and underscores.
1627
+
1628
+ branch : Optional[str], default None
1629
+ The branch to use. If not specified, the branch is set to
1630
+ `user.<username>` unless `production` is set to `True`. This can
1631
+ also be set on the command line using `--branch` as a top-level option.
1632
+ It is an error to specify `branch` in the decorator and on the command line.
1633
+
1634
+ production : bool, default False
1635
+ Whether or not the branch is the production branch. This can also be set on the
1636
+ command line using `--production` as a top-level option. It is an error to specify
1637
+ `production` in the decorator and on the command line.
1638
+ The project branch name will be:
1639
+ - if `branch` is specified:
1640
+ - if `production` is True: `prod.<branch>`
1641
+ - if `production` is False: `test.<branch>`
1642
+ - if `branch` is not specified:
1643
+ - if `production` is True: `prod`
1644
+ - if `production` is False: `user.<username>`
1468
1645
  """
1469
1646
  ...
1470
1647
 
@@ -1511,142 +1688,6 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1511
1688
  """
1512
1689
  ...
1513
1690
 
1514
- @typing.overload
1515
- 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]]:
1516
- """
1517
- Specifies the flow(s) that this flow depends on.
1518
-
1519
- ```
1520
- @trigger_on_finish(flow='FooFlow')
1521
- ```
1522
- or
1523
- ```
1524
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1525
- ```
1526
- This decorator respects the @project decorator and triggers the flow
1527
- when upstream runs within the same namespace complete successfully
1528
-
1529
- Additionally, you can specify project aware upstream flow dependencies
1530
- by specifying the fully qualified project_flow_name.
1531
- ```
1532
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1533
- ```
1534
- or
1535
- ```
1536
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1537
- ```
1538
-
1539
- You can also specify just the project or project branch (other values will be
1540
- inferred from the current project or project branch):
1541
- ```
1542
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1543
- ```
1544
-
1545
- Note that `branch` is typically one of:
1546
- - `prod`
1547
- - `user.bob`
1548
- - `test.my_experiment`
1549
- - `prod.staging`
1550
-
1551
-
1552
- Parameters
1553
- ----------
1554
- flow : Union[str, Dict[str, str]], optional, default None
1555
- Upstream flow dependency for this flow.
1556
- flows : List[Union[str, Dict[str, str]]], default []
1557
- Upstream flow dependencies for this flow.
1558
- options : Dict[str, Any], default {}
1559
- Backend-specific configuration for tuning eventing behavior.
1560
- """
1561
- ...
1562
-
1563
- @typing.overload
1564
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1565
- ...
1566
-
1567
- 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] = {}):
1568
- """
1569
- Specifies the flow(s) that this flow depends on.
1570
-
1571
- ```
1572
- @trigger_on_finish(flow='FooFlow')
1573
- ```
1574
- or
1575
- ```
1576
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1577
- ```
1578
- This decorator respects the @project decorator and triggers the flow
1579
- when upstream runs within the same namespace complete successfully
1580
-
1581
- Additionally, you can specify project aware upstream flow dependencies
1582
- by specifying the fully qualified project_flow_name.
1583
- ```
1584
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1585
- ```
1586
- or
1587
- ```
1588
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1589
- ```
1590
-
1591
- You can also specify just the project or project branch (other values will be
1592
- inferred from the current project or project branch):
1593
- ```
1594
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1595
- ```
1596
-
1597
- Note that `branch` is typically one of:
1598
- - `prod`
1599
- - `user.bob`
1600
- - `test.my_experiment`
1601
- - `prod.staging`
1602
-
1603
-
1604
- Parameters
1605
- ----------
1606
- flow : Union[str, Dict[str, str]], optional, default None
1607
- Upstream flow dependency for this flow.
1608
- flows : List[Union[str, Dict[str, str]]], default []
1609
- Upstream flow dependencies for this flow.
1610
- options : Dict[str, Any], default {}
1611
- Backend-specific configuration for tuning eventing behavior.
1612
- """
1613
- ...
1614
-
1615
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1616
- """
1617
- Specifies what flows belong to the same project.
1618
-
1619
- A project-specific namespace is created for all flows that
1620
- use the same `@project(name)`.
1621
-
1622
-
1623
- Parameters
1624
- ----------
1625
- name : str
1626
- Project name. Make sure that the name is unique amongst all
1627
- projects that use the same production scheduler. The name may
1628
- contain only lowercase alphanumeric characters and underscores.
1629
-
1630
- branch : Optional[str], default None
1631
- The branch to use. If not specified, the branch is set to
1632
- `user.<username>` unless `production` is set to `True`. This can
1633
- also be set on the command line using `--branch` as a top-level option.
1634
- It is an error to specify `branch` in the decorator and on the command line.
1635
-
1636
- production : bool, default False
1637
- Whether or not the branch is the production branch. This can also be set on the
1638
- command line using `--production` as a top-level option. It is an error to specify
1639
- `production` in the decorator and on the command line.
1640
- The project branch name will be:
1641
- - if `branch` is specified:
1642
- - if `production` is True: `prod.<branch>`
1643
- - if `production` is False: `test.<branch>`
1644
- - if `branch` is not specified:
1645
- - if `production` is True: `prod`
1646
- - if `production` is False: `user.<username>`
1647
- """
1648
- ...
1649
-
1650
1691
  @typing.overload
1651
1692
  def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1652
1693
  """
@@ -1740,47 +1781,6 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1740
1781
  """
1741
1782
  ...
1742
1783
 
1743
- @typing.overload
1744
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1745
- """
1746
- Specifies the PyPI packages for all steps of the flow.
1747
-
1748
- Use `@pypi_base` to set common packages required by all
1749
- steps and use `@pypi` to specify step-specific overrides.
1750
-
1751
- Parameters
1752
- ----------
1753
- packages : Dict[str, str], default: {}
1754
- Packages to use for this flow. The key is the name of the package
1755
- and the value is the version to use.
1756
- python : str, optional, default: None
1757
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1758
- that the version used will correspond to the version of the Python interpreter used to start the run.
1759
- """
1760
- ...
1761
-
1762
- @typing.overload
1763
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1764
- ...
1765
-
1766
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1767
- """
1768
- Specifies the PyPI packages for all steps of the flow.
1769
-
1770
- Use `@pypi_base` to set common packages required by all
1771
- steps and use `@pypi` to specify step-specific overrides.
1772
-
1773
- Parameters
1774
- ----------
1775
- packages : Dict[str, str], default: {}
1776
- Packages to use for this flow. The key is the name of the package
1777
- and the value is the version to use.
1778
- python : str, optional, default: None
1779
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1780
- that the version used will correspond to the version of the Python interpreter used to start the run.
1781
- """
1782
- ...
1783
-
1784
1784
  @typing.overload
1785
1785
  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]]:
1786
1786
  """