ob-metaflow-stubs 6.0.4.7__py2.py3-none-any.whl → 6.0.4.8__py2.py3-none-any.whl

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

Potentially problematic release.


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

Files changed (249) hide show
  1. metaflow-stubs/__init__.pyi +883 -883
  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 +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +5 -5
  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 +50 -50
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +3 -3
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +4 -4
  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 +5 -5
  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 +3 -3
  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 +3 -3
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +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 +3 -3
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +5 -5
  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 +2 -2
  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 +5 -5
  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 +4 -4
  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 +3 -3
  89. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +2 -2
  90. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +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 +2 -2
  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 +4 -4
  99. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +3 -3
  100. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +2 -2
  101. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +2 -2
  102. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +2 -2
  103. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  104. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +2 -2
  105. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +2 -2
  106. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  107. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +2 -2
  108. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  109. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  110. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  111. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  112. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  113. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  114. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  115. metaflow-stubs/multicore_utils.pyi +2 -2
  116. metaflow-stubs/ob_internal.pyi +2 -2
  117. metaflow-stubs/parameters.pyi +5 -5
  118. metaflow-stubs/plugins/__init__.pyi +12 -12
  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 +2 -2
  128. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  129. metaflow-stubs/plugins/argo/argo_workflows.pyi +32 -3
  130. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  131. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  132. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +26 -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 +3 -3
  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 +3 -3
  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 +3 -3
  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 +3 -3
  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 +4 -4
  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 +3 -3
  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 +3 -2
  199. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  200. metaflow-stubs/plugins/ollama/__init__.pyi +3 -3
  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 +4 -4
  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 +3 -3
  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 +2 -2
  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 +3 -3
  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 +37 -6
  230. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  231. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  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 +5 -5
  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.7.dist-info → ob_metaflow_stubs-6.0.4.8.dist-info}/METADATA +1 -1
  246. ob_metaflow_stubs-6.0.4.8.dist-info/RECORD +249 -0
  247. ob_metaflow_stubs-6.0.4.7.dist-info/RECORD +0 -249
  248. {ob_metaflow_stubs-6.0.4.7.dist-info → ob_metaflow_stubs-6.0.4.8.dist-info}/WHEEL +0 -0
  249. {ob_metaflow_stubs-6.0.4.7.dist-info → ob_metaflow_stubs-6.0.4.8.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.21.4+obcheckpoint(0.2.4);ob(v1) #
4
- # Generated on 2025-07-25T18:05:15.016391 #
3
+ # MF version: 2.15.21.5+obcheckpoint(0.2.4);ob(v1) #
4
+ # Generated on 2025-07-29T18:33:44.107085 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -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 events as events
39
38
  from . import tuple_util as tuple_util
40
39
  from . import cards as cards
41
40
  from . import metaflow_git as metaflow_git
41
+ from . import events as events
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
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
48
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
+ from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_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
@@ -74,8 +74,8 @@ from .mf_extensions.outerbounds.plugins.checkpoint_datastores.nebius import nebi
74
74
  from .mf_extensions.outerbounds.plugins.checkpoint_datastores.coreweave import coreweave_checkpoints as coreweave_checkpoints
75
75
  from .mf_extensions.outerbounds.plugins.aws.assume_role_decorator import assume_role as assume_role
76
76
  from .mf_extensions.outerbounds.plugins.apps.core.deployer import AppDeployer as AppDeployer
77
- from . import cli_components as cli_components
78
77
  from . import system as system
78
+ from . import cli_components as cli_components
79
79
  from . import pylint_wrapper as pylint_wrapper
80
80
  from . import cli as cli
81
81
  from . import profilers as profilers
@@ -157,30 +157,103 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
157
157
  ...
158
158
 
159
159
  @typing.overload
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]]]:
160
+ 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]]]:
161
161
  """
162
- Specifies secrets to be retrieved and injected as environment variables prior to
163
- the execution of a step.
162
+ Specifies the resources needed when executing this step.
163
+
164
+ Use `@resources` to specify the resource requirements
165
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
166
+
167
+ You can choose the compute layer on the command line by executing e.g.
168
+ ```
169
+ python myflow.py run --with batch
170
+ ```
171
+ or
172
+ ```
173
+ python myflow.py run --with kubernetes
174
+ ```
175
+ which executes the flow on the desired system using the
176
+ requirements specified in `@resources`.
164
177
 
165
178
 
166
179
  Parameters
167
180
  ----------
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
181
+ cpu : int, default 1
182
+ Number of CPUs required for this step.
183
+ gpu : int, optional, default None
184
+ Number of GPUs required for this step.
185
+ disk : int, optional, default None
186
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
187
+ memory : int, default 4096
188
+ Memory size (in MB) required for this step.
189
+ shared_memory : int, optional, default None
190
+ The value for the size (in MiB) of the /dev/shm volume for this step.
191
+ This parameter maps to the `--shm-size` option in Docker.
172
192
  """
173
193
  ...
174
194
 
175
195
  @typing.overload
176
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
196
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
177
197
  ...
178
198
 
179
199
  @typing.overload
180
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
200
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
181
201
  ...
182
202
 
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):
203
+ 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):
204
+ """
205
+ Specifies the resources needed when executing this step.
206
+
207
+ Use `@resources` to specify the resource requirements
208
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
209
+
210
+ You can choose the compute layer on the command line by executing e.g.
211
+ ```
212
+ python myflow.py run --with batch
213
+ ```
214
+ or
215
+ ```
216
+ python myflow.py run --with kubernetes
217
+ ```
218
+ which executes the flow on the desired system using the
219
+ requirements specified in `@resources`.
220
+
221
+
222
+ Parameters
223
+ ----------
224
+ cpu : int, default 1
225
+ Number of CPUs required for this step.
226
+ gpu : int, optional, default None
227
+ Number of GPUs required for this step.
228
+ disk : int, optional, default None
229
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
230
+ memory : int, default 4096
231
+ Memory size (in MB) required for this step.
232
+ shared_memory : int, optional, default None
233
+ The value for the size (in MiB) of the /dev/shm volume for this step.
234
+ This parameter maps to the `--shm-size` option in Docker.
235
+ """
236
+ ...
237
+
238
+ @typing.overload
239
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
240
+ """
241
+ Internal decorator to support Fast bakery
242
+ """
243
+ ...
244
+
245
+ @typing.overload
246
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
247
+ ...
248
+
249
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
250
+ """
251
+ Internal decorator to support Fast bakery
252
+ """
253
+ ...
254
+
255
+ @typing.overload
256
+ 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]]]:
184
257
  """
185
258
  Specifies secrets to be retrieved and injected as environment variables prior to
186
259
  the execution of a step.
@@ -196,21 +269,25 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
196
269
  ...
197
270
 
198
271
  @typing.overload
199
- def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
200
- """
201
- Decorator prototype for all step decorators. This function gets specialized
202
- and imported for all decorators types by _import_plugin_decorators().
203
- """
272
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
204
273
  ...
205
274
 
206
275
  @typing.overload
207
- def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
276
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
208
277
  ...
209
278
 
210
- def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
279
+ 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):
211
280
  """
212
- Decorator prototype for all step decorators. This function gets specialized
213
- and imported for all decorators types by _import_plugin_decorators().
281
+ Specifies secrets to be retrieved and injected as environment variables prior to
282
+ the execution of a step.
283
+
284
+
285
+ Parameters
286
+ ----------
287
+ sources : List[Union[str, Dict[str, Any]]], default: []
288
+ List of secret specs, defining how the secrets are to be retrieved
289
+ role : str, optional, default: None
290
+ Role to use for fetching secrets
214
291
  """
215
292
  ...
216
293
 
@@ -361,148 +438,36 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
361
438
  """
362
439
  ...
363
440
 
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]]]:
365
- """
366
- Specifies that this step should execute on DGX cloud.
367
-
368
-
369
- Parameters
370
- ----------
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.
377
- """
378
- ...
379
-
380
441
  @typing.overload
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]]]:
442
+ 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]]]:
382
443
  """
383
- Enables loading / saving of models within a step.
384
-
385
- > Examples
386
- - Saving Models
387
- ```python
388
- @model
389
- @step
390
- def train(self):
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)
402
-
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
- ```
411
-
412
- - Loading models
413
- ```python
414
- @step
415
- def train(self):
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)
424
- ```
444
+ Specifies environment variables to be set prior to the execution of a step.
425
445
 
426
446
 
427
447
  Parameters
428
448
  ----------
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.
435
-
436
- temp_dir_root : str, default: None
437
- The root directory under which `current.model.loaded` will store loaded models
449
+ vars : Dict[str, str], default {}
450
+ Dictionary of environment variables to set.
438
451
  """
439
452
  ...
440
453
 
441
454
  @typing.overload
442
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
455
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
443
456
  ...
444
457
 
445
458
  @typing.overload
446
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
459
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
447
460
  ...
448
461
 
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):
462
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
450
463
  """
451
- Enables loading / saving of models within a step.
452
-
453
- > Examples
454
- - Saving Models
455
- ```python
456
- @model
457
- @step
458
- def train(self):
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)
470
-
471
- @model(load="my_model")
472
- @step
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
- ```
479
-
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)
492
- ```
464
+ Specifies environment variables to be set prior to the execution of a step.
493
465
 
494
466
 
495
467
  Parameters
496
468
  ----------
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.
503
-
504
- temp_dir_root : str, default: None
505
- The root directory under which `current.model.loaded` will store loaded models
469
+ vars : Dict[str, str], default {}
470
+ Dictionary of environment variables to set.
506
471
  """
507
472
  ...
508
473
 
@@ -565,190 +530,111 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
565
530
  """
566
531
  ...
567
532
 
568
- @typing.overload
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]]]:
533
+ 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]]]:
570
534
  """
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.
586
- """
587
- ...
588
-
589
- @typing.overload
590
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
591
- ...
592
-
593
- @typing.overload
594
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
595
- ...
596
-
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):
598
- """
599
- Creates a human-readable report, a Metaflow Card, after this step completes.
600
-
601
- Note that you may add multiple `@card` decorators in a step with different parameters.
602
-
603
-
604
- Parameters
605
- ----------
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.
614
- """
615
- ...
616
-
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]]]:
619
- """
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`.
535
+ Specifies that this step should execute on Kubernetes.
635
536
 
636
537
 
637
538
  Parameters
638
539
  ----------
639
540
  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.
541
+ Number of CPUs required for this step. If `@resources` is
542
+ also present, the maximum value from all decorators is used.
645
543
  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.
650
- """
651
- ...
652
-
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):
662
- """
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`.
678
-
679
-
680
- Parameters
681
- ----------
682
- cpu : int, default 1
683
- Number of CPUs required for this step.
544
+ Memory size (in MB) required for this step. If
545
+ `@resources` is also present, the maximum value from all decorators is
546
+ used.
547
+ disk : int, default 10240
548
+ Disk size (in MB) required for this step. If
549
+ `@resources` is also present, the maximum value from all decorators is
550
+ used.
551
+ image : str, optional, default None
552
+ Docker image to use when launching on Kubernetes. If not specified, and
553
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
554
+ not, a default Docker image mapping to the current version of Python is used.
555
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
556
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
557
+ image_pull_secrets: List[str], default []
558
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
559
+ Kubernetes image pull secrets to use when pulling container images
560
+ in Kubernetes.
561
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
562
+ Kubernetes service account to use when launching pod in Kubernetes.
563
+ secrets : List[str], optional, default None
564
+ Kubernetes secrets to use when launching pod in Kubernetes. These
565
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
566
+ in Metaflow configuration.
567
+ node_selector: Union[Dict[str,str], str], optional, default None
568
+ Kubernetes node selector(s) to apply to the pod running the task.
569
+ Can be passed in as a comma separated string of values e.g.
570
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
571
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
572
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
573
+ Kubernetes namespace to use when launching pod in Kubernetes.
684
574
  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.
575
+ Number of GPUs required for this step. A value of zero implies that
576
+ the scheduled node should not have GPUs.
577
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
578
+ The vendor of the GPUs to be used for this step.
579
+ tolerations : List[str], default []
580
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
581
+ Kubernetes tolerations to use when launching pod in Kubernetes.
582
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
583
+ Kubernetes labels to use when launching pod in Kubernetes.
584
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
585
+ Kubernetes annotations to use when launching pod in Kubernetes.
586
+ use_tmpfs : bool, default False
587
+ This enables an explicit tmpfs mount for this step.
588
+ tmpfs_tempdir : bool, default True
589
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
590
+ tmpfs_size : int, optional, default: None
591
+ The value for the size (in MiB) of the tmpfs mount for this step.
592
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
593
+ memory allocated for this step.
594
+ tmpfs_path : str, optional, default /metaflow_temp
595
+ Path to tmpfs mount for this step.
596
+ persistent_volume_claims : Dict[str, str], optional, default None
597
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
598
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
599
+ shared_memory: int, optional
600
+ Shared memory size (in MiB) required for this step
601
+ port: int, optional
602
+ Port number to specify in the Kubernetes job object
603
+ compute_pool : str, optional, default None
604
+ Compute pool to be used for for this step.
605
+ If not specified, any accessible compute pool within the perimeter is used.
606
+ hostname_resolution_timeout: int, default 10 * 60
607
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
608
+ Only applicable when @parallel is used.
609
+ qos: str, default: Burstable
610
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
611
+
612
+ security_context: Dict[str, Any], optional, default None
613
+ Container security context. Applies to the task container. Allows the following keys:
614
+ - privileged: bool, optional, default None
615
+ - allow_privilege_escalation: bool, optional, default None
616
+ - run_as_user: int, optional, default None
617
+ - run_as_group: int, optional, default None
618
+ - run_as_non_root: bool, optional, default None
693
619
  """
694
620
  ...
695
621
 
696
622
  @typing.overload
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]]]:
623
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
698
624
  """
699
- Specifies a timeout for your step.
700
-
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.
709
-
710
-
711
- Parameters
712
- ----------
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.
625
+ Decorator prototype for all step decorators. This function gets specialized
626
+ and imported for all decorators types by _import_plugin_decorators().
719
627
  """
720
628
  ...
721
629
 
722
630
  @typing.overload
723
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
724
- ...
725
-
726
- @typing.overload
727
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
631
+ def app_deploy(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
728
632
  ...
729
633
 
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):
634
+ def app_deploy(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
731
635
  """
732
- Specifies a timeout for your step.
733
-
734
- This decorator is useful if this step may hang indefinitely.
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.
739
-
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.
742
-
743
-
744
- Parameters
745
- ----------
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.
636
+ Decorator prototype for all step decorators. This function gets specialized
637
+ and imported for all decorators types by _import_plugin_decorators().
752
638
  """
753
639
  ...
754
640
 
@@ -803,93 +689,145 @@ def vllm(*, model: str, backend: str, openai_api_server: bool, debug: bool, card
803
689
  ...
804
690
 
805
691
  @typing.overload
806
- def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
692
+ 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]]]:
807
693
  """
808
- Specifies the number of times the task corresponding
809
- to a step needs to be retried.
810
-
811
- This decorator is useful for handling transient errors, such as networking issues.
812
- If your task contains operations that can't be retried safely, e.g. database updates,
813
- it is advisable to annotate it with `@retry(times=0)`.
694
+ Enables loading / saving of models within a step.
814
695
 
815
- This can be used in conjunction with the `@catch` decorator. The `@catch`
816
- decorator will execute a no-op task after all retries have been exhausted,
817
- ensuring that the flow execution can continue.
696
+ > Examples
697
+ - Saving Models
698
+ ```python
699
+ @model
700
+ @step
701
+ def train(self):
702
+ # current.model.save returns a dictionary reference to the model saved
703
+ self.my_model = current.model.save(
704
+ path_to_my_model,
705
+ label="my_model",
706
+ metadata={
707
+ "epochs": 10,
708
+ "batch-size": 32,
709
+ "learning-rate": 0.001,
710
+ }
711
+ )
712
+ self.next(self.test)
713
+
714
+ @model(load="my_model")
715
+ @step
716
+ def test(self):
717
+ # `current.model.loaded` returns a dictionary of the loaded models
718
+ # where the key is the name of the artifact and the value is the path to the model
719
+ print(os.listdir(current.model.loaded["my_model"]))
720
+ self.next(self.end)
721
+ ```
722
+
723
+ - Loading models
724
+ ```python
725
+ @step
726
+ def train(self):
727
+ # current.model.load returns the path to the model loaded
728
+ checkpoint_path = current.model.load(
729
+ self.checkpoint_key,
730
+ )
731
+ model_path = current.model.load(
732
+ self.model,
733
+ )
734
+ self.next(self.test)
735
+ ```
818
736
 
819
737
 
820
738
  Parameters
821
739
  ----------
822
- times : int, default 3
823
- Number of times to retry this task.
824
- minutes_between_retries : int, default 2
825
- Number of minutes between retries.
740
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
741
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
742
+ 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`.
743
+ 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
744
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
745
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
746
+
747
+ temp_dir_root : str, default: None
748
+ The root directory under which `current.model.loaded` will store loaded models
826
749
  """
827
750
  ...
828
751
 
829
752
  @typing.overload
830
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
753
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
831
754
  ...
832
755
 
833
756
  @typing.overload
834
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
757
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
835
758
  ...
836
759
 
837
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
760
+ 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):
838
761
  """
839
- Specifies the number of times the task corresponding
840
- to a step needs to be retried.
762
+ Enables loading / saving of models within a step.
841
763
 
842
- This decorator is useful for handling transient errors, such as networking issues.
843
- If your task contains operations that can't be retried safely, e.g. database updates,
844
- it is advisable to annotate it with `@retry(times=0)`.
764
+ > Examples
765
+ - Saving Models
766
+ ```python
767
+ @model
768
+ @step
769
+ def train(self):
770
+ # current.model.save returns a dictionary reference to the model saved
771
+ self.my_model = current.model.save(
772
+ path_to_my_model,
773
+ label="my_model",
774
+ metadata={
775
+ "epochs": 10,
776
+ "batch-size": 32,
777
+ "learning-rate": 0.001,
778
+ }
779
+ )
780
+ self.next(self.test)
845
781
 
846
- This can be used in conjunction with the `@catch` decorator. The `@catch`
847
- decorator will execute a no-op task after all retries have been exhausted,
848
- ensuring that the flow execution can continue.
782
+ @model(load="my_model")
783
+ @step
784
+ def test(self):
785
+ # `current.model.loaded` returns a dictionary of the loaded models
786
+ # where the key is the name of the artifact and the value is the path to the model
787
+ print(os.listdir(current.model.loaded["my_model"]))
788
+ self.next(self.end)
789
+ ```
790
+
791
+ - Loading models
792
+ ```python
793
+ @step
794
+ def train(self):
795
+ # current.model.load returns the path to the model loaded
796
+ checkpoint_path = current.model.load(
797
+ self.checkpoint_key,
798
+ )
799
+ model_path = current.model.load(
800
+ self.model,
801
+ )
802
+ self.next(self.test)
803
+ ```
849
804
 
850
805
 
851
806
  Parameters
852
807
  ----------
853
- times : int, default 3
854
- Number of times to retry this task.
855
- minutes_between_retries : int, default 2
856
- Number of minutes between retries.
857
- """
858
- ...
859
-
860
- @typing.overload
861
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
862
- """
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
808
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
809
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
810
+ 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`.
811
+ 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
812
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
813
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
814
+
815
+ temp_dir_root : str, default: None
816
+ The root directory under which `current.model.loaded` will store loaded models
883
817
  """
884
818
  ...
885
819
 
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):
820
+ 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]]]:
891
821
  """
892
- Internal decorator to support Fast bakery
822
+ Specifies that this step should execute on DGX cloud.
823
+
824
+
825
+ Parameters
826
+ ----------
827
+ gpu : int
828
+ Number of GPUs to use.
829
+ gpu_type : str
830
+ Type of Nvidia GPU to use.
893
831
  """
894
832
  ...
895
833
 
@@ -945,115 +883,184 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
945
883
  ...
946
884
 
947
885
  @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]]]:
886
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
949
887
  """
950
- Specifies environment variables to be set prior to the execution of a step.
951
-
952
-
953
- Parameters
954
- ----------
955
- vars : Dict[str, str], default {}
956
- Dictionary of environment variables to set.
888
+ Decorator prototype for all step decorators. This function gets specialized
889
+ and imported for all decorators types by _import_plugin_decorators().
957
890
  """
958
891
  ...
959
892
 
960
893
  @typing.overload
961
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
894
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
962
895
  ...
963
896
 
964
- @typing.overload
965
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
897
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
898
+ """
899
+ Decorator prototype for all step decorators. This function gets specialized
900
+ and imported for all decorators types by _import_plugin_decorators().
901
+ """
966
902
  ...
967
903
 
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] = {}):
904
+ 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]]]:
969
905
  """
970
- Specifies environment variables to be set prior to the execution of a step.
906
+ Specifies that this step should execute on DGX cloud.
971
907
 
972
908
 
973
909
  Parameters
974
910
  ----------
975
- vars : Dict[str, str], default {}
976
- Dictionary of environment variables to set.
911
+ gpu : int
912
+ Number of GPUs to use.
913
+ gpu_type : str
914
+ Type of Nvidia GPU to use.
915
+ queue_timeout : int
916
+ Time to keep the job in NVCF's queue.
977
917
  """
978
918
  ...
979
919
 
980
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
920
+ 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]]]:
981
921
  """
982
- Decorator that helps cache, version and store models/datasets from huggingface hub.
983
-
984
- > Examples
922
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
985
923
 
986
- **Usage: creating references of models from huggingface that may be loaded in downstream steps**
987
- ```python
988
- @huggingface_hub
989
- @step
990
- def pull_model_from_huggingface(self):
991
- # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
992
- # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
993
- # with the same `repo_id` in the backend storage, it will not download the model again. The return
994
- # value of the function is a reference to the model in the backend storage.
995
- # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
924
+ User code call
925
+ --------------
926
+ @ollama(
927
+ models=[...],
928
+ ...
929
+ )
996
930
 
997
- self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
998
- self.llama_model = current.huggingface_hub.snapshot_download(
999
- repo_id=self.model_id,
1000
- allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1001
- )
1002
- self.next(self.train)
1003
- ```
931
+ Valid backend options
932
+ ---------------------
933
+ - 'local': Run as a separate process on the local task machine.
934
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
935
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
1004
936
 
1005
- **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
1006
- ```python
1007
- @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1008
- @step
1009
- def pull_model_from_huggingface(self):
1010
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1011
- ```
937
+ Valid model options
938
+ -------------------
939
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
1012
940
 
1013
- ```python
1014
- @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1015
- @step
1016
- def finetune_model(self):
1017
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1018
- # path_to_model will be /my-directory
1019
- ```
1020
941
 
1021
- ```python
1022
- # Takes all the arguments passed to `snapshot_download`
1023
- # except for `local_dir`
1024
- @huggingface_hub(load=[
1025
- {
1026
- "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1027
- },
1028
- {
1029
- "repo_id": "myorg/mistral-lora",
1030
- "repo_type": "model",
1031
- },
1032
- ])
1033
- @step
1034
- def finetune_model(self):
1035
- path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1036
- # path_to_model will be /my-directory
1037
- ```
942
+ Parameters
943
+ ----------
944
+ models: list[str]
945
+ List of Ollama containers running models in sidecars.
946
+ backend: str
947
+ Determines where and how to run the Ollama process.
948
+ force_pull: bool
949
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
950
+ cache_update_policy: str
951
+ Cache update policy: "auto", "force", or "never".
952
+ force_cache_update: bool
953
+ Simple override for "force" cache update policy.
954
+ debug: bool
955
+ Whether to turn on verbose debugging logs.
956
+ circuit_breaker_config: dict
957
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
958
+ timeout_config: dict
959
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
960
+ """
961
+ ...
962
+
963
+ @typing.overload
964
+ 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]]]:
965
+ """
966
+ Creates a human-readable report, a Metaflow Card, after this step completes.
967
+
968
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1038
969
 
1039
970
 
1040
971
  Parameters
1041
972
  ----------
1042
- temp_dir_root : str, optional
1043
- The root directory that will hold the temporary directory where objects will be downloaded.
973
+ type : str, default 'default'
974
+ Card type.
975
+ id : str, optional, default None
976
+ If multiple cards are present, use this id to identify this card.
977
+ options : Dict[str, Any], default {}
978
+ Options passed to the card. The contents depend on the card type.
979
+ timeout : int, default 45
980
+ Interrupt reporting if it takes more than this many seconds.
981
+ """
982
+ ...
983
+
984
+ @typing.overload
985
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
986
+ ...
987
+
988
+ @typing.overload
989
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
990
+ ...
991
+
992
+ 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):
993
+ """
994
+ Creates a human-readable report, a Metaflow Card, after this step completes.
1044
995
 
1045
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1046
- The list of repos (models/datasets) to load.
996
+ Note that you may add multiple `@card` decorators in a step with different parameters.
1047
997
 
1048
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1049
998
 
1050
- - If repo (model/dataset) is not found in the datastore:
1051
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1052
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1053
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
999
+ Parameters
1000
+ ----------
1001
+ type : str, default 'default'
1002
+ Card type.
1003
+ id : str, optional, default None
1004
+ If multiple cards are present, use this id to identify this card.
1005
+ options : Dict[str, Any], default {}
1006
+ Options passed to the card. The contents depend on the card type.
1007
+ timeout : int, default 45
1008
+ Interrupt reporting if it takes more than this many seconds.
1009
+ """
1010
+ ...
1011
+
1012
+ @typing.overload
1013
+ def retry(*, times: int = 3, minutes_between_retries: int = 2) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1014
+ """
1015
+ Specifies the number of times the task corresponding
1016
+ to a step needs to be retried.
1054
1017
 
1055
- - If repo is found in the datastore:
1056
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
1018
+ This decorator is useful for handling transient errors, such as networking issues.
1019
+ If your task contains operations that can't be retried safely, e.g. database updates,
1020
+ it is advisable to annotate it with `@retry(times=0)`.
1021
+
1022
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1023
+ decorator will execute a no-op task after all retries have been exhausted,
1024
+ ensuring that the flow execution can continue.
1025
+
1026
+
1027
+ Parameters
1028
+ ----------
1029
+ times : int, default 3
1030
+ Number of times to retry this task.
1031
+ minutes_between_retries : int, default 2
1032
+ Number of minutes between retries.
1033
+ """
1034
+ ...
1035
+
1036
+ @typing.overload
1037
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1038
+ ...
1039
+
1040
+ @typing.overload
1041
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1042
+ ...
1043
+
1044
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
1045
+ """
1046
+ Specifies the number of times the task corresponding
1047
+ to a step needs to be retried.
1048
+
1049
+ This decorator is useful for handling transient errors, such as networking issues.
1050
+ If your task contains operations that can't be retried safely, e.g. database updates,
1051
+ it is advisable to annotate it with `@retry(times=0)`.
1052
+
1053
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1054
+ decorator will execute a no-op task after all retries have been exhausted,
1055
+ ensuring that the flow execution can continue.
1056
+
1057
+
1058
+ Parameters
1059
+ ----------
1060
+ times : int, default 3
1061
+ Number of times to retry this task.
1062
+ minutes_between_retries : int, default 2
1063
+ Number of minutes between retries.
1057
1064
  """
1058
1065
  ...
1059
1066
 
@@ -1076,60 +1083,62 @@ def test_append_card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag]
1076
1083
  """
1077
1084
  ...
1078
1085
 
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]]]:
1086
+ @typing.overload
1087
+ 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]]]:
1080
1088
  """
1081
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
1089
+ Specifies a timeout for your step.
1082
1090
 
1083
- User code call
1084
- --------------
1085
- @ollama(
1086
- models=[...],
1087
- ...
1088
- )
1091
+ This decorator is useful if this step may hang indefinitely.
1089
1092
 
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.
1093
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1094
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1095
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1095
1096
 
1096
- Valid model options
1097
- -------------------
1098
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
1097
+ Note that all the values specified in parameters are added together so if you specify
1098
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1099
1099
 
1100
1100
 
1101
1101
  Parameters
1102
1102
  ----------
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.
1103
+ seconds : int, default 0
1104
+ Number of seconds to wait prior to timing out.
1105
+ minutes : int, default 0
1106
+ Number of minutes to wait prior to timing out.
1107
+ hours : int, default 0
1108
+ Number of hours to wait prior to timing out.
1119
1109
  """
1120
1110
  ...
1121
1111
 
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]]]:
1112
+ @typing.overload
1113
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1114
+ ...
1115
+
1116
+ @typing.overload
1117
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
1118
+ ...
1119
+
1120
+ 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):
1123
1121
  """
1124
- Specifies that this step should execute on DGX cloud.
1122
+ Specifies a timeout for your step.
1123
+
1124
+ This decorator is useful if this step may hang indefinitely.
1125
+
1126
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1127
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
1128
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
1129
+
1130
+ Note that all the values specified in parameters are added together so if you specify
1131
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1125
1132
 
1126
1133
 
1127
1134
  Parameters
1128
1135
  ----------
1129
- gpu : int
1130
- Number of GPUs to use.
1131
- gpu_type : str
1132
- Type of Nvidia GPU to use.
1136
+ seconds : int, default 0
1137
+ Number of seconds to wait prior to timing out.
1138
+ minutes : int, default 0
1139
+ Number of minutes to wait prior to timing out.
1140
+ hours : int, default 0
1141
+ Number of hours to wait prior to timing out.
1133
1142
  """
1134
1143
  ...
1135
1144
 
@@ -1184,92 +1193,291 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
1184
1193
  """
1185
1194
  ...
1186
1195
 
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]]]:
1196
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
1188
1197
  """
1189
- Specifies that this step should execute on Kubernetes.
1198
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
1199
+
1200
+ > Examples
1201
+
1202
+ **Usage: creating references of models from huggingface that may be loaded in downstream steps**
1203
+ ```python
1204
+ @huggingface_hub
1205
+ @step
1206
+ def pull_model_from_huggingface(self):
1207
+ # `current.huggingface_hub.snapshot_download` downloads the model from the Hugging Face Hub
1208
+ # and saves it in the backend storage based on the model's `repo_id`. If there exists a model
1209
+ # with the same `repo_id` in the backend storage, it will not download the model again. The return
1210
+ # value of the function is a reference to the model in the backend storage.
1211
+ # This reference can be used to load the model in the subsequent steps via `@model(load=["llama_model"])`
1212
+
1213
+ self.model_id = "mistralai/Mistral-7B-Instruct-v0.1"
1214
+ self.llama_model = current.huggingface_hub.snapshot_download(
1215
+ repo_id=self.model_id,
1216
+ allow_patterns=["*.safetensors", "*.json", "tokenizer.*"],
1217
+ )
1218
+ self.next(self.train)
1219
+ ```
1220
+
1221
+ **Usage: loading models directly from huggingface hub or from cache (from metaflow's datastore)**
1222
+ ```python
1223
+ @huggingface_hub(load=["mistralai/Mistral-7B-Instruct-v0.1"])
1224
+ @step
1225
+ def pull_model_from_huggingface(self):
1226
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1227
+ ```
1228
+
1229
+ ```python
1230
+ @huggingface_hub(load=[("mistralai/Mistral-7B-Instruct-v0.1", "/my-directory"), ("myorg/mistral-lora, "/my-lora-directory")])
1231
+ @step
1232
+ def finetune_model(self):
1233
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1234
+ # path_to_model will be /my-directory
1235
+ ```
1236
+
1237
+ ```python
1238
+ # Takes all the arguments passed to `snapshot_download`
1239
+ # except for `local_dir`
1240
+ @huggingface_hub(load=[
1241
+ {
1242
+ "repo_id": "mistralai/Mistral-7B-Instruct-v0.1",
1243
+ },
1244
+ {
1245
+ "repo_id": "myorg/mistral-lora",
1246
+ "repo_type": "model",
1247
+ },
1248
+ ])
1249
+ @step
1250
+ def finetune_model(self):
1251
+ path_to_model = current.huggingface_hub.loaded["mistralai/Mistral-7B-Instruct-v0.1"]
1252
+ # path_to_model will be /my-directory
1253
+ ```
1190
1254
 
1191
1255
 
1192
1256
  Parameters
1193
1257
  ----------
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
1258
+ temp_dir_root : str, optional
1259
+ The root directory that will hold the temporary directory where objects will be downloaded.
1260
+
1261
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
1262
+ The list of repos (models/datasets) to load.
1263
+
1264
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
1265
+
1266
+ - If repo (model/dataset) is not found in the datastore:
1267
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
1268
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
1269
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
1270
+
1271
+ - If repo is found in the datastore:
1272
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
1273
+ """
1274
+ ...
1275
+
1276
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1277
+ """
1278
+ Allows setting external datastores to save data for the
1279
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1280
+
1281
+ This decorator is useful when users wish to save data to a different datastore
1282
+ than what is configured in Metaflow. This can be for variety of reasons:
1283
+
1284
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1285
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1286
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1287
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1288
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1289
+
1290
+ Usage:
1291
+ ----------
1292
+
1293
+ - Using a custom IAM role to access the datastore.
1294
+
1295
+ ```python
1296
+ @with_artifact_store(
1297
+ type="s3",
1298
+ config=lambda: {
1299
+ "root": "s3://my-bucket-foo/path/to/root",
1300
+ "role_arn": ROLE,
1301
+ },
1302
+ )
1303
+ class MyFlow(FlowSpec):
1304
+
1305
+ @checkpoint
1306
+ @step
1307
+ def start(self):
1308
+ with open("my_file.txt", "w") as f:
1309
+ f.write("Hello, World!")
1310
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1311
+ self.next(self.end)
1312
+
1313
+ ```
1314
+
1315
+ - Using credentials to access the s3-compatible datastore.
1316
+
1317
+ ```python
1318
+ @with_artifact_store(
1319
+ type="s3",
1320
+ config=lambda: {
1321
+ "root": "s3://my-bucket-foo/path/to/root",
1322
+ "client_params": {
1323
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1324
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1325
+ },
1326
+ },
1327
+ )
1328
+ class MyFlow(FlowSpec):
1329
+
1330
+ @checkpoint
1331
+ @step
1332
+ def start(self):
1333
+ with open("my_file.txt", "w") as f:
1334
+ f.write("Hello, World!")
1335
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1336
+ self.next(self.end)
1337
+
1338
+ ```
1339
+
1340
+ - Accessing objects stored in external datastores after task execution.
1341
+
1342
+ ```python
1343
+ run = Run("CheckpointsTestsFlow/8992")
1344
+ with artifact_store_from(run=run, config={
1345
+ "client_params": {
1346
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1347
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1348
+ },
1349
+ }):
1350
+ with Checkpoint() as cp:
1351
+ latest = cp.list(
1352
+ task=run["start"].task
1353
+ )[0]
1354
+ print(latest)
1355
+ cp.load(
1356
+ latest,
1357
+ "test-checkpoints"
1358
+ )
1359
+
1360
+ task = Task("TorchTuneFlow/8484/train/53673")
1361
+ with artifact_store_from(run=run, config={
1362
+ "client_params": {
1363
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1364
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1365
+ },
1366
+ }):
1367
+ load_model(
1368
+ task.data.model_ref,
1369
+ "test-models"
1370
+ )
1371
+ ```
1372
+ Parameters:
1373
+ ----------
1374
+
1375
+ type: str
1376
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1377
+
1378
+ config: dict or Callable
1379
+ Dictionary of configuration options for the datastore. The following keys are required:
1380
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1381
+ - example: 's3://bucket-name/path/to/root'
1382
+ - example: 'gs://bucket-name/path/to/root'
1383
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1384
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1385
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1386
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1387
+ """
1388
+ ...
1389
+
1390
+ @typing.overload
1391
+ def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1392
+ """
1393
+ Specifies the times when the flow should be run when running on a
1394
+ production scheduler.
1395
+
1396
+
1397
+ Parameters
1398
+ ----------
1399
+ hourly : bool, default False
1400
+ Run the workflow hourly.
1401
+ daily : bool, default True
1402
+ Run the workflow daily.
1403
+ weekly : bool, default False
1404
+ Run the workflow weekly.
1405
+ cron : str, optional, default None
1406
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1407
+ specified by this expression.
1408
+ timezone : str, optional, default None
1409
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1410
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1411
+ """
1412
+ ...
1413
+
1414
+ @typing.overload
1415
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1416
+ ...
1417
+
1418
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1419
+ """
1420
+ Specifies the times when the flow should be run when running on a
1421
+ production scheduler.
1265
1422
 
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
1423
+
1424
+ Parameters
1425
+ ----------
1426
+ hourly : bool, default False
1427
+ Run the workflow hourly.
1428
+ daily : bool, default True
1429
+ Run the workflow daily.
1430
+ weekly : bool, default False
1431
+ Run the workflow weekly.
1432
+ cron : str, optional, default None
1433
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1434
+ specified by this expression.
1435
+ timezone : str, optional, default None
1436
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1437
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1438
+ """
1439
+ ...
1440
+
1441
+ 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]]:
1442
+ """
1443
+ 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.
1444
+ 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.
1445
+
1446
+
1447
+ Parameters
1448
+ ----------
1449
+ timeout : int
1450
+ Time, in seconds before the task times out and fails. (Default: 3600)
1451
+ poke_interval : int
1452
+ Time in seconds that the job should wait in between each try. (Default: 60)
1453
+ mode : str
1454
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1455
+ exponential_backoff : bool
1456
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1457
+ pool : str
1458
+ the slot pool this task should run in,
1459
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1460
+ soft_fail : bool
1461
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1462
+ name : str
1463
+ Name of the sensor on Airflow
1464
+ description : str
1465
+ Description of sensor in the Airflow UI
1466
+ external_dag_id : str
1467
+ The dag_id that contains the task you want to wait for.
1468
+ external_task_ids : List[str]
1469
+ The list of task_ids that you want to wait for.
1470
+ If None (default value) the sensor waits for the DAG. (Default: None)
1471
+ allowed_states : List[str]
1472
+ Iterable of allowed states, (Default: ['success'])
1473
+ failed_states : List[str]
1474
+ Iterable of failed or dis-allowed states. (Default: None)
1475
+ execution_delta : datetime.timedelta
1476
+ time difference with the previous execution to look at,
1477
+ the default is the same logical date as the current task or DAG. (Default: None)
1478
+ check_existence: bool
1479
+ Set to True to check if the external task exists or check if
1480
+ the DAG to wait for exists. (Default: True)
1273
1481
  """
1274
1482
  ...
1275
1483
 
@@ -1316,41 +1524,6 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1316
1524
  """
1317
1525
  ...
1318
1526
 
1319
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1320
- """
1321
- Specifies what flows belong to the same project.
1322
-
1323
- A project-specific namespace is created for all flows that
1324
- use the same `@project(name)`.
1325
-
1326
-
1327
- Parameters
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.
1333
-
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.
1339
-
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
1527
  @typing.overload
1355
1528
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1356
1529
  """
@@ -1393,59 +1566,58 @@ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packag
1393
1566
  ...
1394
1567
 
1395
1568
  @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]]:
1569
+ 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]]:
1397
1570
  """
1398
- Specifies the flow(s) that this flow depends on.
1571
+ Specifies the Conda environment for all steps of the flow.
1399
1572
 
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
1573
+ Use `@conda_base` to set common libraries required by all
1574
+ steps and use `@conda` to specify step-specific additions.
1409
1575
 
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
- ```
1419
1576
 
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
- ```
1577
+ Parameters
1578
+ ----------
1579
+ packages : Dict[str, str], default {}
1580
+ Packages to use for this flow. The key is the name of the package
1581
+ and the value is the version to use.
1582
+ libraries : Dict[str, str], default {}
1583
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1584
+ python : str, optional, default None
1585
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1586
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1587
+ disabled : bool, default False
1588
+ If set to True, disables Conda.
1589
+ """
1590
+ ...
1591
+
1592
+ @typing.overload
1593
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1594
+ ...
1595
+
1596
+ 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):
1597
+ """
1598
+ Specifies the Conda environment for all steps of the flow.
1425
1599
 
1426
- Note that `branch` is typically one of:
1427
- - `prod`
1428
- - `user.bob`
1429
- - `test.my_experiment`
1430
- - `prod.staging`
1600
+ Use `@conda_base` to set common libraries required by all
1601
+ steps and use `@conda` to specify step-specific additions.
1431
1602
 
1432
1603
 
1433
1604
  Parameters
1434
1605
  ----------
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.
1606
+ packages : Dict[str, str], default {}
1607
+ Packages to use for this flow. The key is the name of the package
1608
+ and the value is the version to use.
1609
+ libraries : Dict[str, str], default {}
1610
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1611
+ python : str, optional, default None
1612
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1613
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1614
+ disabled : bool, default False
1615
+ If set to True, disables Conda.
1441
1616
  """
1442
1617
  ...
1443
1618
 
1444
1619
  @typing.overload
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] = {}):
1620
+ 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]]:
1449
1621
  """
1450
1622
  Specifies the flow(s) that this flow depends on.
1451
1623
 
@@ -1493,97 +1665,90 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1493
1665
  """
1494
1666
  ...
1495
1667
 
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]]:
1497
- """
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.
1500
-
1501
-
1502
- Parameters
1503
- ----------
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)
1536
- """
1668
+ @typing.overload
1669
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1537
1670
  ...
1538
1671
 
1539
- @typing.overload
1540
- def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1672
+ 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] = {}):
1541
1673
  """
1542
- Specifies the times when the flow should be run when running on a
1543
- production scheduler.
1674
+ Specifies the flow(s) that this flow depends on.
1675
+
1676
+ ```
1677
+ @trigger_on_finish(flow='FooFlow')
1678
+ ```
1679
+ or
1680
+ ```
1681
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1682
+ ```
1683
+ This decorator respects the @project decorator and triggers the flow
1684
+ when upstream runs within the same namespace complete successfully
1685
+
1686
+ Additionally, you can specify project aware upstream flow dependencies
1687
+ by specifying the fully qualified project_flow_name.
1688
+ ```
1689
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1690
+ ```
1691
+ or
1692
+ ```
1693
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1694
+ ```
1695
+
1696
+ You can also specify just the project or project branch (other values will be
1697
+ inferred from the current project or project branch):
1698
+ ```
1699
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1700
+ ```
1701
+
1702
+ Note that `branch` is typically one of:
1703
+ - `prod`
1704
+ - `user.bob`
1705
+ - `test.my_experiment`
1706
+ - `prod.staging`
1544
1707
 
1545
1708
 
1546
1709
  Parameters
1547
1710
  ----------
1548
- hourly : bool, default False
1549
- Run the workflow hourly.
1550
- daily : bool, default True
1551
- Run the workflow daily.
1552
- weekly : bool, default False
1553
- Run the workflow weekly.
1554
- cron : str, optional, default None
1555
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1556
- specified by this expression.
1557
- timezone : str, optional, default None
1558
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1559
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1711
+ flow : Union[str, Dict[str, str]], optional, default None
1712
+ Upstream flow dependency for this flow.
1713
+ flows : List[Union[str, Dict[str, str]]], default []
1714
+ Upstream flow dependencies for this flow.
1715
+ options : Dict[str, Any], default {}
1716
+ Backend-specific configuration for tuning eventing behavior.
1560
1717
  """
1561
1718
  ...
1562
1719
 
1563
- @typing.overload
1564
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1565
- ...
1566
-
1567
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1720
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1568
1721
  """
1569
- Specifies the times when the flow should be run when running on a
1570
- production scheduler.
1722
+ Specifies what flows belong to the same project.
1723
+
1724
+ A project-specific namespace is created for all flows that
1725
+ use the same `@project(name)`.
1571
1726
 
1572
1727
 
1573
1728
  Parameters
1574
1729
  ----------
1575
- hourly : bool, default False
1576
- Run the workflow hourly.
1577
- daily : bool, default True
1578
- Run the workflow daily.
1579
- weekly : bool, default False
1580
- Run the workflow weekly.
1581
- cron : str, optional, default None
1582
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1583
- specified by this expression.
1584
- timezone : str, optional, default None
1585
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1586
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1730
+ name : str
1731
+ Project name. Make sure that the name is unique amongst all
1732
+ projects that use the same production scheduler. The name may
1733
+ contain only lowercase alphanumeric characters and underscores.
1734
+
1735
+ branch : Optional[str], default None
1736
+ The branch to use. If not specified, the branch is set to
1737
+ `user.<username>` unless `production` is set to `True`. This can
1738
+ also be set on the command line using `--branch` as a top-level option.
1739
+ It is an error to specify `branch` in the decorator and on the command line.
1740
+
1741
+ production : bool, default False
1742
+ Whether or not the branch is the production branch. This can also be set on the
1743
+ command line using `--production` as a top-level option. It is an error to specify
1744
+ `production` in the decorator and on the command line.
1745
+ The project branch name will be:
1746
+ - if `branch` is specified:
1747
+ - if `production` is True: `prod.<branch>`
1748
+ - if `production` is False: `test.<branch>`
1749
+ - if `branch` is not specified:
1750
+ - if `production` is True: `prod`
1751
+ - if `production` is False: `user.<username>`
1587
1752
  """
1588
1753
  ...
1589
1754
 
@@ -1680,170 +1845,5 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1680
1845
  """
1681
1846
  ...
1682
1847
 
1683
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1684
- """
1685
- Allows setting external datastores to save data for the
1686
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1687
-
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:
1690
-
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:
1698
- ----------
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')
1794
- """
1795
- ...
1796
-
1797
- @typing.overload
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]]:
1799
- """
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.
1804
-
1805
-
1806
- Parameters
1807
- ----------
1808
- packages : Dict[str, str], default {}
1809
- Packages to use for this flow. The key is the name of the package
1810
- and the value is the version to use.
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
1814
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
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.
1818
- """
1819
- ...
1820
-
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):
1826
- """
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.
1831
-
1832
-
1833
- Parameters
1834
- ----------
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.
1845
- """
1846
- ...
1847
-
1848
1848
  pkg_name: str
1849
1849