ob-metaflow-stubs 6.0.3.188rc0__py2.py3-none-any.whl → 6.0.3.188rc2__py2.py3-none-any.whl

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