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