ob-metaflow-stubs 6.0.3.180rc0__py2.py3-none-any.whl → 6.0.3.180rc1__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 (220) hide show
  1. metaflow-stubs/__init__.pyi +821 -821
  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 +1 -1
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +114 -114
  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 +2 -2
  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 +4 -4
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +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 +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +2 -2
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  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 +3 -3
  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 +3 -3
  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 +1 -1
  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 +2 -2
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/aws/assume_role_decorator.pyi +1 -1
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  89. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  91. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  92. metaflow-stubs/multicore_utils.pyi +1 -1
  93. metaflow-stubs/ob_internal.pyi +1 -1
  94. metaflow-stubs/parameters.pyi +3 -3
  95. metaflow-stubs/plugins/__init__.pyi +11 -11
  96. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  99. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  100. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  101. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  102. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  103. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  104. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  105. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  106. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  107. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +1 -1
  108. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +2 -2
  109. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  110. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  111. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  112. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  113. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  114. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  115. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  116. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  117. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  118. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +1 -1
  119. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  121. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  122. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  123. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  124. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  125. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  126. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  128. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  129. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +1 -1
  130. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  131. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  132. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  133. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_client.pyi +1 -1
  135. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  140. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  142. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  143. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  144. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  145. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  146. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  147. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  148. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  149. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  151. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  152. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  153. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  154. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  155. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  156. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  157. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  158. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  159. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  160. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  161. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +1 -1
  163. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  164. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  165. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  166. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  169. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  170. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  171. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  172. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  173. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  174. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  175. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/perimeters.pyi +1 -1
  177. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  179. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  181. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  182. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  184. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  185. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  186. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  187. metaflow-stubs/plugins/secrets/__init__.pyi +1 -1
  188. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +1 -1
  189. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  190. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  192. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +1 -1
  193. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  194. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  195. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  196. metaflow-stubs/plugins/uv/uv_environment.pyi +1 -1
  197. metaflow-stubs/profilers/__init__.pyi +1 -1
  198. metaflow-stubs/pylint_wrapper.pyi +1 -1
  199. metaflow-stubs/runner/__init__.pyi +1 -1
  200. metaflow-stubs/runner/deployer.pyi +28 -28
  201. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  202. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  203. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  204. metaflow-stubs/runner/nbrun.pyi +1 -1
  205. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  206. metaflow-stubs/runner/utils.pyi +2 -2
  207. metaflow-stubs/system/__init__.pyi +1 -1
  208. metaflow-stubs/system/system_logger.pyi +2 -2
  209. metaflow-stubs/system/system_monitor.pyi +1 -1
  210. metaflow-stubs/tagging_util.pyi +1 -1
  211. metaflow-stubs/tuple_util.pyi +1 -1
  212. metaflow-stubs/user_configs/__init__.pyi +1 -1
  213. metaflow-stubs/user_configs/config_decorators.pyi +3 -3
  214. metaflow-stubs/user_configs/config_options.pyi +3 -3
  215. metaflow-stubs/user_configs/config_parameters.pyi +4 -4
  216. {ob_metaflow_stubs-6.0.3.180rc0.dist-info → ob_metaflow_stubs-6.0.3.180rc1.dist-info}/METADATA +1 -1
  217. ob_metaflow_stubs-6.0.3.180rc1.dist-info/RECORD +220 -0
  218. ob_metaflow_stubs-6.0.3.180rc0.dist-info/RECORD +0 -220
  219. {ob_metaflow_stubs-6.0.3.180rc0.dist-info → ob_metaflow_stubs-6.0.3.180rc1.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.180rc0.dist-info → ob_metaflow_stubs-6.0.3.180rc1.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.17.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-17T08:34:56.575242 #
4
+ # Generated on 2025-06-17T09:48:39.040005 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,18 +35,18 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import events as events
39
- from . import tuple_util as tuple_util
40
38
  from . import cards as cards
41
39
  from . import metaflow_git as metaflow_git
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 pyproject_toml_parser as pyproject_toml_parser
47
48
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
49
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
49
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_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
@@ -155,28 +155,20 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
155
  """
156
156
  ...
157
157
 
158
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
158
+ @typing.overload
159
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
159
160
  """
160
- Decorator that helps cache, version and store models/datasets from huggingface hub.
161
-
162
-
163
- Parameters
164
- ----------
165
- temp_dir_root : str, optional
166
- The root directory that will hold the temporary directory where objects will be downloaded.
167
-
168
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
169
- The list of repos (models/datasets) to load.
170
-
171
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
172
-
173
- - If repo (model/dataset) is not found in the datastore:
174
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
175
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
176
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
177
-
178
- - If repo is found in the datastore:
179
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
161
+ Internal decorator to support Fast bakery
162
+ """
163
+ ...
164
+
165
+ @typing.overload
166
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
167
+ ...
168
+
169
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
170
+ """
171
+ Internal decorator to support Fast bakery
180
172
  """
181
173
  ...
182
174
 
@@ -224,213 +216,213 @@ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy:
224
216
  ...
225
217
 
226
218
  @typing.overload
227
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> 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
+ 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]]]:
228
220
  """
229
- Specifies secrets to be retrieved and injected as environment variables prior to
230
- the execution of a step.
221
+ Specifies the Conda environment for the step.
222
+
223
+ Information in this decorator will augment any
224
+ attributes set in the `@conda_base` flow-level decorator. Hence,
225
+ you can use `@conda_base` to set packages required by all
226
+ steps and use `@conda` to specify step-specific overrides.
231
227
 
232
228
 
233
229
  Parameters
234
230
  ----------
235
- sources : List[Union[str, Dict[str, Any]]], default: []
236
- List of secret specs, defining how the secrets are to be retrieved
231
+ packages : Dict[str, str], default {}
232
+ Packages to use for this step. The key is the name of the package
233
+ and the value is the version to use.
234
+ libraries : Dict[str, str], default {}
235
+ Supported for backward compatibility. When used with packages, packages will take precedence.
236
+ python : str, optional, default None
237
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
238
+ that the version used will correspond to the version of the Python interpreter used to start the run.
239
+ disabled : bool, default False
240
+ If set to True, disables @conda.
237
241
  """
238
242
  ...
239
243
 
240
244
  @typing.overload
241
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
245
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
242
246
  ...
243
247
 
244
248
  @typing.overload
245
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
249
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
246
250
  ...
247
251
 
248
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
252
+ 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):
249
253
  """
250
- Specifies secrets to be retrieved and injected as environment variables prior to
251
- the execution of a step.
254
+ Specifies the Conda environment for the step.
255
+
256
+ Information in this decorator will augment any
257
+ attributes set in the `@conda_base` flow-level decorator. Hence,
258
+ you can use `@conda_base` to set packages required by all
259
+ steps and use `@conda` to specify step-specific overrides.
252
260
 
253
261
 
254
262
  Parameters
255
263
  ----------
256
- sources : List[Union[str, Dict[str, Any]]], default: []
257
- List of secret specs, defining how the secrets are to be retrieved
264
+ packages : Dict[str, str], default {}
265
+ Packages to use for this step. The key is the name of the package
266
+ and the value is the version to use.
267
+ libraries : Dict[str, str], default {}
268
+ Supported for backward compatibility. When used with packages, packages will take precedence.
269
+ python : str, optional, default None
270
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
271
+ that the version used will correspond to the version of the Python interpreter used to start the run.
272
+ disabled : bool, default False
273
+ If set to True, disables @conda.
258
274
  """
259
275
  ...
260
276
 
261
277
  @typing.overload
262
- 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]]]:
278
+ 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]]]:
263
279
  """
264
- Specifies the number of times the task corresponding
265
- to a step needs to be retried.
266
-
267
- This decorator is useful for handling transient errors, such as networking issues.
268
- If your task contains operations that can't be retried safely, e.g. database updates,
269
- it is advisable to annotate it with `@retry(times=0)`.
270
-
271
- This can be used in conjunction with the `@catch` decorator. The `@catch`
272
- decorator will execute a no-op task after all retries have been exhausted,
273
- ensuring that the flow execution can continue.
280
+ Specifies environment variables to be set prior to the execution of a step.
274
281
 
275
282
 
276
283
  Parameters
277
284
  ----------
278
- times : int, default 3
279
- Number of times to retry this task.
280
- minutes_between_retries : int, default 2
281
- Number of minutes between retries.
285
+ vars : Dict[str, str], default {}
286
+ Dictionary of environment variables to set.
282
287
  """
283
288
  ...
284
289
 
285
290
  @typing.overload
286
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
291
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
287
292
  ...
288
293
 
289
294
  @typing.overload
290
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
295
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
291
296
  ...
292
297
 
293
- 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):
298
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
294
299
  """
295
- Specifies the number of times the task corresponding
296
- to a step needs to be retried.
297
-
298
- This decorator is useful for handling transient errors, such as networking issues.
299
- If your task contains operations that can't be retried safely, e.g. database updates,
300
- it is advisable to annotate it with `@retry(times=0)`.
301
-
302
- This can be used in conjunction with the `@catch` decorator. The `@catch`
303
- decorator will execute a no-op task after all retries have been exhausted,
304
- ensuring that the flow execution can continue.
300
+ Specifies environment variables to be set prior to the execution of a step.
305
301
 
306
302
 
307
303
  Parameters
308
304
  ----------
309
- times : int, default 3
310
- Number of times to retry this task.
311
- minutes_between_retries : int, default 2
312
- Number of minutes between retries.
305
+ vars : Dict[str, str], default {}
306
+ Dictionary of environment variables to set.
313
307
  """
314
308
  ...
315
309
 
316
310
  @typing.overload
317
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
311
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> 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]]]:
318
312
  """
319
- Decorator prototype for all step decorators. This function gets specialized
320
- and imported for all decorators types by _import_plugin_decorators().
313
+ Specifies secrets to be retrieved and injected as environment variables prior to
314
+ the execution of a step.
315
+
316
+
317
+ Parameters
318
+ ----------
319
+ sources : List[Union[str, Dict[str, Any]]], default: []
320
+ List of secret specs, defining how the secrets are to be retrieved
321
321
  """
322
322
  ...
323
323
 
324
324
  @typing.overload
325
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
325
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
326
326
  ...
327
327
 
328
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
329
- """
330
- Decorator prototype for all step decorators. This function gets specialized
331
- and imported for all decorators types by _import_plugin_decorators().
332
- """
328
+ @typing.overload
329
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
333
330
  ...
334
331
 
335
- def vllm(*, model: str, backend: str, debug: bool, kwargs: typing.Any) -> 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]]]:
332
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
336
333
  """
337
- This decorator is used to run vllm APIs as Metaflow task sidecars.
338
-
339
- User code call
340
- --------------
341
- @vllm(
342
- model="...",
343
- ...
344
- )
345
-
346
- Valid backend options
347
- ---------------------
348
- - 'local': Run as a separate process on the local task machine.
349
-
350
- Valid model options
351
- -------------------
352
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
353
-
354
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
355
- If you need multiple models, you must create multiple @vllm decorators.
334
+ Specifies secrets to be retrieved and injected as environment variables prior to
335
+ the execution of a step.
356
336
 
357
337
 
358
338
  Parameters
359
339
  ----------
360
- model: str
361
- HuggingFace model identifier to be served by vLLM.
362
- backend: str
363
- Determines where and how to run the vLLM process.
364
- debug: bool
365
- Whether to turn on verbose debugging logs.
366
- kwargs : Any
367
- Any other keyword arguments are passed directly to the vLLM engine.
368
- This allows for flexible configuration of vLLM server settings.
369
- For example, `tensor_parallel_size=2`.
340
+ sources : List[Union[str, Dict[str, Any]]], default: []
341
+ List of secret specs, defining how the secrets are to be retrieved
370
342
  """
371
343
  ...
372
344
 
373
- def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
345
+ @typing.overload
346
+ 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]]]:
374
347
  """
375
- Specifies that this step should execute on DGX cloud.
348
+ Specifies the resources needed when executing this step.
376
349
 
350
+ Use `@resources` to specify the resource requirements
351
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
377
352
 
378
- Parameters
379
- ----------
380
- gpu : int
381
- Number of GPUs to use.
382
- gpu_type : str
383
- Type of Nvidia GPU to use.
384
- """
385
- ...
386
-
387
- 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]]]:
388
- """
389
- Specifies that this step should execute on DGX cloud.
353
+ You can choose the compute layer on the command line by executing e.g.
354
+ ```
355
+ python myflow.py run --with batch
356
+ ```
357
+ or
358
+ ```
359
+ python myflow.py run --with kubernetes
360
+ ```
361
+ which executes the flow on the desired system using the
362
+ requirements specified in `@resources`.
390
363
 
391
364
 
392
365
  Parameters
393
366
  ----------
394
- gpu : int
395
- Number of GPUs to use.
396
- gpu_type : str
397
- Type of Nvidia GPU to use.
398
- queue_timeout : int
399
- Time to keep the job in NVCF's queue.
367
+ cpu : int, default 1
368
+ Number of CPUs required for this step.
369
+ gpu : int, optional, default None
370
+ Number of GPUs required for this step.
371
+ disk : int, optional, default None
372
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
373
+ memory : int, default 4096
374
+ Memory size (in MB) required for this step.
375
+ shared_memory : int, optional, default None
376
+ The value for the size (in MiB) of the /dev/shm volume for this step.
377
+ This parameter maps to the `--shm-size` option in Docker.
400
378
  """
401
379
  ...
402
380
 
403
381
  @typing.overload
404
- 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]]]:
405
- """
406
- Specifies that the step will success under all circumstances.
407
-
408
- The decorator will create an optional artifact, specified by `var`, which
409
- contains the exception raised. You can use it to detect the presence
410
- of errors, indicating that all happy-path artifacts produced by the step
411
- are missing.
412
-
413
-
414
- Parameters
415
- ----------
416
- var : str, optional, default None
417
- Name of the artifact in which to store the caught exception.
418
- If not specified, the exception is not stored.
419
- print_exception : bool, default True
420
- Determines whether or not the exception is printed to
421
- stdout when caught.
422
- """
382
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
423
383
  ...
424
384
 
425
385
  @typing.overload
426
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
386
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
427
387
  ...
428
388
 
429
- @typing.overload
430
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
389
+ 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):
390
+ """
391
+ Specifies the resources needed when executing this step.
392
+
393
+ Use `@resources` to specify the resource requirements
394
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
395
+
396
+ You can choose the compute layer on the command line by executing e.g.
397
+ ```
398
+ python myflow.py run --with batch
399
+ ```
400
+ or
401
+ ```
402
+ python myflow.py run --with kubernetes
403
+ ```
404
+ which executes the flow on the desired system using the
405
+ requirements specified in `@resources`.
406
+
407
+
408
+ Parameters
409
+ ----------
410
+ cpu : int, default 1
411
+ Number of CPUs required for this step.
412
+ gpu : int, optional, default None
413
+ Number of GPUs required for this step.
414
+ disk : int, optional, default None
415
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
416
+ memory : int, default 4096
417
+ Memory size (in MB) required for this step.
418
+ shared_memory : int, optional, default None
419
+ The value for the size (in MiB) of the /dev/shm volume for this step.
420
+ This parameter maps to the `--shm-size` option in Docker.
421
+ """
431
422
  ...
432
423
 
433
- 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):
424
+ @typing.overload
425
+ 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]]]:
434
426
  """
435
427
  Specifies that the step will success under all circumstances.
436
428
 
@@ -452,110 +444,109 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
452
444
  ...
453
445
 
454
446
  @typing.overload
455
- 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]]]:
447
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
448
+ ...
449
+
450
+ @typing.overload
451
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
452
+ ...
453
+
454
+ 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):
456
455
  """
457
- Specifies the PyPI packages for the step.
456
+ Specifies that the step will success under all circumstances.
458
457
 
459
- Information in this decorator will augment any
460
- attributes set in the `@pyi_base` flow-level decorator. Hence,
461
- you can use `@pypi_base` to set packages required by all
462
- steps and use `@pypi` to specify step-specific overrides.
458
+ The decorator will create an optional artifact, specified by `var`, which
459
+ contains the exception raised. You can use it to detect the presence
460
+ of errors, indicating that all happy-path artifacts produced by the step
461
+ are missing.
463
462
 
464
463
 
465
464
  Parameters
466
465
  ----------
467
- packages : Dict[str, str], default: {}
468
- Packages to use for this step. The key is the name of the package
469
- and the value is the version to use.
470
- python : str, optional, default: None
471
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
472
- that the version used will correspond to the version of the Python interpreter used to start the run.
466
+ var : str, optional, default None
467
+ Name of the artifact in which to store the caught exception.
468
+ If not specified, the exception is not stored.
469
+ print_exception : bool, default True
470
+ Determines whether or not the exception is printed to
471
+ stdout when caught.
473
472
  """
474
473
  ...
475
474
 
476
475
  @typing.overload
477
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
476
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
477
+ """
478
+ Decorator prototype for all step decorators. This function gets specialized
479
+ and imported for all decorators types by _import_plugin_decorators().
480
+ """
478
481
  ...
479
482
 
480
483
  @typing.overload
481
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
484
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
482
485
  ...
483
486
 
484
- 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):
487
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
485
488
  """
486
- Specifies the PyPI packages for the step.
487
-
488
- Information in this decorator will augment any
489
- attributes set in the `@pyi_base` flow-level decorator. Hence,
490
- you can use `@pypi_base` to set packages required by all
491
- steps and use `@pypi` to specify step-specific overrides.
492
-
493
-
494
- Parameters
495
- ----------
496
- packages : Dict[str, str], default: {}
497
- Packages to use for this step. The key is the name of the package
498
- and the value is the version to use.
499
- python : str, optional, default: None
500
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
501
- that the version used will correspond to the version of the Python interpreter used to start the run.
489
+ Decorator prototype for all step decorators. This function gets specialized
490
+ and imported for all decorators types by _import_plugin_decorators().
502
491
  """
503
492
  ...
504
493
 
505
494
  @typing.overload
506
- 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]]]:
495
+ 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]]]:
507
496
  """
508
- Enables loading / saving of models within a step.
497
+ Specifies a timeout for your step.
509
498
 
499
+ This decorator is useful if this step may hang indefinitely.
510
500
 
501
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
502
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
503
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
511
504
 
512
- Parameters
513
- ----------
514
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
515
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
516
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
517
- - `current.checkpoint`
518
- - `current.model`
519
- - `current.huggingface_hub`
505
+ Note that all the values specified in parameters are added together so if you specify
506
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
520
507
 
521
- 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
522
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
523
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
524
508
 
525
- temp_dir_root : str, default: None
526
- The root directory under which `current.model.loaded` will store loaded models
509
+ Parameters
510
+ ----------
511
+ seconds : int, default 0
512
+ Number of seconds to wait prior to timing out.
513
+ minutes : int, default 0
514
+ Number of minutes to wait prior to timing out.
515
+ hours : int, default 0
516
+ Number of hours to wait prior to timing out.
527
517
  """
528
518
  ...
529
519
 
530
520
  @typing.overload
531
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
521
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
532
522
  ...
533
523
 
534
524
  @typing.overload
535
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
525
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
536
526
  ...
537
527
 
538
- 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):
528
+ 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):
539
529
  """
540
- Enables loading / saving of models within a step.
530
+ Specifies a timeout for your step.
541
531
 
532
+ This decorator is useful if this step may hang indefinitely.
542
533
 
534
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
535
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
536
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
543
537
 
544
- Parameters
545
- ----------
546
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
547
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
548
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
549
- - `current.checkpoint`
550
- - `current.model`
551
- - `current.huggingface_hub`
538
+ Note that all the values specified in parameters are added together so if you specify
539
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
552
540
 
553
- 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
554
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
555
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
556
541
 
557
- temp_dir_root : str, default: None
558
- The root directory under which `current.model.loaded` will store loaded models
542
+ Parameters
543
+ ----------
544
+ seconds : int, default 0
545
+ Number of seconds to wait prior to timing out.
546
+ minutes : int, default 0
547
+ Number of minutes to wait prior to timing out.
548
+ hours : int, default 0
549
+ Number of hours to wait prior to timing out.
559
550
  """
560
551
  ...
561
552
 
@@ -575,183 +566,116 @@ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union
575
566
  ...
576
567
 
577
568
  @typing.overload
578
- 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]]]:
569
+ def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
579
570
  """
580
- Specifies environment variables to be set prior to the execution of a step.
571
+ Creates a human-readable report, a Metaflow Card, after this step completes.
572
+
573
+ Note that you may add multiple `@card` decorators in a step with different parameters.
581
574
 
582
575
 
583
576
  Parameters
584
577
  ----------
585
- vars : Dict[str, str], default {}
586
- Dictionary of environment variables to set.
578
+ type : str, default 'default'
579
+ Card type.
580
+ id : str, optional, default None
581
+ If multiple cards are present, use this id to identify this card.
582
+ options : Dict[str, Any], default {}
583
+ Options passed to the card. The contents depend on the card type.
584
+ timeout : int, default 45
585
+ Interrupt reporting if it takes more than this many seconds.
587
586
  """
588
587
  ...
589
588
 
590
589
  @typing.overload
591
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
590
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
592
591
  ...
593
592
 
594
593
  @typing.overload
595
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
594
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
596
595
  ...
597
596
 
598
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
597
+ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, type: str = 'default', id: typing.Optional[str] = None, options: typing.Dict[str, typing.Any] = {}, timeout: int = 45):
599
598
  """
600
- Specifies environment variables to be set prior to the execution of a step.
599
+ Creates a human-readable report, a Metaflow Card, after this step completes.
600
+
601
+ Note that you may add multiple `@card` decorators in a step with different parameters.
601
602
 
602
603
 
603
604
  Parameters
604
605
  ----------
605
- vars : Dict[str, str], default {}
606
- Dictionary of environment variables to set.
606
+ type : str, default 'default'
607
+ Card type.
608
+ id : str, optional, default None
609
+ If multiple cards are present, use this id to identify this card.
610
+ options : Dict[str, Any], default {}
611
+ Options passed to the card. The contents depend on the card type.
612
+ timeout : int, default 45
613
+ Interrupt reporting if it takes more than this many seconds.
607
614
  """
608
615
  ...
609
616
 
610
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
617
+ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
611
618
  """
612
- Specifies that this step should execute on Kubernetes.
619
+ Specifies that this step should execute on DGX cloud.
613
620
 
614
621
 
615
622
  Parameters
616
623
  ----------
617
- cpu : int, default 1
618
- Number of CPUs required for this step. If `@resources` is
619
- also present, the maximum value from all decorators is used.
620
- memory : int, default 4096
621
- Memory size (in MB) required for this step. If
622
- `@resources` is also present, the maximum value from all decorators is
623
- used.
624
- disk : int, default 10240
625
- Disk size (in MB) required for this step. If
626
- `@resources` is also present, the maximum value from all decorators is
627
- used.
628
- image : str, optional, default None
629
- Docker image to use when launching on Kubernetes. If not specified, and
630
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
631
- not, a default Docker image mapping to the current version of Python is used.
632
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
633
- If given, the imagePullPolicy to be applied to the Docker image of the step.
634
- image_pull_secrets: List[str], default []
635
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
636
- Kubernetes image pull secrets to use when pulling container images
637
- in Kubernetes.
638
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
639
- Kubernetes service account to use when launching pod in Kubernetes.
640
- secrets : List[str], optional, default None
641
- Kubernetes secrets to use when launching pod in Kubernetes. These
642
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
643
- in Metaflow configuration.
644
- node_selector: Union[Dict[str,str], str], optional, default None
645
- Kubernetes node selector(s) to apply to the pod running the task.
646
- Can be passed in as a comma separated string of values e.g.
647
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
648
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
649
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
650
- Kubernetes namespace to use when launching pod in Kubernetes.
651
- gpu : int, optional, default None
652
- Number of GPUs required for this step. A value of zero implies that
653
- the scheduled node should not have GPUs.
654
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
655
- The vendor of the GPUs to be used for this step.
656
- tolerations : List[str], default []
657
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
658
- Kubernetes tolerations to use when launching pod in Kubernetes.
659
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
660
- Kubernetes labels to use when launching pod in Kubernetes.
661
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
662
- Kubernetes annotations to use when launching pod in Kubernetes.
663
- use_tmpfs : bool, default False
664
- This enables an explicit tmpfs mount for this step.
665
- tmpfs_tempdir : bool, default True
666
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
667
- tmpfs_size : int, optional, default: None
668
- The value for the size (in MiB) of the tmpfs mount for this step.
669
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
670
- memory allocated for this step.
671
- tmpfs_path : str, optional, default /metaflow_temp
672
- Path to tmpfs mount for this step.
673
- persistent_volume_claims : Dict[str, str], optional, default None
674
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
675
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
676
- shared_memory: int, optional
677
- Shared memory size (in MiB) required for this step
678
- port: int, optional
679
- Port number to specify in the Kubernetes job object
680
- compute_pool : str, optional, default None
681
- Compute pool to be used for for this step.
682
- If not specified, any accessible compute pool within the perimeter is used.
683
- hostname_resolution_timeout: int, default 10 * 60
684
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
685
- Only applicable when @parallel is used.
686
- qos: str, default: Burstable
687
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
688
-
689
- security_context: Dict[str, Any], optional, default None
690
- Container security context. Applies to the task container. Allows the following keys:
691
- - privileged: bool, optional, default None
692
- - allow_privilege_escalation: bool, optional, default None
693
- - run_as_user: int, optional, default None
694
- - run_as_group: int, optional, default None
695
- - run_as_non_root: bool, optional, default None
624
+ gpu : int
625
+ Number of GPUs to use.
626
+ gpu_type : str
627
+ Type of Nvidia GPU to use.
696
628
  """
697
629
  ...
698
630
 
699
631
  @typing.overload
700
- 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]]]:
632
+ 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]]]:
701
633
  """
702
- Specifies a timeout for your step.
703
-
704
- This decorator is useful if this step may hang indefinitely.
705
-
706
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
707
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
708
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
634
+ Specifies the PyPI packages for the step.
709
635
 
710
- Note that all the values specified in parameters are added together so if you specify
711
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
636
+ Information in this decorator will augment any
637
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
638
+ you can use `@pypi_base` to set packages required by all
639
+ steps and use `@pypi` to specify step-specific overrides.
712
640
 
713
641
 
714
642
  Parameters
715
643
  ----------
716
- seconds : int, default 0
717
- Number of seconds to wait prior to timing out.
718
- minutes : int, default 0
719
- Number of minutes to wait prior to timing out.
720
- hours : int, default 0
721
- Number of hours to wait prior to timing out.
644
+ packages : Dict[str, str], default: {}
645
+ Packages to use for this step. The key is the name of the package
646
+ and the value is the version to use.
647
+ python : str, optional, default: None
648
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
649
+ that the version used will correspond to the version of the Python interpreter used to start the run.
722
650
  """
723
651
  ...
724
652
 
725
653
  @typing.overload
726
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
654
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
727
655
  ...
728
656
 
729
657
  @typing.overload
730
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
658
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
731
659
  ...
732
660
 
733
- 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):
661
+ 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):
734
662
  """
735
- Specifies a timeout for your step.
736
-
737
- This decorator is useful if this step may hang indefinitely.
738
-
739
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
740
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
741
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
663
+ Specifies the PyPI packages for the step.
742
664
 
743
- Note that all the values specified in parameters are added together so if you specify
744
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
665
+ Information in this decorator will augment any
666
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
667
+ you can use `@pypi_base` to set packages required by all
668
+ steps and use `@pypi` to specify step-specific overrides.
745
669
 
746
670
 
747
671
  Parameters
748
672
  ----------
749
- seconds : int, default 0
750
- Number of seconds to wait prior to timing out.
751
- minutes : int, default 0
752
- Number of minutes to wait prior to timing out.
753
- hours : int, default 0
754
- Number of hours to wait prior to timing out.
673
+ packages : Dict[str, str], default: {}
674
+ Packages to use for this step. The key is the name of the package
675
+ and the value is the version to use.
676
+ python : str, optional, default: None
677
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
678
+ that the version used will correspond to the version of the Python interpreter used to start the run.
755
679
  """
756
680
  ...
757
681
 
@@ -812,414 +736,345 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
812
736
  """
813
737
  ...
814
738
 
815
- @typing.overload
816
- 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]]]:
739
+ def vllm(*, model: str, backend: str, debug: bool, kwargs: typing.Any) -> 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]]]:
817
740
  """
818
- Creates a human-readable report, a Metaflow Card, after this step completes.
741
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
819
742
 
820
- Note that you may add multiple `@card` decorators in a step with different parameters.
743
+ User code call
744
+ --------------
745
+ @vllm(
746
+ model="...",
747
+ ...
748
+ )
749
+
750
+ Valid backend options
751
+ ---------------------
752
+ - 'local': Run as a separate process on the local task machine.
753
+
754
+ Valid model options
755
+ -------------------
756
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
757
+
758
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
759
+ If you need multiple models, you must create multiple @vllm decorators.
821
760
 
822
761
 
823
762
  Parameters
824
763
  ----------
825
- type : str, default 'default'
826
- Card type.
827
- id : str, optional, default None
828
- If multiple cards are present, use this id to identify this card.
829
- options : Dict[str, Any], default {}
830
- Options passed to the card. The contents depend on the card type.
831
- timeout : int, default 45
832
- Interrupt reporting if it takes more than this many seconds.
764
+ model: str
765
+ HuggingFace model identifier to be served by vLLM.
766
+ backend: str
767
+ Determines where and how to run the vLLM process.
768
+ debug: bool
769
+ Whether to turn on verbose debugging logs.
770
+ kwargs : Any
771
+ Any other keyword arguments are passed directly to the vLLM engine.
772
+ This allows for flexible configuration of vLLM server settings.
773
+ For example, `tensor_parallel_size=2`.
833
774
  """
834
775
  ...
835
776
 
836
- @typing.overload
837
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
838
- ...
839
-
840
- @typing.overload
841
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
842
- ...
843
-
844
- 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):
777
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[str] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
845
778
  """
846
- Creates a human-readable report, a Metaflow Card, after this step completes.
847
-
848
- Note that you may add multiple `@card` decorators in a step with different parameters.
779
+ Specifies that this step should execute on Kubernetes.
849
780
 
850
781
 
851
782
  Parameters
852
783
  ----------
853
- type : str, default 'default'
854
- Card type.
855
- id : str, optional, default None
856
- If multiple cards are present, use this id to identify this card.
857
- options : Dict[str, Any], default {}
858
- Options passed to the card. The contents depend on the card type.
859
- timeout : int, default 45
860
- Interrupt reporting if it takes more than this many seconds.
784
+ cpu : int, default 1
785
+ Number of CPUs required for this step. If `@resources` is
786
+ also present, the maximum value from all decorators is used.
787
+ memory : int, default 4096
788
+ Memory size (in MB) required for this step. If
789
+ `@resources` is also present, the maximum value from all decorators is
790
+ used.
791
+ disk : int, default 10240
792
+ Disk size (in MB) required for this step. If
793
+ `@resources` is also present, the maximum value from all decorators is
794
+ used.
795
+ image : str, optional, default None
796
+ Docker image to use when launching on Kubernetes. If not specified, and
797
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
798
+ not, a default Docker image mapping to the current version of Python is used.
799
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
800
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
801
+ image_pull_secrets: List[str], default []
802
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
803
+ Kubernetes image pull secrets to use when pulling container images
804
+ in Kubernetes.
805
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
806
+ Kubernetes service account to use when launching pod in Kubernetes.
807
+ secrets : List[str], optional, default None
808
+ Kubernetes secrets to use when launching pod in Kubernetes. These
809
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
810
+ in Metaflow configuration.
811
+ node_selector: Union[Dict[str,str], str], optional, default None
812
+ Kubernetes node selector(s) to apply to the pod running the task.
813
+ Can be passed in as a comma separated string of values e.g.
814
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
815
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
816
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
817
+ Kubernetes namespace to use when launching pod in Kubernetes.
818
+ gpu : int, optional, default None
819
+ Number of GPUs required for this step. A value of zero implies that
820
+ the scheduled node should not have GPUs.
821
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
822
+ The vendor of the GPUs to be used for this step.
823
+ tolerations : List[str], default []
824
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
825
+ Kubernetes tolerations to use when launching pod in Kubernetes.
826
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
827
+ Kubernetes labels to use when launching pod in Kubernetes.
828
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
829
+ Kubernetes annotations to use when launching pod in Kubernetes.
830
+ use_tmpfs : bool, default False
831
+ This enables an explicit tmpfs mount for this step.
832
+ tmpfs_tempdir : bool, default True
833
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
834
+ tmpfs_size : int, optional, default: None
835
+ The value for the size (in MiB) of the tmpfs mount for this step.
836
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
837
+ memory allocated for this step.
838
+ tmpfs_path : str, optional, default /metaflow_temp
839
+ Path to tmpfs mount for this step.
840
+ persistent_volume_claims : Dict[str, str], optional, default None
841
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
842
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
843
+ shared_memory: int, optional
844
+ Shared memory size (in MiB) required for this step
845
+ port: int, optional
846
+ Port number to specify in the Kubernetes job object
847
+ compute_pool : str, optional, default None
848
+ Compute pool to be used for for this step.
849
+ If not specified, any accessible compute pool within the perimeter is used.
850
+ hostname_resolution_timeout: int, default 10 * 60
851
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
852
+ Only applicable when @parallel is used.
853
+ qos: str, default: Burstable
854
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
855
+
856
+ security_context: Dict[str, Any], optional, default None
857
+ Container security context. Applies to the task container. Allows the following keys:
858
+ - privileged: bool, optional, default None
859
+ - allow_privilege_escalation: bool, optional, default None
860
+ - run_as_user: int, optional, default None
861
+ - run_as_group: int, optional, default None
862
+ - run_as_non_root: bool, optional, default None
861
863
  """
862
864
  ...
863
865
 
864
- @typing.overload
865
- 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]]]:
866
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
866
867
  """
867
- Specifies the resources needed when executing this step.
868
-
869
- Use `@resources` to specify the resource requirements
870
- independently of the specific compute layer (`@batch`, `@kubernetes`).
871
-
872
- You can choose the compute layer on the command line by executing e.g.
873
- ```
874
- python myflow.py run --with batch
875
- ```
876
- or
877
- ```
878
- python myflow.py run --with kubernetes
879
- ```
880
- which executes the flow on the desired system using the
881
- requirements specified in `@resources`.
868
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
882
869
 
883
870
 
884
871
  Parameters
885
872
  ----------
886
- cpu : int, default 1
887
- Number of CPUs required for this step.
888
- gpu : int, optional, default None
889
- Number of GPUs required for this step.
890
- disk : int, optional, default None
891
- Disk size (in MB) required for this step. Only applies on Kubernetes.
892
- memory : int, default 4096
893
- Memory size (in MB) required for this step.
894
- shared_memory : int, optional, default None
895
- The value for the size (in MiB) of the /dev/shm volume for this step.
896
- This parameter maps to the `--shm-size` option in Docker.
873
+ temp_dir_root : str, optional
874
+ The root directory that will hold the temporary directory where objects will be downloaded.
875
+
876
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
877
+ The list of repos (models/datasets) to load.
878
+
879
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
880
+
881
+ - If repo (model/dataset) is not found in the datastore:
882
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
883
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
884
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
885
+
886
+ - If repo is found in the datastore:
887
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
897
888
  """
898
889
  ...
899
890
 
900
891
  @typing.overload
901
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
902
- ...
903
-
904
- @typing.overload
905
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
906
- ...
907
-
908
- 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):
892
+ 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]]]:
909
893
  """
910
- Specifies the resources needed when executing this step.
911
-
912
- Use `@resources` to specify the resource requirements
913
- independently of the specific compute layer (`@batch`, `@kubernetes`).
894
+ Enables loading / saving of models within a step.
914
895
 
915
- You can choose the compute layer on the command line by executing e.g.
916
- ```
917
- python myflow.py run --with batch
918
- ```
919
- or
920
- ```
921
- python myflow.py run --with kubernetes
922
- ```
923
- which executes the flow on the desired system using the
924
- requirements specified in `@resources`.
925
896
 
926
897
 
927
898
  Parameters
928
899
  ----------
929
- cpu : int, default 1
930
- Number of CPUs required for this step.
931
- gpu : int, optional, default None
932
- Number of GPUs required for this step.
933
- disk : int, optional, default None
934
- Disk size (in MB) required for this step. Only applies on Kubernetes.
935
- memory : int, default 4096
936
- Memory size (in MB) required for this step.
937
- shared_memory : int, optional, default None
938
- The value for the size (in MiB) of the /dev/shm volume for this step.
939
- This parameter maps to the `--shm-size` option in Docker.
900
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
901
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
902
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
903
+ - `current.checkpoint`
904
+ - `current.model`
905
+ - `current.huggingface_hub`
906
+
907
+ 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
908
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
909
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
910
+
911
+ temp_dir_root : str, default: None
912
+ The root directory under which `current.model.loaded` will store loaded models
940
913
  """
941
914
  ...
942
915
 
943
916
  @typing.overload
944
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
945
- """
946
- Internal decorator to support Fast bakery
947
- """
917
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
948
918
  ...
949
919
 
950
920
  @typing.overload
951
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
952
- ...
953
-
954
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
955
- """
956
- Internal decorator to support Fast bakery
957
- """
921
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
958
922
  ...
959
923
 
960
- @typing.overload
961
- 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]]]:
924
+ 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):
962
925
  """
963
- Specifies the Conda environment for the step.
926
+ Enables loading / saving of models within a step.
964
927
 
965
- Information in this decorator will augment any
966
- attributes set in the `@conda_base` flow-level decorator. Hence,
967
- you can use `@conda_base` to set packages required by all
968
- steps and use `@conda` to specify step-specific overrides.
969
928
 
970
929
 
971
930
  Parameters
972
931
  ----------
973
- packages : Dict[str, str], default {}
974
- Packages to use for this step. The key is the name of the package
975
- and the value is the version to use.
976
- libraries : Dict[str, str], default {}
977
- Supported for backward compatibility. When used with packages, packages will take precedence.
978
- python : str, optional, default None
979
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
980
- that the version used will correspond to the version of the Python interpreter used to start the run.
981
- disabled : bool, default False
982
- If set to True, disables @conda.
932
+ load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
933
+ Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
934
+ These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
935
+ - `current.checkpoint`
936
+ - `current.model`
937
+ - `current.huggingface_hub`
938
+
939
+ 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
940
+ the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
941
+ If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
942
+
943
+ temp_dir_root : str, default: None
944
+ The root directory under which `current.model.loaded` will store loaded models
983
945
  """
984
946
  ...
985
947
 
986
- @typing.overload
987
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
988
- ...
989
-
990
- @typing.overload
991
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
992
- ...
993
-
994
- 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):
948
+ 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]]]:
995
949
  """
996
- Specifies the Conda environment for the step.
997
-
998
- Information in this decorator will augment any
999
- attributes set in the `@conda_base` flow-level decorator. Hence,
1000
- you can use `@conda_base` to set packages required by all
1001
- steps and use `@conda` to specify step-specific overrides.
950
+ Specifies that this step should execute on DGX cloud.
1002
951
 
1003
952
 
1004
953
  Parameters
1005
954
  ----------
1006
- packages : Dict[str, str], default {}
1007
- Packages to use for this step. The key is the name of the package
1008
- and the value is the version to use.
1009
- libraries : Dict[str, str], default {}
1010
- Supported for backward compatibility. When used with packages, packages will take precedence.
1011
- python : str, optional, default None
1012
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1013
- that the version used will correspond to the version of the Python interpreter used to start the run.
1014
- disabled : bool, default False
1015
- If set to True, disables @conda.
955
+ gpu : int
956
+ Number of GPUs to use.
957
+ gpu_type : str
958
+ Type of Nvidia GPU to use.
959
+ queue_timeout : int
960
+ Time to keep the job in NVCF's queue.
1016
961
  """
1017
962
  ...
1018
963
 
1019
964
  @typing.overload
1020
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
965
+ 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]]]:
1021
966
  """
1022
- Specifies the PyPI packages for all steps of the flow.
967
+ Specifies the number of times the task corresponding
968
+ to a step needs to be retried.
969
+
970
+ This decorator is useful for handling transient errors, such as networking issues.
971
+ If your task contains operations that can't be retried safely, e.g. database updates,
972
+ it is advisable to annotate it with `@retry(times=0)`.
973
+
974
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
975
+ decorator will execute a no-op task after all retries have been exhausted,
976
+ ensuring that the flow execution can continue.
1023
977
 
1024
- Use `@pypi_base` to set common packages required by all
1025
- steps and use `@pypi` to specify step-specific overrides.
1026
978
 
1027
979
  Parameters
1028
980
  ----------
1029
- packages : Dict[str, str], default: {}
1030
- Packages to use for this flow. The key is the name of the package
1031
- and the value is the version to use.
1032
- python : str, optional, default: None
1033
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1034
- that the version used will correspond to the version of the Python interpreter used to start the run.
981
+ times : int, default 3
982
+ Number of times to retry this task.
983
+ minutes_between_retries : int, default 2
984
+ Number of minutes between retries.
1035
985
  """
1036
986
  ...
1037
987
 
1038
988
  @typing.overload
1039
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
989
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
1040
990
  ...
1041
991
 
1042
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
992
+ @typing.overload
993
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
994
+ ...
995
+
996
+ 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):
1043
997
  """
1044
- Specifies the PyPI packages for all steps of the flow.
998
+ Specifies the number of times the task corresponding
999
+ to a step needs to be retried.
1000
+
1001
+ This decorator is useful for handling transient errors, such as networking issues.
1002
+ If your task contains operations that can't be retried safely, e.g. database updates,
1003
+ it is advisable to annotate it with `@retry(times=0)`.
1004
+
1005
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
1006
+ decorator will execute a no-op task after all retries have been exhausted,
1007
+ ensuring that the flow execution can continue.
1045
1008
 
1046
- Use `@pypi_base` to set common packages required by all
1047
- steps and use `@pypi` to specify step-specific overrides.
1048
1009
 
1049
1010
  Parameters
1050
1011
  ----------
1051
- packages : Dict[str, str], default: {}
1052
- Packages to use for this flow. The key is the name of the package
1053
- and the value is the version to use.
1054
- python : str, optional, default: None
1055
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1056
- that the version used will correspond to the version of the Python interpreter used to start the run.
1012
+ times : int, default 3
1013
+ Number of times to retry this task.
1014
+ minutes_between_retries : int, default 2
1015
+ Number of minutes between retries.
1057
1016
  """
1058
1017
  ...
1059
1018
 
1060
1019
  @typing.overload
1061
- 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]]:
1020
+ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1062
1021
  """
1063
- Specifies the flow(s) that this flow depends on.
1022
+ Specifies the event(s) that this flow depends on.
1064
1023
 
1065
1024
  ```
1066
- @trigger_on_finish(flow='FooFlow')
1025
+ @trigger(event='foo')
1067
1026
  ```
1068
1027
  or
1069
1028
  ```
1070
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1029
+ @trigger(events=['foo', 'bar'])
1071
1030
  ```
1072
- This decorator respects the @project decorator and triggers the flow
1073
- when upstream runs within the same namespace complete successfully
1074
1031
 
1075
- Additionally, you can specify project aware upstream flow dependencies
1076
- by specifying the fully qualified project_flow_name.
1032
+ Additionally, you can specify the parameter mappings
1033
+ to map event payload to Metaflow parameters for the flow.
1077
1034
  ```
1078
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1035
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1079
1036
  ```
1080
1037
  or
1081
1038
  ```
1082
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1039
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1040
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1083
1041
  ```
1084
1042
 
1085
- You can also specify just the project or project branch (other values will be
1086
- inferred from the current project or project branch):
1043
+ 'parameters' can also be a list of strings and tuples like so:
1087
1044
  ```
1088
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1045
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1046
+ ```
1047
+ This is equivalent to:
1048
+ ```
1049
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1089
1050
  ```
1090
-
1091
- Note that `branch` is typically one of:
1092
- - `prod`
1093
- - `user.bob`
1094
- - `test.my_experiment`
1095
- - `prod.staging`
1096
1051
 
1097
1052
 
1098
1053
  Parameters
1099
1054
  ----------
1100
- flow : Union[str, Dict[str, str]], optional, default None
1101
- Upstream flow dependency for this flow.
1102
- flows : List[Union[str, Dict[str, str]]], default []
1103
- Upstream flow dependencies for this flow.
1055
+ event : Union[str, Dict[str, Any]], optional, default None
1056
+ Event dependency for this flow.
1057
+ events : List[Union[str, Dict[str, Any]]], default []
1058
+ Events dependency for this flow.
1104
1059
  options : Dict[str, Any], default {}
1105
1060
  Backend-specific configuration for tuning eventing behavior.
1106
1061
  """
1107
1062
  ...
1108
1063
 
1109
1064
  @typing.overload
1110
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1065
+ def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1111
1066
  ...
1112
1067
 
1113
- 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] = {}):
1068
+ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1114
1069
  """
1115
- Specifies the flow(s) that this flow depends on.
1070
+ Specifies the event(s) that this flow depends on.
1116
1071
 
1117
1072
  ```
1118
- @trigger_on_finish(flow='FooFlow')
1073
+ @trigger(event='foo')
1119
1074
  ```
1120
1075
  or
1121
1076
  ```
1122
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1123
- ```
1124
- This decorator respects the @project decorator and triggers the flow
1125
- when upstream runs within the same namespace complete successfully
1126
-
1127
- Additionally, you can specify project aware upstream flow dependencies
1128
- by specifying the fully qualified project_flow_name.
1129
- ```
1130
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1131
- ```
1132
- or
1133
- ```
1134
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1135
- ```
1136
-
1137
- You can also specify just the project or project branch (other values will be
1138
- inferred from the current project or project branch):
1139
- ```
1140
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1141
- ```
1142
-
1143
- Note that `branch` is typically one of:
1144
- - `prod`
1145
- - `user.bob`
1146
- - `test.my_experiment`
1147
- - `prod.staging`
1148
-
1149
-
1150
- Parameters
1151
- ----------
1152
- flow : Union[str, Dict[str, str]], optional, default None
1153
- Upstream flow dependency for this flow.
1154
- flows : List[Union[str, Dict[str, str]]], default []
1155
- Upstream flow dependencies for this flow.
1156
- options : Dict[str, Any], default {}
1157
- Backend-specific configuration for tuning eventing behavior.
1158
- """
1159
- ...
1160
-
1161
- @typing.overload
1162
- 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]]:
1163
- """
1164
- Specifies the Conda environment for all steps of the flow.
1165
-
1166
- Use `@conda_base` to set common libraries required by all
1167
- steps and use `@conda` to specify step-specific additions.
1168
-
1169
-
1170
- Parameters
1171
- ----------
1172
- packages : Dict[str, str], default {}
1173
- Packages to use for this flow. The key is the name of the package
1174
- and the value is the version to use.
1175
- libraries : Dict[str, str], default {}
1176
- Supported for backward compatibility. When used with packages, packages will take precedence.
1177
- python : str, optional, default None
1178
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1179
- that the version used will correspond to the version of the Python interpreter used to start the run.
1180
- disabled : bool, default False
1181
- If set to True, disables Conda.
1182
- """
1183
- ...
1184
-
1185
- @typing.overload
1186
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1187
- ...
1188
-
1189
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1190
- """
1191
- Specifies the Conda environment for all steps of the flow.
1192
-
1193
- Use `@conda_base` to set common libraries required by all
1194
- steps and use `@conda` to specify step-specific additions.
1195
-
1196
-
1197
- Parameters
1198
- ----------
1199
- packages : Dict[str, str], default {}
1200
- Packages to use for this flow. The key is the name of the package
1201
- and the value is the version to use.
1202
- libraries : Dict[str, str], default {}
1203
- Supported for backward compatibility. When used with packages, packages will take precedence.
1204
- python : str, optional, default None
1205
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1206
- that the version used will correspond to the version of the Python interpreter used to start the run.
1207
- disabled : bool, default False
1208
- If set to True, disables Conda.
1209
- """
1210
- ...
1211
-
1212
- @typing.overload
1213
- def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1214
- """
1215
- Specifies the event(s) that this flow depends on.
1216
-
1217
- ```
1218
- @trigger(event='foo')
1219
- ```
1220
- or
1221
- ```
1222
- @trigger(events=['foo', 'bar'])
1077
+ @trigger(events=['foo', 'bar'])
1223
1078
  ```
1224
1079
 
1225
1080
  Additionally, you can specify the parameter mappings
@@ -1254,58 +1109,127 @@ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = No
1254
1109
  """
1255
1110
  ...
1256
1111
 
1257
- @typing.overload
1258
- def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1259
- ...
1260
-
1261
- def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}):
1112
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1262
1113
  """
1263
- Specifies the event(s) that this flow depends on.
1114
+ Allows setting external datastores to save data for the
1115
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1264
1116
 
1265
- ```
1266
- @trigger(event='foo')
1267
- ```
1268
- or
1269
- ```
1270
- @trigger(events=['foo', 'bar'])
1271
- ```
1117
+ This decorator is useful when users wish to save data to a different datastore
1118
+ than what is configured in Metaflow. This can be for variety of reasons:
1272
1119
 
1273
- Additionally, you can specify the parameter mappings
1274
- to map event payload to Metaflow parameters for the flow.
1275
- ```
1276
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1277
- ```
1278
- or
1279
- ```
1280
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1281
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1282
- ```
1120
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1121
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1122
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1123
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1124
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1283
1125
 
1284
- 'parameters' can also be a list of strings and tuples like so:
1285
- ```
1286
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1287
- ```
1288
- This is equivalent to:
1289
- ```
1290
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1291
- ```
1126
+ Usage:
1127
+ ----------
1292
1128
 
1129
+ - Using a custom IAM role to access the datastore.
1293
1130
 
1294
- Parameters
1131
+ ```python
1132
+ @with_artifact_store(
1133
+ type="s3",
1134
+ config=lambda: {
1135
+ "root": "s3://my-bucket-foo/path/to/root",
1136
+ "role_arn": ROLE,
1137
+ },
1138
+ )
1139
+ class MyFlow(FlowSpec):
1140
+
1141
+ @checkpoint
1142
+ @step
1143
+ def start(self):
1144
+ with open("my_file.txt", "w") as f:
1145
+ f.write("Hello, World!")
1146
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1147
+ self.next(self.end)
1148
+
1149
+ ```
1150
+
1151
+ - Using credentials to access the s3-compatible datastore.
1152
+
1153
+ ```python
1154
+ @with_artifact_store(
1155
+ type="s3",
1156
+ config=lambda: {
1157
+ "root": "s3://my-bucket-foo/path/to/root",
1158
+ "client_params": {
1159
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1160
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1161
+ },
1162
+ },
1163
+ )
1164
+ class MyFlow(FlowSpec):
1165
+
1166
+ @checkpoint
1167
+ @step
1168
+ def start(self):
1169
+ with open("my_file.txt", "w") as f:
1170
+ f.write("Hello, World!")
1171
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1172
+ self.next(self.end)
1173
+
1174
+ ```
1175
+
1176
+ - Accessing objects stored in external datastores after task execution.
1177
+
1178
+ ```python
1179
+ run = Run("CheckpointsTestsFlow/8992")
1180
+ with artifact_store_from(run=run, config={
1181
+ "client_params": {
1182
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1183
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1184
+ },
1185
+ }):
1186
+ with Checkpoint() as cp:
1187
+ latest = cp.list(
1188
+ task=run["start"].task
1189
+ )[0]
1190
+ print(latest)
1191
+ cp.load(
1192
+ latest,
1193
+ "test-checkpoints"
1194
+ )
1195
+
1196
+ task = Task("TorchTuneFlow/8484/train/53673")
1197
+ with artifact_store_from(run=run, config={
1198
+ "client_params": {
1199
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1200
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1201
+ },
1202
+ }):
1203
+ load_model(
1204
+ task.data.model_ref,
1205
+ "test-models"
1206
+ )
1207
+ ```
1208
+ Parameters:
1295
1209
  ----------
1296
- event : Union[str, Dict[str, Any]], optional, default None
1297
- Event dependency for this flow.
1298
- events : List[Union[str, Dict[str, Any]]], default []
1299
- Events dependency for this flow.
1300
- options : Dict[str, Any], default {}
1301
- Backend-specific configuration for tuning eventing behavior.
1210
+
1211
+ type: str
1212
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1213
+
1214
+ config: dict or Callable
1215
+ Dictionary of configuration options for the datastore. The following keys are required:
1216
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1217
+ - example: 's3://bucket-name/path/to/root'
1218
+ - example: 'gs://bucket-name/path/to/root'
1219
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1220
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1221
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1222
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1302
1223
  """
1303
1224
  ...
1304
1225
 
1305
- 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]]:
1226
+ 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]]:
1306
1227
  """
1307
- 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.
1308
- 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.
1228
+ 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)
1229
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1230
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1231
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1232
+ starts only after all sensors finish.
1309
1233
 
1310
1234
 
1311
1235
  Parameters
@@ -1327,21 +1251,18 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1327
1251
  Name of the sensor on Airflow
1328
1252
  description : str
1329
1253
  Description of sensor in the Airflow UI
1330
- external_dag_id : str
1331
- The dag_id that contains the task you want to wait for.
1332
- external_task_ids : List[str]
1333
- The list of task_ids that you want to wait for.
1334
- If None (default value) the sensor waits for the DAG. (Default: None)
1335
- allowed_states : List[str]
1336
- Iterable of allowed states, (Default: ['success'])
1337
- failed_states : List[str]
1338
- Iterable of failed or dis-allowed states. (Default: None)
1339
- execution_delta : datetime.timedelta
1340
- time difference with the previous execution to look at,
1341
- the default is the same logical date as the current task or DAG. (Default: None)
1342
- check_existence: bool
1343
- Set to True to check if the external task exists or check if
1344
- the DAG to wait for exists. (Default: True)
1254
+ bucket_key : Union[str, List[str]]
1255
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1256
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1257
+ bucket_name : str
1258
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1259
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1260
+ wildcard_match : bool
1261
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1262
+ aws_conn_id : str
1263
+ a reference to the s3 connection on Airflow. (Default: None)
1264
+ verify : bool
1265
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1345
1266
  """
1346
1267
  ...
1347
1268
 
@@ -1380,6 +1301,100 @@ def project(*, name: str, branch: typing.Optional[str] = None, production: bool
1380
1301
  """
1381
1302
  ...
1382
1303
 
1304
+ @typing.overload
1305
+ 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]]:
1306
+ """
1307
+ Specifies the Conda environment for all steps of the flow.
1308
+
1309
+ Use `@conda_base` to set common libraries required by all
1310
+ steps and use `@conda` to specify step-specific additions.
1311
+
1312
+
1313
+ Parameters
1314
+ ----------
1315
+ packages : Dict[str, str], default {}
1316
+ Packages to use for this flow. The key is the name of the package
1317
+ and the value is the version to use.
1318
+ libraries : Dict[str, str], default {}
1319
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1320
+ python : str, optional, default None
1321
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1322
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1323
+ disabled : bool, default False
1324
+ If set to True, disables Conda.
1325
+ """
1326
+ ...
1327
+
1328
+ @typing.overload
1329
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1330
+ ...
1331
+
1332
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1333
+ """
1334
+ Specifies the Conda environment for all steps of the flow.
1335
+
1336
+ Use `@conda_base` to set common libraries required by all
1337
+ steps and use `@conda` to specify step-specific additions.
1338
+
1339
+
1340
+ Parameters
1341
+ ----------
1342
+ packages : Dict[str, str], default {}
1343
+ Packages to use for this flow. The key is the name of the package
1344
+ and the value is the version to use.
1345
+ libraries : Dict[str, str], default {}
1346
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1347
+ python : str, optional, default None
1348
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1349
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1350
+ disabled : bool, default False
1351
+ If set to True, disables Conda.
1352
+ """
1353
+ ...
1354
+
1355
+ 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]]:
1356
+ """
1357
+ 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.
1358
+ 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.
1359
+
1360
+
1361
+ Parameters
1362
+ ----------
1363
+ timeout : int
1364
+ Time, in seconds before the task times out and fails. (Default: 3600)
1365
+ poke_interval : int
1366
+ Time in seconds that the job should wait in between each try. (Default: 60)
1367
+ mode : str
1368
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1369
+ exponential_backoff : bool
1370
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1371
+ pool : str
1372
+ the slot pool this task should run in,
1373
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1374
+ soft_fail : bool
1375
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1376
+ name : str
1377
+ Name of the sensor on Airflow
1378
+ description : str
1379
+ Description of sensor in the Airflow UI
1380
+ external_dag_id : str
1381
+ The dag_id that contains the task you want to wait for.
1382
+ external_task_ids : List[str]
1383
+ The list of task_ids that you want to wait for.
1384
+ If None (default value) the sensor waits for the DAG. (Default: None)
1385
+ allowed_states : List[str]
1386
+ Iterable of allowed states, (Default: ['success'])
1387
+ failed_states : List[str]
1388
+ Iterable of failed or dis-allowed states. (Default: None)
1389
+ execution_delta : datetime.timedelta
1390
+ time difference with the previous execution to look at,
1391
+ the default is the same logical date as the current task or DAG. (Default: None)
1392
+ check_existence: bool
1393
+ Set to True to check if the external task exists or check if
1394
+ the DAG to wait for exists. (Default: True)
1395
+ """
1396
+ ...
1397
+
1383
1398
  @typing.overload
1384
1399
  def schedule(*, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1385
1400
  """
@@ -1431,160 +1446,145 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1431
1446
  """
1432
1447
  ...
1433
1448
 
1434
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1449
+ @typing.overload
1450
+ 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]]:
1435
1451
  """
1436
- Allows setting external datastores to save data for the
1437
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1438
-
1439
- This decorator is useful when users wish to save data to a different datastore
1440
- than what is configured in Metaflow. This can be for variety of reasons:
1441
-
1442
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1443
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1444
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1445
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1446
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1452
+ Specifies the flow(s) that this flow depends on.
1447
1453
 
1448
- Usage:
1449
- ----------
1454
+ ```
1455
+ @trigger_on_finish(flow='FooFlow')
1456
+ ```
1457
+ or
1458
+ ```
1459
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1460
+ ```
1461
+ This decorator respects the @project decorator and triggers the flow
1462
+ when upstream runs within the same namespace complete successfully
1450
1463
 
1451
- - Using a custom IAM role to access the datastore.
1464
+ Additionally, you can specify project aware upstream flow dependencies
1465
+ by specifying the fully qualified project_flow_name.
1466
+ ```
1467
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1468
+ ```
1469
+ or
1470
+ ```
1471
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1472
+ ```
1452
1473
 
1453
- ```python
1454
- @with_artifact_store(
1455
- type="s3",
1456
- config=lambda: {
1457
- "root": "s3://my-bucket-foo/path/to/root",
1458
- "role_arn": ROLE,
1459
- },
1460
- )
1461
- class MyFlow(FlowSpec):
1474
+ You can also specify just the project or project branch (other values will be
1475
+ inferred from the current project or project branch):
1476
+ ```
1477
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1478
+ ```
1462
1479
 
1463
- @checkpoint
1464
- @step
1465
- def start(self):
1466
- with open("my_file.txt", "w") as f:
1467
- f.write("Hello, World!")
1468
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1469
- self.next(self.end)
1480
+ Note that `branch` is typically one of:
1481
+ - `prod`
1482
+ - `user.bob`
1483
+ - `test.my_experiment`
1484
+ - `prod.staging`
1470
1485
 
1471
- ```
1472
1486
 
1473
- - Using credentials to access the s3-compatible datastore.
1487
+ Parameters
1488
+ ----------
1489
+ flow : Union[str, Dict[str, str]], optional, default None
1490
+ Upstream flow dependency for this flow.
1491
+ flows : List[Union[str, Dict[str, str]]], default []
1492
+ Upstream flow dependencies for this flow.
1493
+ options : Dict[str, Any], default {}
1494
+ Backend-specific configuration for tuning eventing behavior.
1495
+ """
1496
+ ...
1497
+
1498
+ @typing.overload
1499
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1500
+ ...
1501
+
1502
+ 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] = {}):
1503
+ """
1504
+ Specifies the flow(s) that this flow depends on.
1474
1505
 
1475
- ```python
1476
- @with_artifact_store(
1477
- type="s3",
1478
- config=lambda: {
1479
- "root": "s3://my-bucket-foo/path/to/root",
1480
- "client_params": {
1481
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1482
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1483
- },
1484
- },
1485
- )
1486
- class MyFlow(FlowSpec):
1506
+ ```
1507
+ @trigger_on_finish(flow='FooFlow')
1508
+ ```
1509
+ or
1510
+ ```
1511
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1512
+ ```
1513
+ This decorator respects the @project decorator and triggers the flow
1514
+ when upstream runs within the same namespace complete successfully
1487
1515
 
1488
- @checkpoint
1489
- @step
1490
- def start(self):
1491
- with open("my_file.txt", "w") as f:
1492
- f.write("Hello, World!")
1493
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1494
- self.next(self.end)
1516
+ Additionally, you can specify project aware upstream flow dependencies
1517
+ by specifying the fully qualified project_flow_name.
1518
+ ```
1519
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1520
+ ```
1521
+ or
1522
+ ```
1523
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1524
+ ```
1495
1525
 
1496
- ```
1526
+ You can also specify just the project or project branch (other values will be
1527
+ inferred from the current project or project branch):
1528
+ ```
1529
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1530
+ ```
1497
1531
 
1498
- - Accessing objects stored in external datastores after task execution.
1532
+ Note that `branch` is typically one of:
1533
+ - `prod`
1534
+ - `user.bob`
1535
+ - `test.my_experiment`
1536
+ - `prod.staging`
1499
1537
 
1500
- ```python
1501
- run = Run("CheckpointsTestsFlow/8992")
1502
- with artifact_store_from(run=run, config={
1503
- "client_params": {
1504
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1505
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1506
- },
1507
- }):
1508
- with Checkpoint() as cp:
1509
- latest = cp.list(
1510
- task=run["start"].task
1511
- )[0]
1512
- print(latest)
1513
- cp.load(
1514
- latest,
1515
- "test-checkpoints"
1516
- )
1517
1538
 
1518
- task = Task("TorchTuneFlow/8484/train/53673")
1519
- with artifact_store_from(run=run, config={
1520
- "client_params": {
1521
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1522
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1523
- },
1524
- }):
1525
- load_model(
1526
- task.data.model_ref,
1527
- "test-models"
1528
- )
1529
- ```
1530
- Parameters:
1539
+ Parameters
1531
1540
  ----------
1541
+ flow : Union[str, Dict[str, str]], optional, default None
1542
+ Upstream flow dependency for this flow.
1543
+ flows : List[Union[str, Dict[str, str]]], default []
1544
+ Upstream flow dependencies for this flow.
1545
+ options : Dict[str, Any], default {}
1546
+ Backend-specific configuration for tuning eventing behavior.
1547
+ """
1548
+ ...
1549
+
1550
+ @typing.overload
1551
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1552
+ """
1553
+ Specifies the PyPI packages for all steps of the flow.
1532
1554
 
1533
- type: str
1534
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1555
+ Use `@pypi_base` to set common packages required by all
1556
+ steps and use `@pypi` to specify step-specific overrides.
1535
1557
 
1536
- config: dict or Callable
1537
- Dictionary of configuration options for the datastore. The following keys are required:
1538
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1539
- - example: 's3://bucket-name/path/to/root'
1540
- - example: 'gs://bucket-name/path/to/root'
1541
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1542
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1543
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1544
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1558
+ Parameters
1559
+ ----------
1560
+ packages : Dict[str, str], default: {}
1561
+ Packages to use for this flow. The key is the name of the package
1562
+ and the value is the version to use.
1563
+ python : str, optional, default: None
1564
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1565
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1545
1566
  """
1546
1567
  ...
1547
1568
 
1548
- 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]]:
1569
+ @typing.overload
1570
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1571
+ ...
1572
+
1573
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1549
1574
  """
1550
- 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)
1551
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1552
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1553
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1554
- starts only after all sensors finish.
1575
+ Specifies the PyPI packages for all steps of the flow.
1555
1576
 
1577
+ Use `@pypi_base` to set common packages required by all
1578
+ steps and use `@pypi` to specify step-specific overrides.
1556
1579
 
1557
1580
  Parameters
1558
1581
  ----------
1559
- timeout : int
1560
- Time, in seconds before the task times out and fails. (Default: 3600)
1561
- poke_interval : int
1562
- Time in seconds that the job should wait in between each try. (Default: 60)
1563
- mode : str
1564
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1565
- exponential_backoff : bool
1566
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1567
- pool : str
1568
- the slot pool this task should run in,
1569
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1570
- soft_fail : bool
1571
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1572
- name : str
1573
- Name of the sensor on Airflow
1574
- description : str
1575
- Description of sensor in the Airflow UI
1576
- bucket_key : Union[str, List[str]]
1577
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1578
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1579
- bucket_name : str
1580
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1581
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1582
- wildcard_match : bool
1583
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1584
- aws_conn_id : str
1585
- a reference to the s3 connection on Airflow. (Default: None)
1586
- verify : bool
1587
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1582
+ packages : Dict[str, str], default: {}
1583
+ Packages to use for this flow. The key is the name of the package
1584
+ and the value is the version to use.
1585
+ python : str, optional, default: None
1586
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1587
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1588
1588
  """
1589
1589
  ...
1590
1590