ob-metaflow-stubs 6.0.3.176rc4__py2.py3-none-any.whl → 6.0.3.176rc5__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 (218) hide show
  1. metaflow-stubs/__init__.pyi +702 -703
  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 +3 -3
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +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 +84 -84
  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 +1 -1
  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 +1 -1
  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 +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/constructors.pyi +1 -1
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +1 -1
  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 +1 -1
  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 +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +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/card_utilities/__init__.pyi +1 -1
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/async_cards.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +1 -1
  69. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/__init__.pyi +1 -1
  70. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/coreweave.pyi +1 -1
  71. metaflow-stubs/mf_extensions/outerbounds/plugins/checkpoint_datastores/nebius.pyi +1 -1
  72. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  73. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  74. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -2
  86. metaflow-stubs/multicore_utils.pyi +1 -1
  87. metaflow-stubs/parameters.pyi +2 -2
  88. metaflow-stubs/plugins/__init__.pyi +10 -10
  89. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  90. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  91. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  92. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  93. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  94. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  95. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  96. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  97. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  98. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  99. metaflow-stubs/plugins/argo/argo_workflows.pyi +1 -1
  100. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  101. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +1 -1
  102. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +2 -2
  103. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  104. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  105. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  106. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  107. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  108. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  109. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  110. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  111. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +1 -1
  112. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  113. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  114. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  115. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  116. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  117. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +2 -2
  118. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  119. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  120. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  121. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  122. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +1 -1
  123. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  124. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  125. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  126. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  128. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  129. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  130. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  131. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  132. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  133. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_modules/components.pyi +2 -2
  135. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  139. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  140. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  141. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  142. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  143. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  144. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  145. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  146. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  147. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  148. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  149. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  150. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  151. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  152. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  153. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  154. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  155. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +1 -1
  156. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  157. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  158. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  159. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  160. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  161. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +1 -1
  162. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  163. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  164. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  165. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  166. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  167. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  168. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  169. metaflow-stubs/plugins/perimeters.pyi +1 -1
  170. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  171. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  173. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  174. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  175. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  177. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  178. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  179. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  180. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  181. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  182. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  184. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  185. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  186. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  187. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  188. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  189. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  190. metaflow-stubs/profilers/__init__.pyi +1 -1
  191. metaflow-stubs/pylint_wrapper.pyi +1 -1
  192. metaflow-stubs/runner/__init__.pyi +1 -1
  193. metaflow-stubs/runner/deployer.pyi +26 -26
  194. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  195. metaflow-stubs/runner/metaflow_runner.pyi +1 -1
  196. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  197. metaflow-stubs/runner/nbrun.pyi +1 -1
  198. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  199. metaflow-stubs/runner/utils.pyi +2 -2
  200. metaflow-stubs/system/__init__.pyi +1 -1
  201. metaflow-stubs/system/system_logger.pyi +2 -2
  202. metaflow-stubs/system/system_monitor.pyi +1 -1
  203. metaflow-stubs/tagging_util.pyi +1 -1
  204. metaflow-stubs/tuple_util.pyi +1 -1
  205. metaflow-stubs/user_configs/__init__.pyi +1 -1
  206. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  207. metaflow-stubs/user_configs/config_options.pyi +1 -1
  208. metaflow-stubs/user_configs/config_parameters.pyi +4 -4
  209. {ob_metaflow_stubs-6.0.3.176rc4.dist-info → ob_metaflow_stubs-6.0.3.176rc5.dist-info}/METADATA +1 -1
  210. ob_metaflow_stubs-6.0.3.176rc5.dist-info/RECORD +213 -0
  211. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/__init__.pyi +0 -6
  212. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +0 -51
  213. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +0 -65
  214. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +0 -74
  215. metaflow-stubs/ob_internal.pyi +0 -11
  216. ob_metaflow_stubs-6.0.3.176rc4.dist-info/RECORD +0 -218
  217. {ob_metaflow_stubs-6.0.3.176rc4.dist-info → ob_metaflow_stubs-6.0.3.176rc5.dist-info}/WHEEL +0 -0
  218. {ob_metaflow_stubs-6.0.3.176rc4.dist-info → ob_metaflow_stubs-6.0.3.176rc5.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.14.1+obcheckpoint(0.2.1);ob(v1) #
4
- # Generated on 2025-06-04T19:55:23.864727 #
4
+ # Generated on 2025-06-04T21:50:56.035194 #
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 tuple_util as tuple_util
38
39
  from . import cards as cards
39
40
  from . import events as events
40
41
  from . import metaflow_git as metaflow_git
41
- from . import tuple_util as tuple_util
42
42
  from . import runner as runner
43
43
  from . import plugins as plugins
44
44
  from .mf_extensions.outerbounds.toplevel.global_aliases_for_metaflow_package import S3 as S3
45
45
  from . import includefile as includefile
46
46
  from .includefile import IncludeFile as IncludeFile
47
47
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
48
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_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
@@ -77,7 +77,6 @@ from . import system as system
77
77
  from . import pylint_wrapper as pylint_wrapper
78
78
  from . import cli as cli
79
79
  from . import profilers as profilers
80
- from . import ob_internal as ob_internal
81
80
 
82
81
  EXT_PKG: str
83
82
 
@@ -154,58 +153,28 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
154
153
  """
155
154
  ...
156
155
 
157
- @typing.overload
158
- 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]]]:
156
+ 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
157
  """
160
- Specifies the number of times the task corresponding
161
- to a step needs to be retried.
162
-
163
- This decorator is useful for handling transient errors, such as networking issues.
164
- If your task contains operations that can't be retried safely, e.g. database updates,
165
- it is advisable to annotate it with `@retry(times=0)`.
166
-
167
- This can be used in conjunction with the `@catch` decorator. The `@catch`
168
- decorator will execute a no-op task after all retries have been exhausted,
169
- ensuring that the flow execution can continue.
158
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
170
159
 
171
160
 
172
161
  Parameters
173
162
  ----------
174
- times : int, default 3
175
- Number of times to retry this task.
176
- minutes_between_retries : int, default 2
177
- Number of minutes between retries.
178
- """
179
- ...
180
-
181
- @typing.overload
182
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
183
- ...
184
-
185
- @typing.overload
186
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
187
- ...
188
-
189
- 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):
190
- """
191
- Specifies the number of times the task corresponding
192
- to a step needs to be retried.
163
+ temp_dir_root : str, optional
164
+ The root directory that will hold the temporary directory where objects will be downloaded.
193
165
 
194
- This decorator is useful for handling transient errors, such as networking issues.
195
- If your task contains operations that can't be retried safely, e.g. database updates,
196
- it is advisable to annotate it with `@retry(times=0)`.
166
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
167
+ The list of repos (models/datasets) to load.
197
168
 
198
- This can be used in conjunction with the `@catch` decorator. The `@catch`
199
- decorator will execute a no-op task after all retries have been exhausted,
200
- ensuring that the flow execution can continue.
169
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
201
170
 
171
+ - If repo (model/dataset) is not found in the datastore:
172
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
173
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
174
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
202
175
 
203
- Parameters
204
- ----------
205
- times : int, default 3
206
- Number of times to retry this task.
207
- minutes_between_retries : int, default 2
208
- Number of minutes between retries.
176
+ - If repo is found in the datastore:
177
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
209
178
  """
210
179
  ...
211
180
 
@@ -223,61 +192,22 @@ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Cal
223
192
  """
224
193
  ...
225
194
 
226
- 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]]]:
195
+ @typing.overload
196
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
227
197
  """
228
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
229
-
230
- User code call
231
- --------------
232
- @ollama(
233
- models=[...],
234
- ...
235
- )
236
-
237
- Valid backend options
238
- ---------------------
239
- - 'local': Run as a separate process on the local task machine.
240
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
241
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
242
-
243
- Valid model options
244
- -------------------
245
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
246
-
247
-
248
- Parameters
249
- ----------
250
- models: list[str]
251
- List of Ollama containers running models in sidecars.
252
- backend: str
253
- Determines where and how to run the Ollama process.
254
- force_pull: bool
255
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
256
- cache_update_policy: str
257
- Cache update policy: "auto", "force", or "never".
258
- force_cache_update: bool
259
- Simple override for "force" cache update policy.
260
- debug: bool
261
- Whether to turn on verbose debugging logs.
262
- circuit_breaker_config: dict
263
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
264
- timeout_config: dict
265
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
198
+ Decorator prototype for all step decorators. This function gets specialized
199
+ and imported for all decorators types by _import_plugin_decorators().
266
200
  """
267
201
  ...
268
202
 
269
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
203
+ @typing.overload
204
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
205
+ ...
206
+
207
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
270
208
  """
271
- Specifies that this step is used to deploy an instance of the app.
272
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
273
-
274
-
275
- Parameters
276
- ----------
277
- app_port : int
278
- Number of GPUs to use.
279
- app_name : str
280
- Name of the app to deploy.
209
+ Decorator prototype for all step decorators. This function gets specialized
210
+ and imported for all decorators types by _import_plugin_decorators().
281
211
  """
282
212
  ...
283
213
 
@@ -317,132 +247,35 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
317
247
  ...
318
248
 
319
249
  @typing.overload
320
- 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]]]:
321
- """
322
- Specifies the PyPI packages for the step.
323
-
324
- Information in this decorator will augment any
325
- attributes set in the `@pyi_base` flow-level decorator. Hence,
326
- you can use `@pypi_base` to set packages required by all
327
- steps and use `@pypi` to specify step-specific overrides.
328
-
329
-
330
- Parameters
331
- ----------
332
- packages : Dict[str, str], default: {}
333
- Packages to use for this step. The key is the name of the package
334
- and the value is the version to use.
335
- python : str, optional, default: None
336
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
337
- that the version used will correspond to the version of the Python interpreter used to start the run.
338
- """
339
- ...
340
-
341
- @typing.overload
342
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
343
- ...
344
-
345
- @typing.overload
346
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
347
- ...
348
-
349
- 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):
350
- """
351
- Specifies the PyPI packages for the step.
352
-
353
- Information in this decorator will augment any
354
- attributes set in the `@pyi_base` flow-level decorator. Hence,
355
- you can use `@pypi_base` to set packages required by all
356
- steps and use `@pypi` to specify step-specific overrides.
357
-
358
-
359
- Parameters
360
- ----------
361
- packages : Dict[str, str], default: {}
362
- Packages to use for this step. The key is the name of the package
363
- and the value is the version to use.
364
- python : str, optional, default: None
365
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
366
- that the version used will correspond to the version of the Python interpreter used to start the run.
367
- """
368
- ...
369
-
370
- @typing.overload
371
- 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]]]:
250
+ 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]]]:
372
251
  """
373
- Specifies the resources needed when executing this step.
374
-
375
- Use `@resources` to specify the resource requirements
376
- independently of the specific compute layer (`@batch`, `@kubernetes`).
377
-
378
- You can choose the compute layer on the command line by executing e.g.
379
- ```
380
- python myflow.py run --with batch
381
- ```
382
- or
383
- ```
384
- python myflow.py run --with kubernetes
385
- ```
386
- which executes the flow on the desired system using the
387
- requirements specified in `@resources`.
252
+ Specifies environment variables to be set prior to the execution of a step.
388
253
 
389
254
 
390
255
  Parameters
391
256
  ----------
392
- cpu : int, default 1
393
- Number of CPUs required for this step.
394
- gpu : int, optional, default None
395
- Number of GPUs required for this step.
396
- disk : int, optional, default None
397
- Disk size (in MB) required for this step. Only applies on Kubernetes.
398
- memory : int, default 4096
399
- Memory size (in MB) required for this step.
400
- shared_memory : int, optional, default None
401
- The value for the size (in MiB) of the /dev/shm volume for this step.
402
- This parameter maps to the `--shm-size` option in Docker.
257
+ vars : Dict[str, str], default {}
258
+ Dictionary of environment variables to set.
403
259
  """
404
260
  ...
405
261
 
406
262
  @typing.overload
407
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
263
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
408
264
  ...
409
265
 
410
266
  @typing.overload
411
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
267
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
412
268
  ...
413
269
 
414
- 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):
270
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
415
271
  """
416
- Specifies the resources needed when executing this step.
417
-
418
- Use `@resources` to specify the resource requirements
419
- independently of the specific compute layer (`@batch`, `@kubernetes`).
420
-
421
- You can choose the compute layer on the command line by executing e.g.
422
- ```
423
- python myflow.py run --with batch
424
- ```
425
- or
426
- ```
427
- python myflow.py run --with kubernetes
428
- ```
429
- which executes the flow on the desired system using the
430
- requirements specified in `@resources`.
272
+ Specifies environment variables to be set prior to the execution of a step.
431
273
 
432
274
 
433
275
  Parameters
434
276
  ----------
435
- cpu : int, default 1
436
- Number of CPUs required for this step.
437
- gpu : int, optional, default None
438
- Number of GPUs required for this step.
439
- disk : int, optional, default None
440
- Disk size (in MB) required for this step. Only applies on Kubernetes.
441
- memory : int, default 4096
442
- Memory size (in MB) required for this step.
443
- shared_memory : int, optional, default None
444
- The value for the size (in MiB) of the /dev/shm volume for this step.
445
- This parameter maps to the `--shm-size` option in Docker.
277
+ vars : Dict[str, str], default {}
278
+ Dictionary of environment variables to set.
446
279
  """
447
280
  ...
448
281
 
@@ -506,62 +339,11 @@ def timeout(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
506
339
  ...
507
340
 
508
341
  @typing.overload
509
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
342
+ 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]]]:
510
343
  """
511
- Specifies that the step will success under all circumstances.
344
+ Creates a human-readable report, a Metaflow Card, after this step completes.
512
345
 
513
- The decorator will create an optional artifact, specified by `var`, which
514
- contains the exception raised. You can use it to detect the presence
515
- of errors, indicating that all happy-path artifacts produced by the step
516
- are missing.
517
-
518
-
519
- Parameters
520
- ----------
521
- var : str, optional, default None
522
- Name of the artifact in which to store the caught exception.
523
- If not specified, the exception is not stored.
524
- print_exception : bool, default True
525
- Determines whether or not the exception is printed to
526
- stdout when caught.
527
- """
528
- ...
529
-
530
- @typing.overload
531
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
532
- ...
533
-
534
- @typing.overload
535
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
536
- ...
537
-
538
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
539
- """
540
- Specifies that the step will success under all circumstances.
541
-
542
- The decorator will create an optional artifact, specified by `var`, which
543
- contains the exception raised. You can use it to detect the presence
544
- of errors, indicating that all happy-path artifacts produced by the step
545
- are missing.
546
-
547
-
548
- Parameters
549
- ----------
550
- var : str, optional, default None
551
- Name of the artifact in which to store the caught exception.
552
- If not specified, the exception is not stored.
553
- print_exception : bool, default True
554
- Determines whether or not the exception is printed to
555
- stdout when caught.
556
- """
557
- ...
558
-
559
- @typing.overload
560
- 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]]]:
561
- """
562
- Creates a human-readable report, a Metaflow Card, after this step completes.
563
-
564
- Note that you may add multiple `@card` decorators in a step with different parameters.
346
+ Note that you may add multiple `@card` decorators in a step with different parameters.
565
347
 
566
348
 
567
349
  Parameters
@@ -606,21 +388,61 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
606
388
  ...
607
389
 
608
390
  @typing.overload
609
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
391
+ def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
610
392
  """
611
- Decorator prototype for all step decorators. This function gets specialized
612
- and imported for all decorators types by _import_plugin_decorators().
393
+ Specifies the Conda environment for the step.
394
+
395
+ Information in this decorator will augment any
396
+ attributes set in the `@conda_base` flow-level decorator. Hence,
397
+ you can use `@conda_base` to set packages required by all
398
+ steps and use `@conda` to specify step-specific overrides.
399
+
400
+
401
+ Parameters
402
+ ----------
403
+ packages : Dict[str, str], default {}
404
+ Packages to use for this step. The key is the name of the package
405
+ and the value is the version to use.
406
+ libraries : Dict[str, str], default {}
407
+ Supported for backward compatibility. When used with packages, packages will take precedence.
408
+ python : str, optional, default None
409
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
410
+ that the version used will correspond to the version of the Python interpreter used to start the run.
411
+ disabled : bool, default False
412
+ If set to True, disables @conda.
613
413
  """
614
414
  ...
615
415
 
616
416
  @typing.overload
617
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
417
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
618
418
  ...
619
419
 
620
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
420
+ @typing.overload
421
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
422
+ ...
423
+
424
+ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
621
425
  """
622
- Decorator prototype for all step decorators. This function gets specialized
623
- and imported for all decorators types by _import_plugin_decorators().
426
+ Specifies the Conda environment for the step.
427
+
428
+ Information in this decorator will augment any
429
+ attributes set in the `@conda_base` flow-level decorator. Hence,
430
+ you can use `@conda_base` to set packages required by all
431
+ steps and use `@conda` to specify step-specific overrides.
432
+
433
+
434
+ Parameters
435
+ ----------
436
+ packages : Dict[str, str], default {}
437
+ Packages to use for this step. The key is the name of the package
438
+ and the value is the version to use.
439
+ libraries : Dict[str, str], default {}
440
+ Supported for backward compatibility. When used with packages, packages will take precedence.
441
+ python : str, optional, default None
442
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
443
+ that the version used will correspond to the version of the Python interpreter used to start the run.
444
+ disabled : bool, default False
445
+ If set to True, disables @conda.
624
446
  """
625
447
  ...
626
448
 
@@ -709,210 +531,371 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
709
531
  """
710
532
  ...
711
533
 
712
- 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]]]:
534
+ @typing.overload
535
+ 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]]]:
713
536
  """
714
- Specifies that this step should execute on DGX cloud.
537
+ Specifies the resources needed when executing this step.
538
+
539
+ Use `@resources` to specify the resource requirements
540
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
541
+
542
+ You can choose the compute layer on the command line by executing e.g.
543
+ ```
544
+ python myflow.py run --with batch
545
+ ```
546
+ or
547
+ ```
548
+ python myflow.py run --with kubernetes
549
+ ```
550
+ which executes the flow on the desired system using the
551
+ requirements specified in `@resources`.
715
552
 
716
553
 
717
554
  Parameters
718
555
  ----------
719
- gpu : int
720
- Number of GPUs to use.
721
- gpu_type : str
722
- Type of Nvidia GPU to use.
723
- queue_timeout : int
724
- Time to keep the job in NVCF's queue.
556
+ cpu : int, default 1
557
+ Number of CPUs required for this step.
558
+ gpu : int, optional, default None
559
+ Number of GPUs required for this step.
560
+ disk : int, optional, default None
561
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
562
+ memory : int, default 4096
563
+ Memory size (in MB) required for this step.
564
+ shared_memory : int, optional, default None
565
+ The value for the size (in MiB) of the /dev/shm volume for this step.
566
+ This parameter maps to the `--shm-size` option in Docker.
725
567
  """
726
568
  ...
727
569
 
728
570
  @typing.overload
729
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
730
- """
731
- Internal decorator to support Fast bakery
732
- """
571
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
733
572
  ...
734
573
 
735
574
  @typing.overload
736
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
737
- ...
738
-
739
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
740
- """
741
- Internal decorator to support Fast bakery
742
- """
575
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
743
576
  ...
744
577
 
745
- @typing.overload
746
- 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]]]:
578
+ 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):
747
579
  """
748
- Enables checkpointing for a step.
580
+ Specifies the resources needed when executing this step.
581
+
582
+ Use `@resources` to specify the resource requirements
583
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
749
584
 
585
+ You can choose the compute layer on the command line by executing e.g.
586
+ ```
587
+ python myflow.py run --with batch
588
+ ```
589
+ or
590
+ ```
591
+ python myflow.py run --with kubernetes
592
+ ```
593
+ which executes the flow on the desired system using the
594
+ requirements specified in `@resources`.
750
595
 
751
596
 
752
597
  Parameters
753
598
  ----------
754
- load_policy : str, default: "fresh"
755
- The policy for loading the checkpoint. The following policies are supported:
756
- - "eager": Loads the the latest available checkpoint within the namespace.
757
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
758
- will be loaded at the start of the task.
759
- - "none": Do not load any checkpoint
760
- - "fresh": Loads the lastest checkpoint created within the running Task.
761
- This mode helps loading checkpoints across various retry attempts of the same task.
762
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
763
- created within the task will be loaded when the task is retries execution on failure.
764
-
765
- temp_dir_root : str, default: None
766
- The root directory under which `current.checkpoint.directory` will be created.
599
+ cpu : int, default 1
600
+ Number of CPUs required for this step.
601
+ gpu : int, optional, default None
602
+ Number of GPUs required for this step.
603
+ disk : int, optional, default None
604
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
605
+ memory : int, default 4096
606
+ Memory size (in MB) required for this step.
607
+ shared_memory : int, optional, default None
608
+ The value for the size (in MiB) of the /dev/shm volume for this step.
609
+ This parameter maps to the `--shm-size` option in Docker.
767
610
  """
768
611
  ...
769
612
 
770
- @typing.overload
771
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
772
- ...
773
-
774
- @typing.overload
775
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
776
- ...
777
-
778
- 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):
613
+ 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]]]:
779
614
  """
780
- Enables checkpointing for a step.
615
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
616
+
617
+ User code call
618
+ --------------
619
+ @ollama(
620
+ models=[...],
621
+ ...
622
+ )
623
+
624
+ Valid backend options
625
+ ---------------------
626
+ - 'local': Run as a separate process on the local task machine.
627
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
628
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
781
629
 
630
+ Valid model options
631
+ -------------------
632
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
782
633
 
783
634
 
784
635
  Parameters
785
636
  ----------
786
- load_policy : str, default: "fresh"
787
- The policy for loading the checkpoint. The following policies are supported:
788
- - "eager": Loads the the latest available checkpoint within the namespace.
789
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
790
- will be loaded at the start of the task.
791
- - "none": Do not load any checkpoint
792
- - "fresh": Loads the lastest checkpoint created within the running Task.
793
- This mode helps loading checkpoints across various retry attempts of the same task.
794
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
795
- created within the task will be loaded when the task is retries execution on failure.
796
-
797
- temp_dir_root : str, default: None
798
- The root directory under which `current.checkpoint.directory` will be created.
637
+ models: list[str]
638
+ List of Ollama containers running models in sidecars.
639
+ backend: str
640
+ Determines where and how to run the Ollama process.
641
+ force_pull: bool
642
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
643
+ cache_update_policy: str
644
+ Cache update policy: "auto", "force", or "never".
645
+ force_cache_update: bool
646
+ Simple override for "force" cache update policy.
647
+ debug: bool
648
+ Whether to turn on verbose debugging logs.
649
+ circuit_breaker_config: dict
650
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
651
+ timeout_config: dict
652
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
799
653
  """
800
654
  ...
801
655
 
802
656
  @typing.overload
803
- def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
657
+ 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]]]:
804
658
  """
805
- Specifies the Conda environment for the step.
659
+ Specifies the PyPI packages for the step.
806
660
 
807
661
  Information in this decorator will augment any
808
- attributes set in the `@conda_base` flow-level decorator. Hence,
809
- you can use `@conda_base` to set packages required by all
810
- steps and use `@conda` to specify step-specific overrides.
662
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
663
+ you can use `@pypi_base` to set packages required by all
664
+ steps and use `@pypi` to specify step-specific overrides.
811
665
 
812
666
 
813
667
  Parameters
814
668
  ----------
815
- packages : Dict[str, str], default {}
669
+ packages : Dict[str, str], default: {}
816
670
  Packages to use for this step. The key is the name of the package
817
671
  and the value is the version to use.
818
- libraries : Dict[str, str], default {}
819
- Supported for backward compatibility. When used with packages, packages will take precedence.
820
- python : str, optional, default None
672
+ python : str, optional, default: None
821
673
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
822
674
  that the version used will correspond to the version of the Python interpreter used to start the run.
823
- disabled : bool, default False
824
- If set to True, disables @conda.
825
675
  """
826
676
  ...
827
677
 
828
678
  @typing.overload
829
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
679
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
830
680
  ...
831
681
 
832
682
  @typing.overload
833
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
683
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
834
684
  ...
835
685
 
836
- def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
686
+ 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):
837
687
  """
838
- Specifies the Conda environment for the step.
688
+ Specifies the PyPI packages for the step.
839
689
 
840
690
  Information in this decorator will augment any
841
- attributes set in the `@conda_base` flow-level decorator. Hence,
842
- you can use `@conda_base` to set packages required by all
843
- steps and use `@conda` to specify step-specific overrides.
691
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
692
+ you can use `@pypi_base` to set packages required by all
693
+ steps and use `@pypi` to specify step-specific overrides.
844
694
 
845
695
 
846
696
  Parameters
847
697
  ----------
848
- packages : Dict[str, str], default {}
698
+ packages : Dict[str, str], default: {}
849
699
  Packages to use for this step. The key is the name of the package
850
700
  and the value is the version to use.
851
- libraries : Dict[str, str], default {}
852
- Supported for backward compatibility. When used with packages, packages will take precedence.
853
- python : str, optional, default None
701
+ python : str, optional, default: None
854
702
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
855
703
  that the version used will correspond to the version of the Python interpreter used to start the run.
856
- disabled : bool, default False
857
- If set to True, disables @conda.
858
704
  """
859
705
  ...
860
706
 
861
707
  @typing.overload
862
- 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]]]:
708
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
863
709
  """
864
- Specifies environment variables to be set prior to the execution of a step.
710
+ Internal decorator to support Fast bakery
711
+ """
712
+ ...
713
+
714
+ @typing.overload
715
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
716
+ ...
717
+
718
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
719
+ """
720
+ Internal decorator to support Fast bakery
721
+ """
722
+ ...
723
+
724
+ @typing.overload
725
+ 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]]]:
726
+ """
727
+ Specifies the number of times the task corresponding
728
+ to a step needs to be retried.
729
+
730
+ This decorator is useful for handling transient errors, such as networking issues.
731
+ If your task contains operations that can't be retried safely, e.g. database updates,
732
+ it is advisable to annotate it with `@retry(times=0)`.
733
+
734
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
735
+ decorator will execute a no-op task after all retries have been exhausted,
736
+ ensuring that the flow execution can continue.
865
737
 
866
738
 
867
739
  Parameters
868
740
  ----------
869
- vars : Dict[str, str], default {}
870
- Dictionary of environment variables to set.
741
+ times : int, default 3
742
+ Number of times to retry this task.
743
+ minutes_between_retries : int, default 2
744
+ Number of minutes between retries.
871
745
  """
872
746
  ...
873
747
 
874
748
  @typing.overload
875
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
749
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
876
750
  ...
877
751
 
878
752
  @typing.overload
879
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
753
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
880
754
  ...
881
755
 
882
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
756
+ 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):
883
757
  """
884
- Specifies environment variables to be set prior to the execution of a step.
758
+ Specifies the number of times the task corresponding
759
+ to a step needs to be retried.
760
+
761
+ This decorator is useful for handling transient errors, such as networking issues.
762
+ If your task contains operations that can't be retried safely, e.g. database updates,
763
+ it is advisable to annotate it with `@retry(times=0)`.
764
+
765
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
766
+ decorator will execute a no-op task after all retries have been exhausted,
767
+ ensuring that the flow execution can continue.
885
768
 
886
769
 
887
770
  Parameters
888
771
  ----------
889
- vars : Dict[str, str], default {}
890
- Dictionary of environment variables to set.
772
+ times : int, default 3
773
+ Number of times to retry this task.
774
+ minutes_between_retries : int, default 2
775
+ Number of minutes between retries.
891
776
  """
892
777
  ...
893
778
 
894
- 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]]]:
779
+ @typing.overload
780
+ def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
895
781
  """
896
- Decorator that helps cache, version and store models/datasets from huggingface hub.
782
+ Specifies that the step will success under all circumstances.
783
+
784
+ The decorator will create an optional artifact, specified by `var`, which
785
+ contains the exception raised. You can use it to detect the presence
786
+ of errors, indicating that all happy-path artifacts produced by the step
787
+ are missing.
897
788
 
898
789
 
899
790
  Parameters
900
791
  ----------
901
- temp_dir_root : str, optional
902
- The root directory that will hold the temporary directory where objects will be downloaded.
792
+ var : str, optional, default None
793
+ Name of the artifact in which to store the caught exception.
794
+ If not specified, the exception is not stored.
795
+ print_exception : bool, default True
796
+ Determines whether or not the exception is printed to
797
+ stdout when caught.
798
+ """
799
+ ...
800
+
801
+ @typing.overload
802
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
803
+ ...
804
+
805
+ @typing.overload
806
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
807
+ ...
808
+
809
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
810
+ """
811
+ Specifies that the step will success under all circumstances.
903
812
 
904
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
905
- The list of repos (models/datasets) to load.
813
+ The decorator will create an optional artifact, specified by `var`, which
814
+ contains the exception raised. You can use it to detect the presence
815
+ of errors, indicating that all happy-path artifacts produced by the step
816
+ are missing.
906
817
 
907
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
908
818
 
909
- - If repo (model/dataset) is not found in the datastore:
910
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
911
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
912
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
819
+ Parameters
820
+ ----------
821
+ var : str, optional, default None
822
+ Name of the artifact in which to store the caught exception.
823
+ If not specified, the exception is not stored.
824
+ print_exception : bool, default True
825
+ Determines whether or not the exception is printed to
826
+ stdout when caught.
827
+ """
828
+ ...
829
+
830
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
831
+ """
832
+ Specifies that this step is used to deploy an instance of the app.
833
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
913
834
 
914
- - If repo is found in the datastore:
915
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
835
+
836
+ Parameters
837
+ ----------
838
+ app_port : int
839
+ Number of GPUs to use.
840
+ app_name : str
841
+ Name of the app to deploy.
842
+ """
843
+ ...
844
+
845
+ @typing.overload
846
+ 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]]]:
847
+ """
848
+ Enables checkpointing for a step.
849
+
850
+
851
+
852
+ Parameters
853
+ ----------
854
+ load_policy : str, default: "fresh"
855
+ The policy for loading the checkpoint. The following policies are supported:
856
+ - "eager": Loads the the latest available checkpoint within the namespace.
857
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
858
+ will be loaded at the start of the task.
859
+ - "none": Do not load any checkpoint
860
+ - "fresh": Loads the lastest checkpoint created within the running Task.
861
+ This mode helps loading checkpoints across various retry attempts of the same task.
862
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
863
+ created within the task will be loaded when the task is retries execution on failure.
864
+
865
+ temp_dir_root : str, default: None
866
+ The root directory under which `current.checkpoint.directory` will be created.
867
+ """
868
+ ...
869
+
870
+ @typing.overload
871
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
872
+ ...
873
+
874
+ @typing.overload
875
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
876
+ ...
877
+
878
+ 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):
879
+ """
880
+ Enables checkpointing for a step.
881
+
882
+
883
+
884
+ Parameters
885
+ ----------
886
+ load_policy : str, default: "fresh"
887
+ The policy for loading the checkpoint. The following policies are supported:
888
+ - "eager": Loads the the latest available checkpoint within the namespace.
889
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
890
+ will be loaded at the start of the task.
891
+ - "none": Do not load any checkpoint
892
+ - "fresh": Loads the lastest checkpoint created within the running Task.
893
+ This mode helps loading checkpoints across various retry attempts of the same task.
894
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
895
+ created within the task will be loaded when the task is retries execution on failure.
896
+
897
+ temp_dir_root : str, default: None
898
+ The root directory under which `current.checkpoint.directory` will be created.
916
899
  """
917
900
  ...
918
901
 
@@ -973,186 +956,23 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
973
956
  """
974
957
  ...
975
958
 
976
- 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]]:
959
+ 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]]]:
977
960
  """
978
- 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.
979
- 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.
961
+ Specifies that this step should execute on DGX cloud.
980
962
 
981
963
 
982
964
  Parameters
983
965
  ----------
984
- timeout : int
985
- Time, in seconds before the task times out and fails. (Default: 3600)
986
- poke_interval : int
987
- Time in seconds that the job should wait in between each try. (Default: 60)
988
- mode : str
989
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
990
- exponential_backoff : bool
991
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
992
- pool : str
993
- the slot pool this task should run in,
994
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
995
- soft_fail : bool
996
- Set to true to mark the task as SKIPPED on failure. (Default: False)
997
- name : str
998
- Name of the sensor on Airflow
999
- description : str
1000
- Description of sensor in the Airflow UI
1001
- external_dag_id : str
1002
- The dag_id that contains the task you want to wait for.
1003
- external_task_ids : List[str]
1004
- The list of task_ids that you want to wait for.
1005
- If None (default value) the sensor waits for the DAG. (Default: None)
1006
- allowed_states : List[str]
1007
- Iterable of allowed states, (Default: ['success'])
1008
- failed_states : List[str]
1009
- Iterable of failed or dis-allowed states. (Default: None)
1010
- execution_delta : datetime.timedelta
1011
- time difference with the previous execution to look at,
1012
- the default is the same logical date as the current task or DAG. (Default: None)
1013
- check_existence: bool
1014
- Set to True to check if the external task exists or check if
1015
- the DAG to wait for exists. (Default: True)
966
+ gpu : int
967
+ Number of GPUs to use.
968
+ gpu_type : str
969
+ Type of Nvidia GPU to use.
970
+ queue_timeout : int
971
+ Time to keep the job in NVCF's queue.
1016
972
  """
1017
973
  ...
1018
974
 
1019
- @typing.overload
1020
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1021
- """
1022
- Specifies the flow(s) that this flow depends on.
1023
-
1024
- ```
1025
- @trigger_on_finish(flow='FooFlow')
1026
- ```
1027
- or
1028
- ```
1029
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1030
- ```
1031
- This decorator respects the @project decorator and triggers the flow
1032
- when upstream runs within the same namespace complete successfully
1033
-
1034
- Additionally, you can specify project aware upstream flow dependencies
1035
- by specifying the fully qualified project_flow_name.
1036
- ```
1037
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1038
- ```
1039
- or
1040
- ```
1041
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1042
- ```
1043
-
1044
- You can also specify just the project or project branch (other values will be
1045
- inferred from the current project or project branch):
1046
- ```
1047
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1048
- ```
1049
-
1050
- Note that `branch` is typically one of:
1051
- - `prod`
1052
- - `user.bob`
1053
- - `test.my_experiment`
1054
- - `prod.staging`
1055
-
1056
-
1057
- Parameters
1058
- ----------
1059
- flow : Union[str, Dict[str, str]], optional, default None
1060
- Upstream flow dependency for this flow.
1061
- flows : List[Union[str, Dict[str, str]]], default []
1062
- Upstream flow dependencies for this flow.
1063
- options : Dict[str, Any], default {}
1064
- Backend-specific configuration for tuning eventing behavior.
1065
- """
1066
- ...
1067
-
1068
- @typing.overload
1069
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1070
- ...
1071
-
1072
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1073
- """
1074
- Specifies the flow(s) that this flow depends on.
1075
-
1076
- ```
1077
- @trigger_on_finish(flow='FooFlow')
1078
- ```
1079
- or
1080
- ```
1081
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1082
- ```
1083
- This decorator respects the @project decorator and triggers the flow
1084
- when upstream runs within the same namespace complete successfully
1085
-
1086
- Additionally, you can specify project aware upstream flow dependencies
1087
- by specifying the fully qualified project_flow_name.
1088
- ```
1089
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1090
- ```
1091
- or
1092
- ```
1093
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1094
- ```
1095
-
1096
- You can also specify just the project or project branch (other values will be
1097
- inferred from the current project or project branch):
1098
- ```
1099
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1100
- ```
1101
-
1102
- Note that `branch` is typically one of:
1103
- - `prod`
1104
- - `user.bob`
1105
- - `test.my_experiment`
1106
- - `prod.staging`
1107
-
1108
-
1109
- Parameters
1110
- ----------
1111
- flow : Union[str, Dict[str, str]], optional, default None
1112
- Upstream flow dependency for this flow.
1113
- flows : List[Union[str, Dict[str, str]]], default []
1114
- Upstream flow dependencies for this flow.
1115
- options : Dict[str, Any], default {}
1116
- Backend-specific configuration for tuning eventing behavior.
1117
- """
1118
- ...
1119
-
1120
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1121
- """
1122
- Specifies what flows belong to the same project.
1123
-
1124
- A project-specific namespace is created for all flows that
1125
- use the same `@project(name)`.
1126
-
1127
-
1128
- Parameters
1129
- ----------
1130
- name : str
1131
- Project name. Make sure that the name is unique amongst all
1132
- projects that use the same production scheduler. The name may
1133
- contain only lowercase alphanumeric characters and underscores.
1134
-
1135
- branch : Optional[str], default None
1136
- The branch to use. If not specified, the branch is set to
1137
- `user.<username>` unless `production` is set to `True`. This can
1138
- also be set on the command line using `--branch` as a top-level option.
1139
- It is an error to specify `branch` in the decorator and on the command line.
1140
-
1141
- production : bool, default False
1142
- Whether or not the branch is the production branch. This can also be set on the
1143
- command line using `--production` as a top-level option. It is an error to specify
1144
- `production` in the decorator and on the command line.
1145
- The project branch name will be:
1146
- - if `branch` is specified:
1147
- - if `production` is True: `prod.<branch>`
1148
- - if `production` is False: `test.<branch>`
1149
- - if `branch` is not specified:
1150
- - if `production` is True: `prod`
1151
- - if `production` is False: `user.<username>`
1152
- """
1153
- ...
1154
-
1155
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
975
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1156
976
  """
1157
977
  Allows setting external datastores to save data for the
1158
978
  `@checkpoint`/`@model`/`@huggingface_hub` decorators.
@@ -1251,18 +1071,305 @@ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None)
1251
1071
  Parameters:
1252
1072
  ----------
1253
1073
 
1254
- type: str
1255
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1074
+ type: str
1075
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1076
+
1077
+ config: dict or Callable
1078
+ Dictionary of configuration options for the datastore. The following keys are required:
1079
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1080
+ - example: 's3://bucket-name/path/to/root'
1081
+ - example: 'gs://bucket-name/path/to/root'
1082
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1083
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1084
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1085
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1086
+ """
1087
+ ...
1088
+
1089
+ @typing.overload
1090
+ 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]]:
1091
+ """
1092
+ Specifies the Conda environment for all steps of the flow.
1093
+
1094
+ Use `@conda_base` to set common libraries required by all
1095
+ steps and use `@conda` to specify step-specific additions.
1096
+
1097
+
1098
+ Parameters
1099
+ ----------
1100
+ packages : Dict[str, str], default {}
1101
+ Packages to use for this flow. The key is the name of the package
1102
+ and the value is the version to use.
1103
+ libraries : Dict[str, str], default {}
1104
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1105
+ python : str, optional, default None
1106
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1107
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1108
+ disabled : bool, default False
1109
+ If set to True, disables Conda.
1110
+ """
1111
+ ...
1112
+
1113
+ @typing.overload
1114
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1115
+ ...
1116
+
1117
+ 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):
1118
+ """
1119
+ Specifies the Conda environment for all steps of the flow.
1120
+
1121
+ Use `@conda_base` to set common libraries required by all
1122
+ steps and use `@conda` to specify step-specific additions.
1123
+
1124
+
1125
+ Parameters
1126
+ ----------
1127
+ packages : Dict[str, str], default {}
1128
+ Packages to use for this flow. The key is the name of the package
1129
+ and the value is the version to use.
1130
+ libraries : Dict[str, str], default {}
1131
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1132
+ python : str, optional, default None
1133
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1134
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1135
+ disabled : bool, default False
1136
+ If set to True, disables Conda.
1137
+ """
1138
+ ...
1139
+
1140
+ @typing.overload
1141
+ 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]]:
1142
+ """
1143
+ Specifies the times when the flow should be run when running on a
1144
+ production scheduler.
1145
+
1146
+
1147
+ Parameters
1148
+ ----------
1149
+ hourly : bool, default False
1150
+ Run the workflow hourly.
1151
+ daily : bool, default True
1152
+ Run the workflow daily.
1153
+ weekly : bool, default False
1154
+ Run the workflow weekly.
1155
+ cron : str, optional, default None
1156
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1157
+ specified by this expression.
1158
+ timezone : str, optional, default None
1159
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1160
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1161
+ """
1162
+ ...
1163
+
1164
+ @typing.overload
1165
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1166
+ ...
1167
+
1168
+ 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):
1169
+ """
1170
+ Specifies the times when the flow should be run when running on a
1171
+ production scheduler.
1172
+
1173
+
1174
+ Parameters
1175
+ ----------
1176
+ hourly : bool, default False
1177
+ Run the workflow hourly.
1178
+ daily : bool, default True
1179
+ Run the workflow daily.
1180
+ weekly : bool, default False
1181
+ Run the workflow weekly.
1182
+ cron : str, optional, default None
1183
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1184
+ specified by this expression.
1185
+ timezone : str, optional, default None
1186
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1187
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1188
+ """
1189
+ ...
1190
+
1191
+ @typing.overload
1192
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1193
+ """
1194
+ Specifies the PyPI packages for all steps of the flow.
1195
+
1196
+ Use `@pypi_base` to set common packages required by all
1197
+ steps and use `@pypi` to specify step-specific overrides.
1198
+
1199
+ Parameters
1200
+ ----------
1201
+ packages : Dict[str, str], default: {}
1202
+ Packages to use for this flow. The key is the name of the package
1203
+ and the value is the version to use.
1204
+ python : str, optional, default: None
1205
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1206
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1207
+ """
1208
+ ...
1209
+
1210
+ @typing.overload
1211
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1212
+ ...
1213
+
1214
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1215
+ """
1216
+ Specifies the PyPI packages for all steps of the flow.
1217
+
1218
+ Use `@pypi_base` to set common packages required by all
1219
+ steps and use `@pypi` to specify step-specific overrides.
1220
+
1221
+ Parameters
1222
+ ----------
1223
+ packages : Dict[str, str], default: {}
1224
+ Packages to use for this flow. The key is the name of the package
1225
+ and the value is the version to use.
1226
+ python : str, optional, default: None
1227
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1228
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1229
+ """
1230
+ ...
1231
+
1232
+ @typing.overload
1233
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1234
+ """
1235
+ Specifies the flow(s) that this flow depends on.
1236
+
1237
+ ```
1238
+ @trigger_on_finish(flow='FooFlow')
1239
+ ```
1240
+ or
1241
+ ```
1242
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1243
+ ```
1244
+ This decorator respects the @project decorator and triggers the flow
1245
+ when upstream runs within the same namespace complete successfully
1246
+
1247
+ Additionally, you can specify project aware upstream flow dependencies
1248
+ by specifying the fully qualified project_flow_name.
1249
+ ```
1250
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1251
+ ```
1252
+ or
1253
+ ```
1254
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1255
+ ```
1256
+
1257
+ You can also specify just the project or project branch (other values will be
1258
+ inferred from the current project or project branch):
1259
+ ```
1260
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1261
+ ```
1262
+
1263
+ Note that `branch` is typically one of:
1264
+ - `prod`
1265
+ - `user.bob`
1266
+ - `test.my_experiment`
1267
+ - `prod.staging`
1268
+
1269
+
1270
+ Parameters
1271
+ ----------
1272
+ flow : Union[str, Dict[str, str]], optional, default None
1273
+ Upstream flow dependency for this flow.
1274
+ flows : List[Union[str, Dict[str, str]]], default []
1275
+ Upstream flow dependencies for this flow.
1276
+ options : Dict[str, Any], default {}
1277
+ Backend-specific configuration for tuning eventing behavior.
1278
+ """
1279
+ ...
1280
+
1281
+ @typing.overload
1282
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1283
+ ...
1284
+
1285
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1286
+ """
1287
+ Specifies the flow(s) that this flow depends on.
1288
+
1289
+ ```
1290
+ @trigger_on_finish(flow='FooFlow')
1291
+ ```
1292
+ or
1293
+ ```
1294
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1295
+ ```
1296
+ This decorator respects the @project decorator and triggers the flow
1297
+ when upstream runs within the same namespace complete successfully
1298
+
1299
+ Additionally, you can specify project aware upstream flow dependencies
1300
+ by specifying the fully qualified project_flow_name.
1301
+ ```
1302
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1303
+ ```
1304
+ or
1305
+ ```
1306
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1307
+ ```
1308
+
1309
+ You can also specify just the project or project branch (other values will be
1310
+ inferred from the current project or project branch):
1311
+ ```
1312
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1313
+ ```
1314
+
1315
+ Note that `branch` is typically one of:
1316
+ - `prod`
1317
+ - `user.bob`
1318
+ - `test.my_experiment`
1319
+ - `prod.staging`
1320
+
1321
+
1322
+ Parameters
1323
+ ----------
1324
+ flow : Union[str, Dict[str, str]], optional, default None
1325
+ Upstream flow dependency for this flow.
1326
+ flows : List[Union[str, Dict[str, str]]], default []
1327
+ Upstream flow dependencies for this flow.
1328
+ options : Dict[str, Any], default {}
1329
+ Backend-specific configuration for tuning eventing behavior.
1330
+ """
1331
+ ...
1332
+
1333
+ 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]]:
1334
+ """
1335
+ 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)
1336
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1337
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1338
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1339
+ starts only after all sensors finish.
1340
+
1256
1341
 
1257
- config: dict or Callable
1258
- Dictionary of configuration options for the datastore. The following keys are required:
1259
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1260
- - example: 's3://bucket-name/path/to/root'
1261
- - example: 'gs://bucket-name/path/to/root'
1262
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1263
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1264
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1265
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1342
+ Parameters
1343
+ ----------
1344
+ timeout : int
1345
+ Time, in seconds before the task times out and fails. (Default: 3600)
1346
+ poke_interval : int
1347
+ Time in seconds that the job should wait in between each try. (Default: 60)
1348
+ mode : str
1349
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1350
+ exponential_backoff : bool
1351
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1352
+ pool : str
1353
+ the slot pool this task should run in,
1354
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1355
+ soft_fail : bool
1356
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1357
+ name : str
1358
+ Name of the sensor on Airflow
1359
+ description : str
1360
+ Description of sensor in the Airflow UI
1361
+ bucket_key : Union[str, List[str]]
1362
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1363
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1364
+ bucket_name : str
1365
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1366
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1367
+ wildcard_match : bool
1368
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1369
+ aws_conn_id : str
1370
+ a reference to the s3 connection on Airflow. (Default: None)
1371
+ verify : bool
1372
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1266
1373
  """
1267
1374
  ...
1268
1375
 
@@ -1359,115 +1466,45 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1359
1466
  """
1360
1467
  ...
1361
1468
 
1362
- @typing.overload
1363
- 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]]:
1364
- """
1365
- Specifies the times when the flow should be run when running on a
1366
- production scheduler.
1367
-
1368
-
1369
- Parameters
1370
- ----------
1371
- hourly : bool, default False
1372
- Run the workflow hourly.
1373
- daily : bool, default True
1374
- Run the workflow daily.
1375
- weekly : bool, default False
1376
- Run the workflow weekly.
1377
- cron : str, optional, default None
1378
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1379
- specified by this expression.
1380
- timezone : str, optional, default None
1381
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1382
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1383
- """
1384
- ...
1385
-
1386
- @typing.overload
1387
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1388
- ...
1389
-
1390
- 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):
1391
- """
1392
- Specifies the times when the flow should be run when running on a
1393
- production scheduler.
1394
-
1395
-
1396
- Parameters
1397
- ----------
1398
- hourly : bool, default False
1399
- Run the workflow hourly.
1400
- daily : bool, default True
1401
- Run the workflow daily.
1402
- weekly : bool, default False
1403
- Run the workflow weekly.
1404
- cron : str, optional, default None
1405
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1406
- specified by this expression.
1407
- timezone : str, optional, default None
1408
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1409
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1410
- """
1411
- ...
1412
-
1413
- @typing.overload
1414
- 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]]:
1469
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1415
1470
  """
1416
- Specifies the Conda environment for all steps of the flow.
1471
+ Specifies what flows belong to the same project.
1417
1472
 
1418
- Use `@conda_base` to set common libraries required by all
1419
- steps and use `@conda` to specify step-specific additions.
1473
+ A project-specific namespace is created for all flows that
1474
+ use the same `@project(name)`.
1420
1475
 
1421
1476
 
1422
1477
  Parameters
1423
1478
  ----------
1424
- packages : Dict[str, str], default {}
1425
- Packages to use for this flow. The key is the name of the package
1426
- and the value is the version to use.
1427
- libraries : Dict[str, str], default {}
1428
- Supported for backward compatibility. When used with packages, packages will take precedence.
1429
- python : str, optional, default None
1430
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1431
- that the version used will correspond to the version of the Python interpreter used to start the run.
1432
- disabled : bool, default False
1433
- If set to True, disables Conda.
1434
- """
1435
- ...
1436
-
1437
- @typing.overload
1438
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1439
- ...
1440
-
1441
- 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):
1442
- """
1443
- Specifies the Conda environment for all steps of the flow.
1444
-
1445
- Use `@conda_base` to set common libraries required by all
1446
- steps and use `@conda` to specify step-specific additions.
1479
+ name : str
1480
+ Project name. Make sure that the name is unique amongst all
1481
+ projects that use the same production scheduler. The name may
1482
+ contain only lowercase alphanumeric characters and underscores.
1447
1483
 
1484
+ branch : Optional[str], default None
1485
+ The branch to use. If not specified, the branch is set to
1486
+ `user.<username>` unless `production` is set to `True`. This can
1487
+ also be set on the command line using `--branch` as a top-level option.
1488
+ It is an error to specify `branch` in the decorator and on the command line.
1448
1489
 
1449
- Parameters
1450
- ----------
1451
- packages : Dict[str, str], default {}
1452
- Packages to use for this flow. The key is the name of the package
1453
- and the value is the version to use.
1454
- libraries : Dict[str, str], default {}
1455
- Supported for backward compatibility. When used with packages, packages will take precedence.
1456
- python : str, optional, default None
1457
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1458
- that the version used will correspond to the version of the Python interpreter used to start the run.
1459
- disabled : bool, default False
1460
- If set to True, disables Conda.
1490
+ production : bool, default False
1491
+ Whether or not the branch is the production branch. This can also be set on the
1492
+ command line using `--production` as a top-level option. It is an error to specify
1493
+ `production` in the decorator and on the command line.
1494
+ The project branch name will be:
1495
+ - if `branch` is specified:
1496
+ - if `production` is True: `prod.<branch>`
1497
+ - if `production` is False: `test.<branch>`
1498
+ - if `branch` is not specified:
1499
+ - if `production` is True: `prod`
1500
+ - if `production` is False: `user.<username>`
1461
1501
  """
1462
1502
  ...
1463
1503
 
1464
- 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]]:
1504
+ 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]]:
1465
1505
  """
1466
- 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)
1467
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1468
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1469
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1470
- starts only after all sensors finish.
1506
+ 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.
1507
+ 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.
1471
1508
 
1472
1509
 
1473
1510
  Parameters
@@ -1489,59 +1526,21 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1489
1526
  Name of the sensor on Airflow
1490
1527
  description : str
1491
1528
  Description of sensor in the Airflow UI
1492
- bucket_key : Union[str, List[str]]
1493
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1494
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1495
- bucket_name : str
1496
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1497
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1498
- wildcard_match : bool
1499
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1500
- aws_conn_id : str
1501
- a reference to the s3 connection on Airflow. (Default: None)
1502
- verify : bool
1503
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1504
- """
1505
- ...
1506
-
1507
- @typing.overload
1508
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1509
- """
1510
- Specifies the PyPI packages for all steps of the flow.
1511
-
1512
- Use `@pypi_base` to set common packages required by all
1513
- steps and use `@pypi` to specify step-specific overrides.
1514
-
1515
- Parameters
1516
- ----------
1517
- packages : Dict[str, str], default: {}
1518
- Packages to use for this flow. The key is the name of the package
1519
- and the value is the version to use.
1520
- python : str, optional, default: None
1521
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1522
- that the version used will correspond to the version of the Python interpreter used to start the run.
1523
- """
1524
- ...
1525
-
1526
- @typing.overload
1527
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1528
- ...
1529
-
1530
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1531
- """
1532
- Specifies the PyPI packages for all steps of the flow.
1533
-
1534
- Use `@pypi_base` to set common packages required by all
1535
- steps and use `@pypi` to specify step-specific overrides.
1536
-
1537
- Parameters
1538
- ----------
1539
- packages : Dict[str, str], default: {}
1540
- Packages to use for this flow. The key is the name of the package
1541
- and the value is the version to use.
1542
- python : str, optional, default: None
1543
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1544
- that the version used will correspond to the version of the Python interpreter used to start the run.
1529
+ external_dag_id : str
1530
+ The dag_id that contains the task you want to wait for.
1531
+ external_task_ids : List[str]
1532
+ The list of task_ids that you want to wait for.
1533
+ If None (default value) the sensor waits for the DAG. (Default: None)
1534
+ allowed_states : List[str]
1535
+ Iterable of allowed states, (Default: ['success'])
1536
+ failed_states : List[str]
1537
+ Iterable of failed or dis-allowed states. (Default: None)
1538
+ execution_delta : datetime.timedelta
1539
+ time difference with the previous execution to look at,
1540
+ the default is the same logical date as the current task or DAG. (Default: None)
1541
+ check_existence: bool
1542
+ Set to True to check if the external task exists or check if
1543
+ the DAG to wait for exists. (Default: True)
1545
1544
  """
1546
1545
  ...
1547
1546