ob-metaflow-stubs 6.0.3.179rc5__py2.py3-none-any.whl → 6.0.3.180rc0__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 +707 -669
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +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 +2 -2
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +1 -1
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +83 -83
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +1 -1
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +1 -1
  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 +1 -1
  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 +1 -1
  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 +1 -1
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +1 -1
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +1 -1
  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 +1 -1
  63. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +1 -1
  64. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +1 -1
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +1 -1
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/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 +3 -3
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +2 -2
  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 +10 -10
  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 +1 -1
  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 +1 -1
  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 +2 -2
  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 +2 -2
  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 +2 -2
  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 +1 -1
  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 +1 -1
  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 +2 -2
  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 +2 -2
  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 +4 -4
  201. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  202. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  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 +1 -1
  207. metaflow-stubs/system/__init__.pyi +1 -1
  208. metaflow-stubs/system/system_logger.pyi +1 -1
  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 +6 -6
  214. metaflow-stubs/user_configs/config_options.pyi +2 -2
  215. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  216. {ob_metaflow_stubs-6.0.3.179rc5.dist-info → ob_metaflow_stubs-6.0.3.180rc0.dist-info}/METADATA +1 -1
  217. ob_metaflow_stubs-6.0.3.180rc0.dist-info/RECORD +220 -0
  218. ob_metaflow_stubs-6.0.3.179rc5.dist-info/RECORD +0 -220
  219. {ob_metaflow_stubs-6.0.3.179rc5.dist-info → ob_metaflow_stubs-6.0.3.180rc0.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.179rc5.dist-info → ob_metaflow_stubs-6.0.3.180rc0.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
3
  # MF version: 2.15.17.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-13T21:29:38.839071 #
4
+ # Generated on 2025-06-17T08:34:56.575242 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -35,18 +35,18 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
+ from . import events as events
38
39
  from . import tuple_util as tuple_util
39
40
  from . import cards as cards
40
41
  from . import metaflow_git as metaflow_git
41
- from . import events as events
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
47
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
49
48
  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,93 +155,248 @@ 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]]]:
159
+ """
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)
180
+ """
181
+ ...
182
+
183
+ def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
184
+ """
185
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
186
+
187
+ User code call
188
+ --------------
189
+ @ollama(
190
+ models=[...],
191
+ ...
192
+ )
193
+
194
+ Valid backend options
195
+ ---------------------
196
+ - 'local': Run as a separate process on the local task machine.
197
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
198
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
199
+
200
+ Valid model options
201
+ -------------------
202
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
203
+
204
+
205
+ Parameters
206
+ ----------
207
+ models: list[str]
208
+ List of Ollama containers running models in sidecars.
209
+ backend: str
210
+ Determines where and how to run the Ollama process.
211
+ force_pull: bool
212
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
213
+ cache_update_policy: str
214
+ Cache update policy: "auto", "force", or "never".
215
+ force_cache_update: bool
216
+ Simple override for "force" cache update policy.
217
+ debug: bool
218
+ Whether to turn on verbose debugging logs.
219
+ circuit_breaker_config: dict
220
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
221
+ timeout_config: dict
222
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
223
+ """
224
+ ...
225
+
158
226
  @typing.overload
159
- 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]]]:
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]]]:
160
228
  """
161
- Specifies environment variables to be set prior to the execution of a step.
229
+ Specifies secrets to be retrieved and injected as environment variables prior to
230
+ the execution of a step.
162
231
 
163
232
 
164
233
  Parameters
165
234
  ----------
166
- vars : Dict[str, str], default {}
167
- Dictionary of environment variables to set.
235
+ sources : List[Union[str, Dict[str, Any]]], default: []
236
+ List of secret specs, defining how the secrets are to be retrieved
168
237
  """
169
238
  ...
170
239
 
171
240
  @typing.overload
172
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
241
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
173
242
  ...
174
243
 
175
244
  @typing.overload
176
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
245
+ def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
177
246
  ...
178
247
 
179
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
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]]] = []):
180
249
  """
181
- Specifies environment variables to be set prior to the execution of a step.
250
+ Specifies secrets to be retrieved and injected as environment variables prior to
251
+ the execution of a step.
182
252
 
183
253
 
184
254
  Parameters
185
255
  ----------
186
- vars : Dict[str, str], default {}
187
- Dictionary of environment variables to set.
256
+ sources : List[Union[str, Dict[str, Any]]], default: []
257
+ List of secret specs, defining how the secrets are to be retrieved
188
258
  """
189
259
  ...
190
260
 
191
261
  @typing.overload
192
- def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
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]]]:
193
263
  """
194
- Enables checkpointing for a step.
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)`.
195
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.
196
274
 
197
275
 
198
276
  Parameters
199
277
  ----------
200
- load_policy : str, default: "fresh"
201
- The policy for loading the checkpoint. The following policies are supported:
202
- - "eager": Loads the the latest available checkpoint within the namespace.
203
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
204
- will be loaded at the start of the task.
205
- - "none": Do not load any checkpoint
206
- - "fresh": Loads the lastest checkpoint created within the running Task.
207
- This mode helps loading checkpoints across various retry attempts of the same task.
208
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
209
- created within the task will be loaded when the task is retries execution on failure.
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.
282
+ """
283
+ ...
284
+
285
+ @typing.overload
286
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
287
+ ...
288
+
289
+ @typing.overload
290
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
291
+ ...
292
+
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):
294
+ """
295
+ Specifies the number of times the task corresponding
296
+ to a step needs to be retried.
210
297
 
211
- temp_dir_root : str, default: None
212
- The root directory under which `current.checkpoint.directory` will be created.
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.
305
+
306
+
307
+ Parameters
308
+ ----------
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.
213
313
  """
214
314
  ...
215
315
 
216
316
  @typing.overload
217
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
317
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
318
+ """
319
+ Decorator prototype for all step decorators. This function gets specialized
320
+ and imported for all decorators types by _import_plugin_decorators().
321
+ """
218
322
  ...
219
323
 
220
324
  @typing.overload
221
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
325
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
222
326
  ...
223
327
 
224
- def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
328
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
225
329
  """
226
- Enables checkpointing for a step.
330
+ Decorator prototype for all step decorators. This function gets specialized
331
+ and imported for all decorators types by _import_plugin_decorators().
332
+ """
333
+ ...
334
+
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]]]:
336
+ """
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'
227
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.
228
356
 
229
357
 
230
358
  Parameters
231
359
  ----------
232
- load_policy : str, default: "fresh"
233
- The policy for loading the checkpoint. The following policies are supported:
234
- - "eager": Loads the the latest available checkpoint within the namespace.
235
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
236
- will be loaded at the start of the task.
237
- - "none": Do not load any checkpoint
238
- - "fresh": Loads the lastest checkpoint created within the running Task.
239
- This mode helps loading checkpoints across various retry attempts of the same task.
240
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
241
- created within the task will be loaded when the task is retries execution on failure.
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`.
370
+ """
371
+ ...
372
+
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]]]:
374
+ """
375
+ Specifies that this step should execute on DGX cloud.
242
376
 
243
- temp_dir_root : str, default: None
244
- The root directory under which `current.checkpoint.directory` will be created.
377
+
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.
390
+
391
+
392
+ Parameters
393
+ ----------
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.
245
400
  """
246
401
  ...
247
402
 
@@ -296,220 +451,111 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
296
451
  """
297
452
  ...
298
453
 
299
- 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]]]:
454
+ @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]]]:
300
456
  """
301
- Specifies that this step should execute on Kubernetes.
457
+ Specifies the PyPI packages for the step.
458
+
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.
302
463
 
303
464
 
304
465
  Parameters
305
466
  ----------
306
- cpu : int, default 1
307
- Number of CPUs required for this step. If `@resources` is
308
- also present, the maximum value from all decorators is used.
309
- memory : int, default 4096
310
- Memory size (in MB) required for this step. If
311
- `@resources` is also present, the maximum value from all decorators is
312
- used.
313
- disk : int, default 10240
314
- Disk size (in MB) required for this step. If
315
- `@resources` is also present, the maximum value from all decorators is
316
- used.
317
- image : str, optional, default None
318
- Docker image to use when launching on Kubernetes. If not specified, and
319
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
320
- not, a default Docker image mapping to the current version of Python is used.
321
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
322
- If given, the imagePullPolicy to be applied to the Docker image of the step.
323
- image_pull_secrets: List[str], default []
324
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
325
- Kubernetes image pull secrets to use when pulling container images
326
- in Kubernetes.
327
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
328
- Kubernetes service account to use when launching pod in Kubernetes.
329
- secrets : List[str], optional, default None
330
- Kubernetes secrets to use when launching pod in Kubernetes. These
331
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
332
- in Metaflow configuration.
333
- node_selector: Union[Dict[str,str], str], optional, default None
334
- Kubernetes node selector(s) to apply to the pod running the task.
335
- Can be passed in as a comma separated string of values e.g.
336
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
337
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
338
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
339
- Kubernetes namespace to use when launching pod in Kubernetes.
340
- gpu : int, optional, default None
341
- Number of GPUs required for this step. A value of zero implies that
342
- the scheduled node should not have GPUs.
343
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
344
- The vendor of the GPUs to be used for this step.
345
- tolerations : List[str], default []
346
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
347
- Kubernetes tolerations to use when launching pod in Kubernetes.
348
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
349
- Kubernetes labels to use when launching pod in Kubernetes.
350
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
351
- Kubernetes annotations to use when launching pod in Kubernetes.
352
- use_tmpfs : bool, default False
353
- This enables an explicit tmpfs mount for this step.
354
- tmpfs_tempdir : bool, default True
355
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
356
- tmpfs_size : int, optional, default: None
357
- The value for the size (in MiB) of the tmpfs mount for this step.
358
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
359
- memory allocated for this step.
360
- tmpfs_path : str, optional, default /metaflow_temp
361
- Path to tmpfs mount for this step.
362
- persistent_volume_claims : Dict[str, str], optional, default None
363
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
364
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
365
- shared_memory: int, optional
366
- Shared memory size (in MiB) required for this step
367
- port: int, optional
368
- Port number to specify in the Kubernetes job object
369
- compute_pool : str, optional, default None
370
- Compute pool to be used for for this step.
371
- If not specified, any accessible compute pool within the perimeter is used.
372
- hostname_resolution_timeout: int, default 10 * 60
373
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
374
- Only applicable when @parallel is used.
375
- qos: str, default: Burstable
376
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
377
-
378
- security_context: Dict[str, Any], optional, default None
379
- Container security context. Applies to the task container. Allows the following keys:
380
- - privileged: bool, optional, default None
381
- - allow_privilege_escalation: bool, optional, default None
382
- - run_as_user: int, optional, default None
383
- - run_as_group: int, optional, default None
384
- - run_as_non_root: bool, optional, default None
385
- """
386
- ...
387
-
388
- @typing.overload
389
- 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]]]:
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.
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.
421
473
  """
422
474
  ...
423
475
 
424
476
  @typing.overload
425
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
477
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
426
478
  ...
427
479
 
428
480
  @typing.overload
429
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
481
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
430
482
  ...
431
483
 
432
- 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):
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):
433
485
  """
434
- Specifies the resources needed when executing this step.
435
-
436
- Use `@resources` to specify the resource requirements
437
- independently of the specific compute layer (`@batch`, `@kubernetes`).
486
+ Specifies the PyPI packages for the step.
438
487
 
439
- You can choose the compute layer on the command line by executing e.g.
440
- ```
441
- python myflow.py run --with batch
442
- ```
443
- or
444
- ```
445
- python myflow.py run --with kubernetes
446
- ```
447
- which executes the flow on the desired system using the
448
- requirements specified in `@resources`.
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.
449
492
 
450
493
 
451
494
  Parameters
452
495
  ----------
453
- cpu : int, default 1
454
- Number of CPUs required for this step.
455
- gpu : int, optional, default None
456
- Number of GPUs required for this step.
457
- disk : int, optional, default None
458
- Disk size (in MB) required for this step. Only applies on Kubernetes.
459
- memory : int, default 4096
460
- Memory size (in MB) required for this step.
461
- shared_memory : int, optional, default None
462
- The value for the size (in MiB) of the /dev/shm volume for this step.
463
- This parameter maps to the `--shm-size` option in Docker.
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.
464
502
  """
465
503
  ...
466
504
 
467
505
  @typing.overload
468
- 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]]]:
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]]]:
469
507
  """
470
- Creates a human-readable report, a Metaflow Card, after this step completes.
508
+ Enables loading / saving of models within a step.
471
509
 
472
- Note that you may add multiple `@card` decorators in a step with different parameters.
473
510
 
474
511
 
475
512
  Parameters
476
513
  ----------
477
- type : str, default 'default'
478
- Card type.
479
- id : str, optional, default None
480
- If multiple cards are present, use this id to identify this card.
481
- options : Dict[str, Any], default {}
482
- Options passed to the card. The contents depend on the card type.
483
- timeout : int, default 45
484
- Interrupt reporting if it takes more than this many seconds.
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`
520
+
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
+
525
+ temp_dir_root : str, default: None
526
+ The root directory under which `current.model.loaded` will store loaded models
485
527
  """
486
528
  ...
487
529
 
488
530
  @typing.overload
489
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
531
+ def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
490
532
  ...
491
533
 
492
534
  @typing.overload
493
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
535
+ def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
494
536
  ...
495
537
 
496
- 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):
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):
497
539
  """
498
- Creates a human-readable report, a Metaflow Card, after this step completes.
540
+ Enables loading / saving of models within a step.
499
541
 
500
- Note that you may add multiple `@card` decorators in a step with different parameters.
501
542
 
502
543
 
503
544
  Parameters
504
545
  ----------
505
- type : str, default 'default'
506
- Card type.
507
- id : str, optional, default None
508
- If multiple cards are present, use this id to identify this card.
509
- options : Dict[str, Any], default {}
510
- Options passed to the card. The contents depend on the card type.
511
- timeout : int, default 45
512
- Interrupt reporting if it takes more than this many seconds.
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`
552
+
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
+
557
+ temp_dir_root : str, default: None
558
+ The root directory under which `current.model.loaded` will store loaded models
513
559
  """
514
560
  ...
515
561
 
@@ -528,141 +574,125 @@ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union
528
574
  """
529
575
  ...
530
576
 
531
- 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]]]:
577
+ @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]]]:
532
579
  """
533
- Decorator that helps cache, version and store models/datasets from huggingface hub.
580
+ Specifies environment variables to be set prior to the execution of a step.
534
581
 
535
582
 
536
583
  Parameters
537
584
  ----------
538
- temp_dir_root : str, optional
539
- The root directory that will hold the temporary directory where objects will be downloaded.
540
-
541
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
542
- The list of repos (models/datasets) to load.
543
-
544
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
545
-
546
- - If repo (model/dataset) is not found in the datastore:
547
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
548
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
549
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
550
-
551
- - If repo is found in the datastore:
552
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
585
+ vars : Dict[str, str], default {}
586
+ Dictionary of environment variables to set.
553
587
  """
554
588
  ...
555
589
 
556
590
  @typing.overload
557
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
558
- """
559
- Decorator prototype for all step decorators. This function gets specialized
560
- and imported for all decorators types by _import_plugin_decorators().
561
- """
591
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
562
592
  ...
563
593
 
564
594
  @typing.overload
565
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
566
- ...
567
-
568
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
569
- """
570
- Decorator prototype for all step decorators. This function gets specialized
571
- and imported for all decorators types by _import_plugin_decorators().
572
- """
595
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
573
596
  ...
574
597
 
575
- @typing.overload
576
- 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]]]:
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] = {}):
577
599
  """
578
- Specifies the PyPI packages for the step.
579
-
580
- Information in this decorator will augment any
581
- attributes set in the `@pyi_base` flow-level decorator. Hence,
582
- you can use `@pypi_base` to set packages required by all
583
- steps and use `@pypi` to specify step-specific overrides.
600
+ Specifies environment variables to be set prior to the execution of a step.
584
601
 
585
602
 
586
603
  Parameters
587
604
  ----------
588
- packages : Dict[str, str], default: {}
589
- Packages to use for this step. The key is the name of the package
590
- and the value is the version to use.
591
- python : str, optional, default: None
592
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
593
- that the version used will correspond to the version of the Python interpreter used to start the run.
605
+ vars : Dict[str, str], default {}
606
+ Dictionary of environment variables to set.
594
607
  """
595
608
  ...
596
609
 
597
- @typing.overload
598
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
599
- ...
600
-
601
- @typing.overload
602
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
603
- ...
604
-
605
- 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):
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]]]:
606
611
  """
607
- Specifies the PyPI packages for the step.
608
-
609
- Information in this decorator will augment any
610
- attributes set in the `@pyi_base` flow-level decorator. Hence,
611
- you can use `@pypi_base` to set packages required by all
612
- steps and use `@pypi` to specify step-specific overrides.
612
+ Specifies that this step should execute on Kubernetes.
613
613
 
614
614
 
615
615
  Parameters
616
616
  ----------
617
- packages : Dict[str, str], default: {}
618
- Packages to use for this step. The key is the name of the package
619
- and the value is the version to use.
620
- python : str, optional, default: None
621
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
622
- that the version used will correspond to the version of the Python interpreter used to start the run.
623
- """
624
- ...
625
-
626
- def ollama(*, models: list, backend: str, force_pull: bool, cache_update_policy: str, force_cache_update: bool, debug: bool, circuit_breaker_config: dict, timeout_config: dict) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
627
- """
628
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
629
-
630
- User code call
631
- --------------
632
- @ollama(
633
- models=[...],
634
- ...
635
- )
636
-
637
- Valid backend options
638
- ---------------------
639
- - 'local': Run as a separate process on the local task machine.
640
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
641
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
642
-
643
- Valid model options
644
- -------------------
645
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
646
-
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
647
688
 
648
- Parameters
649
- ----------
650
- models: list[str]
651
- List of Ollama containers running models in sidecars.
652
- backend: str
653
- Determines where and how to run the Ollama process.
654
- force_pull: bool
655
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
656
- cache_update_policy: str
657
- Cache update policy: "auto", "force", or "never".
658
- force_cache_update: bool
659
- Simple override for "force" cache update policy.
660
- debug: bool
661
- Whether to turn on verbose debugging logs.
662
- circuit_breaker_config: dict
663
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
664
- timeout_config: dict
665
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
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
666
696
  """
667
697
  ...
668
698
 
@@ -726,149 +756,204 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
726
756
  ...
727
757
 
728
758
  @typing.overload
729
- 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]]]:
759
+ def checkpoint(*, load_policy: str = 'fresh', temp_dir_root: str = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
730
760
  """
731
- Specifies the number of times the task corresponding
732
- to a step needs to be retried.
733
-
734
- This decorator is useful for handling transient errors, such as networking issues.
735
- If your task contains operations that can't be retried safely, e.g. database updates,
736
- it is advisable to annotate it with `@retry(times=0)`.
761
+ Enables checkpointing for a step.
737
762
 
738
- This can be used in conjunction with the `@catch` decorator. The `@catch`
739
- decorator will execute a no-op task after all retries have been exhausted,
740
- ensuring that the flow execution can continue.
741
763
 
742
764
 
743
765
  Parameters
744
766
  ----------
745
- times : int, default 3
746
- Number of times to retry this task.
747
- minutes_between_retries : int, default 2
748
- Number of minutes between retries.
767
+ load_policy : str, default: "fresh"
768
+ The policy for loading the checkpoint. The following policies are supported:
769
+ - "eager": Loads the the latest available checkpoint within the namespace.
770
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
771
+ will be loaded at the start of the task.
772
+ - "none": Do not load any checkpoint
773
+ - "fresh": Loads the lastest checkpoint created within the running Task.
774
+ This mode helps loading checkpoints across various retry attempts of the same task.
775
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
776
+ created within the task will be loaded when the task is retries execution on failure.
777
+
778
+ temp_dir_root : str, default: None
779
+ The root directory under which `current.checkpoint.directory` will be created.
749
780
  """
750
781
  ...
751
782
 
752
783
  @typing.overload
753
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
784
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
754
785
  ...
755
786
 
756
787
  @typing.overload
757
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
788
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
758
789
  ...
759
790
 
760
- 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):
791
+ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, load_policy: str = 'fresh', temp_dir_root: str = None):
761
792
  """
762
- Specifies the number of times the task corresponding
763
- to a step needs to be retried.
764
-
765
- This decorator is useful for handling transient errors, such as networking issues.
766
- If your task contains operations that can't be retried safely, e.g. database updates,
767
- it is advisable to annotate it with `@retry(times=0)`.
793
+ Enables checkpointing for a step.
768
794
 
769
- This can be used in conjunction with the `@catch` decorator. The `@catch`
770
- decorator will execute a no-op task after all retries have been exhausted,
771
- ensuring that the flow execution can continue.
772
795
 
773
796
 
774
797
  Parameters
775
798
  ----------
776
- times : int, default 3
777
- Number of times to retry this task.
778
- minutes_between_retries : int, default 2
779
- Number of minutes between retries.
799
+ load_policy : str, default: "fresh"
800
+ The policy for loading the checkpoint. The following policies are supported:
801
+ - "eager": Loads the the latest available checkpoint within the namespace.
802
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
803
+ will be loaded at the start of the task.
804
+ - "none": Do not load any checkpoint
805
+ - "fresh": Loads the lastest checkpoint created within the running Task.
806
+ This mode helps loading checkpoints across various retry attempts of the same task.
807
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
808
+ created within the task will be loaded when the task is retries execution on failure.
809
+
810
+ temp_dir_root : str, default: None
811
+ The root directory under which `current.checkpoint.directory` will be created.
780
812
  """
781
813
  ...
782
814
 
783
815
  @typing.overload
784
- 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]]]:
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]]]:
785
817
  """
786
- Enables loading / saving of models within a step.
818
+ Creates a human-readable report, a Metaflow Card, after this step completes.
787
819
 
820
+ Note that you may add multiple `@card` decorators in a step with different parameters.
788
821
 
789
822
 
790
823
  Parameters
791
824
  ----------
792
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
793
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
794
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
795
- - `current.checkpoint`
796
- - `current.model`
797
- - `current.huggingface_hub`
798
-
799
- 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
800
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
801
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
802
-
803
- temp_dir_root : str, default: None
804
- The root directory under which `current.model.loaded` will store loaded models
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.
805
833
  """
806
834
  ...
807
835
 
808
836
  @typing.overload
809
- def model(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
837
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
810
838
  ...
811
839
 
812
840
  @typing.overload
813
- def model(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
841
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
814
842
  ...
815
843
 
816
- 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):
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):
817
845
  """
818
- Enables loading / saving of models within a step.
846
+ Creates a human-readable report, a Metaflow Card, after this step completes.
819
847
 
848
+ Note that you may add multiple `@card` decorators in a step with different parameters.
820
849
 
821
850
 
822
851
  Parameters
823
852
  ----------
824
- load : Union[List[str],str,List[Tuple[str,Union[str,None]]]], default: None
825
- Artifact name/s referencing the models/checkpoints to load. Artifact names refer to the names of the instance variables set to `self`.
826
- These artifact names give to `load` be reference objects or reference `key` string's from objects created by:
827
- - `current.checkpoint`
828
- - `current.model`
829
- - `current.huggingface_hub`
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.
861
+ """
862
+ ...
863
+
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
+ """
867
+ Specifies the resources needed when executing this step.
830
868
 
831
- 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
832
- the local filesystem. If the second element is None, the artifact will be unpacked in the current working directory.
833
- If a string is provided, then the artifact corresponding to that name will be loaded in the current working directory.
869
+ Use `@resources` to specify the resource requirements
870
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
834
871
 
835
- temp_dir_root : str, default: None
836
- The root directory under which `current.model.loaded` will store loaded models
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`.
882
+
883
+
884
+ Parameters
885
+ ----------
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.
837
897
  """
838
898
  ...
839
899
 
840
900
  @typing.overload
841
- 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]]]:
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):
842
909
  """
843
- Specifies secrets to be retrieved and injected as environment variables prior to
844
- the execution of a step.
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`).
914
+
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`.
845
925
 
846
926
 
847
927
  Parameters
848
928
  ----------
849
- sources : List[Union[str, Dict[str, Any]]], default: []
850
- List of secret specs, defining how the secrets are to be retrieved
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.
851
940
  """
852
941
  ...
853
942
 
854
943
  @typing.overload
855
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
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
+ """
856
948
  ...
857
949
 
858
950
  @typing.overload
859
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
951
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
860
952
  ...
861
953
 
862
- 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]]] = []):
954
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
863
955
  """
864
- Specifies secrets to be retrieved and injected as environment variables prior to
865
- the execution of a step.
866
-
867
-
868
- Parameters
869
- ----------
870
- sources : List[Union[str, Dict[str, Any]]], default: []
871
- List of secret specs, defining how the secrets are to be retrieved
956
+ Internal decorator to support Fast bakery
872
957
  """
873
958
  ...
874
959
 
@@ -931,93 +1016,44 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
931
1016
  """
932
1017
  ...
933
1018
 
934
- 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]]]:
1019
+ @typing.overload
1020
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
935
1021
  """
936
- Specifies that this step should execute on DGX cloud.
1022
+ Specifies the PyPI packages for all steps of the flow.
937
1023
 
1024
+ Use `@pypi_base` to set common packages required by all
1025
+ steps and use `@pypi` to specify step-specific overrides.
938
1026
 
939
1027
  Parameters
940
1028
  ----------
941
- gpu : int
942
- Number of GPUs to use.
943
- gpu_type : str
944
- Type of Nvidia GPU to use.
945
- """
946
- ...
947
-
948
- @typing.overload
949
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
950
- """
951
- Internal decorator to support Fast bakery
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.
952
1035
  """
953
1036
  ...
954
1037
 
955
1038
  @typing.overload
956
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
957
- ...
958
-
959
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
960
- """
961
- Internal decorator to support Fast bakery
962
- """
963
- ...
964
-
965
- 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]]]:
966
- """
967
- Specifies that this step should execute on DGX cloud.
968
-
969
-
970
- Parameters
971
- ----------
972
- gpu : int
973
- Number of GPUs to use.
974
- gpu_type : str
975
- Type of Nvidia GPU to use.
976
- queue_timeout : int
977
- Time to keep the job in NVCF's queue.
978
- """
1039
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
979
1040
  ...
980
1041
 
981
- 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]]:
1042
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
982
1043
  """
983
- 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.
984
- 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.
1044
+ Specifies the PyPI packages for all steps of the flow.
985
1045
 
1046
+ Use `@pypi_base` to set common packages required by all
1047
+ steps and use `@pypi` to specify step-specific overrides.
986
1048
 
987
1049
  Parameters
988
1050
  ----------
989
- timeout : int
990
- Time, in seconds before the task times out and fails. (Default: 3600)
991
- poke_interval : int
992
- Time in seconds that the job should wait in between each try. (Default: 60)
993
- mode : str
994
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
995
- exponential_backoff : bool
996
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
997
- pool : str
998
- the slot pool this task should run in,
999
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1000
- soft_fail : bool
1001
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1002
- name : str
1003
- Name of the sensor on Airflow
1004
- description : str
1005
- Description of sensor in the Airflow UI
1006
- external_dag_id : str
1007
- The dag_id that contains the task you want to wait for.
1008
- external_task_ids : List[str]
1009
- The list of task_ids that you want to wait for.
1010
- If None (default value) the sensor waits for the DAG. (Default: None)
1011
- allowed_states : List[str]
1012
- Iterable of allowed states, (Default: ['success'])
1013
- failed_states : List[str]
1014
- Iterable of failed or dis-allowed states. (Default: None)
1015
- execution_delta : datetime.timedelta
1016
- time difference with the previous execution to look at,
1017
- the default is the same logical date as the current task or DAG. (Default: None)
1018
- check_existence: bool
1019
- Set to True to check if the external task exists or check if
1020
- the DAG to wait for exists. (Default: True)
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.
1021
1057
  """
1022
1058
  ...
1023
1059
 
@@ -1096,215 +1132,80 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1096
1132
  or
1097
1133
  ```
1098
1134
  @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1099
- ```
1100
-
1101
- You can also specify just the project or project branch (other values will be
1102
- inferred from the current project or project branch):
1103
- ```
1104
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1105
- ```
1106
-
1107
- Note that `branch` is typically one of:
1108
- - `prod`
1109
- - `user.bob`
1110
- - `test.my_experiment`
1111
- - `prod.staging`
1112
-
1113
-
1114
- Parameters
1115
- ----------
1116
- flow : Union[str, Dict[str, str]], optional, default None
1117
- Upstream flow dependency for this flow.
1118
- flows : List[Union[str, Dict[str, str]]], default []
1119
- Upstream flow dependencies for this flow.
1120
- options : Dict[str, Any], default {}
1121
- Backend-specific configuration for tuning eventing behavior.
1122
- """
1123
- ...
1124
-
1125
- @typing.overload
1126
- 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]]:
1127
- """
1128
- Specifies the Conda environment for all steps of the flow.
1129
-
1130
- Use `@conda_base` to set common libraries required by all
1131
- steps and use `@conda` to specify step-specific additions.
1132
-
1133
-
1134
- Parameters
1135
- ----------
1136
- packages : Dict[str, str], default {}
1137
- Packages to use for this flow. The key is the name of the package
1138
- and the value is the version to use.
1139
- libraries : Dict[str, str], default {}
1140
- Supported for backward compatibility. When used with packages, packages will take precedence.
1141
- python : str, optional, default None
1142
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1143
- that the version used will correspond to the version of the Python interpreter used to start the run.
1144
- disabled : bool, default False
1145
- If set to True, disables Conda.
1146
- """
1147
- ...
1148
-
1149
- @typing.overload
1150
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1151
- ...
1152
-
1153
- 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):
1154
- """
1155
- Specifies the Conda environment for all steps of the flow.
1156
-
1157
- Use `@conda_base` to set common libraries required by all
1158
- steps and use `@conda` to specify step-specific additions.
1159
-
1160
-
1161
- Parameters
1162
- ----------
1163
- packages : Dict[str, str], default {}
1164
- Packages to use for this flow. The key is the name of the package
1165
- and the value is the version to use.
1166
- libraries : Dict[str, str], default {}
1167
- Supported for backward compatibility. When used with packages, packages will take precedence.
1168
- python : str, optional, default None
1169
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1170
- that the version used will correspond to the version of the Python interpreter used to start the run.
1171
- disabled : bool, default False
1172
- If set to True, disables Conda.
1173
- """
1174
- ...
1175
-
1176
- 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]]:
1177
- """
1178
- 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)
1179
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1180
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1181
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1182
- starts only after all sensors finish.
1183
-
1184
-
1185
- Parameters
1186
- ----------
1187
- timeout : int
1188
- Time, in seconds before the task times out and fails. (Default: 3600)
1189
- poke_interval : int
1190
- Time in seconds that the job should wait in between each try. (Default: 60)
1191
- mode : str
1192
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1193
- exponential_backoff : bool
1194
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1195
- pool : str
1196
- the slot pool this task should run in,
1197
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1198
- soft_fail : bool
1199
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1200
- name : str
1201
- Name of the sensor on Airflow
1202
- description : str
1203
- Description of sensor in the Airflow UI
1204
- bucket_key : Union[str, List[str]]
1205
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1206
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1207
- bucket_name : str
1208
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1209
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1210
- wildcard_match : bool
1211
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1212
- aws_conn_id : str
1213
- a reference to the s3 connection on Airflow. (Default: None)
1214
- verify : bool
1215
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1216
- """
1217
- ...
1218
-
1219
- @typing.overload
1220
- 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]]:
1221
- """
1222
- Specifies the times when the flow should be run when running on a
1223
- production scheduler.
1224
-
1225
-
1226
- Parameters
1227
- ----------
1228
- hourly : bool, default False
1229
- Run the workflow hourly.
1230
- daily : bool, default True
1231
- Run the workflow daily.
1232
- weekly : bool, default False
1233
- Run the workflow weekly.
1234
- cron : str, optional, default None
1235
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1236
- specified by this expression.
1237
- timezone : str, optional, default None
1238
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1239
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1240
- """
1241
- ...
1242
-
1243
- @typing.overload
1244
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1245
- ...
1246
-
1247
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1248
- """
1249
- Specifies the times when the flow should be run when running on a
1250
- production scheduler.
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`
1251
1148
 
1252
1149
 
1253
1150
  Parameters
1254
1151
  ----------
1255
- hourly : bool, default False
1256
- Run the workflow hourly.
1257
- daily : bool, default True
1258
- Run the workflow daily.
1259
- weekly : bool, default False
1260
- Run the workflow weekly.
1261
- cron : str, optional, default None
1262
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1263
- specified by this expression.
1264
- timezone : str, optional, default None
1265
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1266
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
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.
1267
1158
  """
1268
1159
  ...
1269
1160
 
1270
1161
  @typing.overload
1271
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
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]]:
1272
1163
  """
1273
- Specifies the PyPI packages for all steps of the flow.
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.
1274
1168
 
1275
- Use `@pypi_base` to set common packages required by all
1276
- steps and use `@pypi` to specify step-specific overrides.
1277
1169
 
1278
1170
  Parameters
1279
1171
  ----------
1280
- packages : Dict[str, str], default: {}
1172
+ packages : Dict[str, str], default {}
1281
1173
  Packages to use for this flow. The key is the name of the package
1282
1174
  and the value is the version to use.
1283
- python : str, optional, default: None
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
1284
1178
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1285
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.
1286
1182
  """
1287
1183
  ...
1288
1184
 
1289
1185
  @typing.overload
1290
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1186
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1291
1187
  ...
1292
1188
 
1293
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
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):
1294
1190
  """
1295
- Specifies the PyPI packages for all steps of the flow.
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.
1296
1195
 
1297
- Use `@pypi_base` to set common packages required by all
1298
- steps and use `@pypi` to specify step-specific overrides.
1299
1196
 
1300
1197
  Parameters
1301
1198
  ----------
1302
- packages : Dict[str, str], default: {}
1199
+ packages : Dict[str, str], default {}
1303
1200
  Packages to use for this flow. The key is the name of the package
1304
1201
  and the value is the version to use.
1305
- python : str, optional, default: None
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
1306
1205
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
1307
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.
1308
1209
  """
1309
1210
  ...
1310
1211
 
@@ -1401,6 +1302,135 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1401
1302
  """
1402
1303
  ...
1403
1304
 
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]]:
1306
+ """
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.
1309
+
1310
+
1311
+ Parameters
1312
+ ----------
1313
+ timeout : int
1314
+ Time, in seconds before the task times out and fails. (Default: 3600)
1315
+ poke_interval : int
1316
+ Time in seconds that the job should wait in between each try. (Default: 60)
1317
+ mode : str
1318
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1319
+ exponential_backoff : bool
1320
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1321
+ pool : str
1322
+ the slot pool this task should run in,
1323
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1324
+ soft_fail : bool
1325
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1326
+ name : str
1327
+ Name of the sensor on Airflow
1328
+ description : str
1329
+ 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)
1345
+ """
1346
+ ...
1347
+
1348
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1349
+ """
1350
+ Specifies what flows belong to the same project.
1351
+
1352
+ A project-specific namespace is created for all flows that
1353
+ use the same `@project(name)`.
1354
+
1355
+
1356
+ Parameters
1357
+ ----------
1358
+ name : str
1359
+ Project name. Make sure that the name is unique amongst all
1360
+ projects that use the same production scheduler. The name may
1361
+ contain only lowercase alphanumeric characters and underscores.
1362
+
1363
+ branch : Optional[str], default None
1364
+ The branch to use. If not specified, the branch is set to
1365
+ `user.<username>` unless `production` is set to `True`. This can
1366
+ also be set on the command line using `--branch` as a top-level option.
1367
+ It is an error to specify `branch` in the decorator and on the command line.
1368
+
1369
+ production : bool, default False
1370
+ Whether or not the branch is the production branch. This can also be set on the
1371
+ command line using `--production` as a top-level option. It is an error to specify
1372
+ `production` in the decorator and on the command line.
1373
+ The project branch name will be:
1374
+ - if `branch` is specified:
1375
+ - if `production` is True: `prod.<branch>`
1376
+ - if `production` is False: `test.<branch>`
1377
+ - if `branch` is not specified:
1378
+ - if `production` is True: `prod`
1379
+ - if `production` is False: `user.<username>`
1380
+ """
1381
+ ...
1382
+
1383
+ @typing.overload
1384
+ 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
+ """
1386
+ Specifies the times when the flow should be run when running on a
1387
+ production scheduler.
1388
+
1389
+
1390
+ Parameters
1391
+ ----------
1392
+ hourly : bool, default False
1393
+ Run the workflow hourly.
1394
+ daily : bool, default True
1395
+ Run the workflow daily.
1396
+ weekly : bool, default False
1397
+ Run the workflow weekly.
1398
+ cron : str, optional, default None
1399
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1400
+ specified by this expression.
1401
+ timezone : str, optional, default None
1402
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1403
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1404
+ """
1405
+ ...
1406
+
1407
+ @typing.overload
1408
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1409
+ ...
1410
+
1411
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1412
+ """
1413
+ Specifies the times when the flow should be run when running on a
1414
+ production scheduler.
1415
+
1416
+
1417
+ Parameters
1418
+ ----------
1419
+ hourly : bool, default False
1420
+ Run the workflow hourly.
1421
+ daily : bool, default True
1422
+ Run the workflow daily.
1423
+ weekly : bool, default False
1424
+ Run the workflow weekly.
1425
+ cron : str, optional, default None
1426
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1427
+ specified by this expression.
1428
+ timezone : str, optional, default None
1429
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1430
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1431
+ """
1432
+ ...
1433
+
1404
1434
  def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1405
1435
  """
1406
1436
  Allows setting external datastores to save data for the
@@ -1515,38 +1545,46 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1515
1545
  """
1516
1546
  ...
1517
1547
 
1518
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
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]]:
1519
1549
  """
1520
- Specifies what flows belong to the same project.
1521
-
1522
- A project-specific namespace is created for all flows that
1523
- use the same `@project(name)`.
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.
1524
1555
 
1525
1556
 
1526
1557
  Parameters
1527
1558
  ----------
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)
1528
1572
  name : str
1529
- Project name. Make sure that the name is unique amongst all
1530
- projects that use the same production scheduler. The name may
1531
- contain only lowercase alphanumeric characters and underscores.
1532
-
1533
- branch : Optional[str], default None
1534
- The branch to use. If not specified, the branch is set to
1535
- `user.<username>` unless `production` is set to `True`. This can
1536
- also be set on the command line using `--branch` as a top-level option.
1537
- It is an error to specify `branch` in the decorator and on the command line.
1538
-
1539
- production : bool, default False
1540
- Whether or not the branch is the production branch. This can also be set on the
1541
- command line using `--production` as a top-level option. It is an error to specify
1542
- `production` in the decorator and on the command line.
1543
- The project branch name will be:
1544
- - if `branch` is specified:
1545
- - if `production` is True: `prod.<branch>`
1546
- - if `production` is False: `test.<branch>`
1547
- - if `branch` is not specified:
1548
- - if `production` is True: `prod`
1549
- - if `production` is False: `user.<username>`
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)
1550
1588
  """
1551
1589
  ...
1552
1590