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