ob-metaflow-stubs 6.0.3.180rc2__py2.py3-none-any.whl → 6.0.3.180rc3__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (220) hide show
  1. metaflow-stubs/__init__.pyi +626 -626
  2. metaflow-stubs/cards.pyi +1 -1
  3. metaflow-stubs/cli.pyi +1 -1
  4. metaflow-stubs/cli_components/__init__.pyi +1 -1
  5. metaflow-stubs/cli_components/utils.pyi +1 -1
  6. metaflow-stubs/client/__init__.pyi +1 -1
  7. metaflow-stubs/client/core.pyi +4 -4
  8. metaflow-stubs/client/filecache.pyi +1 -1
  9. metaflow-stubs/events.pyi +1 -1
  10. metaflow-stubs/exception.pyi +1 -1
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/info_file.pyi +1 -1
  15. metaflow-stubs/metadata_provider/__init__.pyi +1 -1
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +1 -1
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +1 -1
  19. metaflow-stubs/metaflow_config.pyi +1 -1
  20. metaflow-stubs/metaflow_current.pyi +35 -35
  21. metaflow-stubs/metaflow_git.pyi +1 -1
  22. metaflow-stubs/mf_extensions/__init__.pyi +1 -1
  23. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +1 -1
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +1 -1
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +1 -1
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +1 -1
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +1 -1
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +1 -1
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +1 -1
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +3 -3
  33. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +1 -1
  34. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/checkpoint_storage.pyi +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 +2 -2
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +3 -3
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +1 -1
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +1 -1
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +1 -1
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +1 -1
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/context.pyi +1 -1
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +1 -1
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/decorator.pyi +1 -1
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +1 -1
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +1 -1
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +1 -1
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +1 -1
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +1 -1
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +1 -1
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +3 -3
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +1 -1
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +1 -1
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +1 -1
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +1 -1
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +1 -1
  61. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +1 -1
  62. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +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/fast_bakery/__init__.pyi +1 -1
  75. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/baker.pyi +2 -2
  76. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/docker_environment.pyi +1 -1
  77. metaflow-stubs/mf_extensions/outerbounds/plugins/fast_bakery/fast_bakery.pyi +1 -1
  78. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/__init__.pyi +1 -1
  79. metaflow-stubs/mf_extensions/outerbounds/plugins/kubernetes/pod_killer.pyi +1 -1
  80. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +1 -1
  81. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/constants.pyi +1 -1
  82. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/exceptions.pyi +1 -1
  83. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +1 -1
  84. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/status_card.pyi +1 -1
  85. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +1 -1
  86. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +1 -1
  87. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +1 -1
  88. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +1 -1
  89. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +1 -1
  90. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +1 -1
  91. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +1 -1
  92. metaflow-stubs/multicore_utils.pyi +1 -1
  93. metaflow-stubs/ob_internal.pyi +1 -1
  94. metaflow-stubs/parameters.pyi +2 -2
  95. metaflow-stubs/plugins/__init__.pyi +10 -10
  96. metaflow-stubs/plugins/airflow/__init__.pyi +1 -1
  97. metaflow-stubs/plugins/airflow/airflow_utils.pyi +1 -1
  98. metaflow-stubs/plugins/airflow/exception.pyi +1 -1
  99. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +1 -1
  100. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +1 -1
  101. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +1 -1
  102. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +1 -1
  103. metaflow-stubs/plugins/argo/__init__.pyi +1 -1
  104. metaflow-stubs/plugins/argo/argo_client.pyi +1 -1
  105. metaflow-stubs/plugins/argo/argo_events.pyi +1 -1
  106. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  107. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  108. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  109. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  110. metaflow-stubs/plugins/aws/__init__.pyi +1 -1
  111. metaflow-stubs/plugins/aws/aws_client.pyi +1 -1
  112. metaflow-stubs/plugins/aws/aws_utils.pyi +1 -1
  113. metaflow-stubs/plugins/aws/batch/__init__.pyi +1 -1
  114. metaflow-stubs/plugins/aws/batch/batch.pyi +1 -1
  115. metaflow-stubs/plugins/aws/batch/batch_client.pyi +1 -1
  116. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +1 -1
  117. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +1 -1
  118. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +2 -2
  119. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +1 -1
  120. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +1 -1
  121. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +1 -1
  122. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +1 -1
  123. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +1 -1
  124. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  125. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  126. metaflow-stubs/plugins/azure/__init__.pyi +1 -1
  127. metaflow-stubs/plugins/azure/azure_credential.pyi +1 -1
  128. metaflow-stubs/plugins/azure/azure_exceptions.pyi +1 -1
  129. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +2 -2
  130. metaflow-stubs/plugins/azure/azure_utils.pyi +1 -1
  131. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +1 -1
  132. metaflow-stubs/plugins/azure/includefile_support.pyi +1 -1
  133. metaflow-stubs/plugins/cards/__init__.pyi +1 -1
  134. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  135. metaflow-stubs/plugins/cards/card_creator.pyi +1 -1
  136. metaflow-stubs/plugins/cards/card_datastore.pyi +1 -1
  137. metaflow-stubs/plugins/cards/card_decorator.pyi +1 -1
  138. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +1 -1
  139. metaflow-stubs/plugins/cards/card_modules/basic.pyi +1 -1
  140. metaflow-stubs/plugins/cards/card_modules/card.pyi +1 -1
  141. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  142. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +1 -1
  143. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +1 -1
  144. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +1 -1
  145. metaflow-stubs/plugins/cards/card_resolver.pyi +1 -1
  146. metaflow-stubs/plugins/cards/component_serializer.pyi +1 -1
  147. metaflow-stubs/plugins/cards/exception.pyi +1 -1
  148. metaflow-stubs/plugins/catch_decorator.pyi +1 -1
  149. metaflow-stubs/plugins/datatools/__init__.pyi +1 -1
  150. metaflow-stubs/plugins/datatools/local.pyi +1 -1
  151. metaflow-stubs/plugins/datatools/s3/__init__.pyi +1 -1
  152. metaflow-stubs/plugins/datatools/s3/s3.pyi +2 -2
  153. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +1 -1
  154. metaflow-stubs/plugins/datatools/s3/s3util.pyi +1 -1
  155. metaflow-stubs/plugins/debug_logger.pyi +1 -1
  156. metaflow-stubs/plugins/debug_monitor.pyi +1 -1
  157. metaflow-stubs/plugins/environment_decorator.pyi +1 -1
  158. metaflow-stubs/plugins/events_decorator.pyi +1 -1
  159. metaflow-stubs/plugins/frameworks/__init__.pyi +1 -1
  160. metaflow-stubs/plugins/frameworks/pytorch.pyi +1 -1
  161. metaflow-stubs/plugins/gcp/__init__.pyi +1 -1
  162. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +2 -2
  163. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +1 -1
  164. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +1 -1
  165. metaflow-stubs/plugins/gcp/gs_utils.pyi +1 -1
  166. metaflow-stubs/plugins/gcp/includefile_support.pyi +1 -1
  167. metaflow-stubs/plugins/kubernetes/__init__.pyi +1 -1
  168. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  169. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +1 -1
  170. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +1 -1
  171. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +1 -1
  172. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +1 -1
  173. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +1 -1
  174. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  175. metaflow-stubs/plugins/parallel_decorator.pyi +1 -1
  176. metaflow-stubs/plugins/perimeters.pyi +1 -1
  177. metaflow-stubs/plugins/project_decorator.pyi +1 -1
  178. metaflow-stubs/plugins/pypi/__init__.pyi +1 -1
  179. metaflow-stubs/plugins/pypi/conda_decorator.pyi +1 -1
  180. metaflow-stubs/plugins/pypi/conda_environment.pyi +4 -4
  181. metaflow-stubs/plugins/pypi/parsers.pyi +1 -1
  182. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +1 -1
  183. metaflow-stubs/plugins/pypi/pypi_environment.pyi +1 -1
  184. metaflow-stubs/plugins/pypi/utils.pyi +1 -1
  185. metaflow-stubs/plugins/resources_decorator.pyi +1 -1
  186. metaflow-stubs/plugins/retry_decorator.pyi +1 -1
  187. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  188. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +2 -2
  189. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +1 -1
  190. metaflow-stubs/plugins/snowflake/__init__.pyi +1 -1
  191. metaflow-stubs/plugins/storage_executor.pyi +1 -1
  192. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  193. metaflow-stubs/plugins/timeout_decorator.pyi +1 -1
  194. metaflow-stubs/plugins/torchtune/__init__.pyi +1 -1
  195. metaflow-stubs/plugins/uv/__init__.pyi +1 -1
  196. metaflow-stubs/plugins/uv/uv_environment.pyi +1 -1
  197. metaflow-stubs/profilers/__init__.pyi +1 -1
  198. metaflow-stubs/pylint_wrapper.pyi +1 -1
  199. metaflow-stubs/runner/__init__.pyi +1 -1
  200. metaflow-stubs/runner/deployer.pyi +30 -30
  201. metaflow-stubs/runner/deployer_impl.pyi +1 -1
  202. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  203. metaflow-stubs/runner/nbdeploy.pyi +1 -1
  204. metaflow-stubs/runner/nbrun.pyi +1 -1
  205. metaflow-stubs/runner/subprocess_manager.pyi +1 -1
  206. metaflow-stubs/runner/utils.pyi +3 -3
  207. metaflow-stubs/system/__init__.pyi +1 -1
  208. metaflow-stubs/system/system_logger.pyi +2 -2
  209. metaflow-stubs/system/system_monitor.pyi +1 -1
  210. metaflow-stubs/tagging_util.pyi +1 -1
  211. metaflow-stubs/tuple_util.pyi +1 -1
  212. metaflow-stubs/user_configs/__init__.pyi +1 -1
  213. metaflow-stubs/user_configs/config_decorators.pyi +4 -4
  214. metaflow-stubs/user_configs/config_options.pyi +2 -2
  215. metaflow-stubs/user_configs/config_parameters.pyi +5 -5
  216. {ob_metaflow_stubs-6.0.3.180rc2.dist-info → ob_metaflow_stubs-6.0.3.180rc3.dist-info}/METADATA +1 -1
  217. ob_metaflow_stubs-6.0.3.180rc3.dist-info/RECORD +220 -0
  218. ob_metaflow_stubs-6.0.3.180rc2.dist-info/RECORD +0 -220
  219. {ob_metaflow_stubs-6.0.3.180rc2.dist-info → ob_metaflow_stubs-6.0.3.180rc3.dist-info}/WHEEL +0 -0
  220. {ob_metaflow_stubs-6.0.3.180rc2.dist-info → ob_metaflow_stubs-6.0.3.180rc3.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-17T09:53:19.248370 #
4
+ # Generated on 2025-06-17T10:52:54.470733 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import datetime
12
11
  import typing
12
+ import datetime
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,17 +35,17 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import tuple_util as tuple_util
39
38
  from . import cards as cards
40
- from . import events as events
39
+ from . import tuple_util as tuple_util
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
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
48
47
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
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
@@ -155,70 +155,97 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
155
155
  """
156
156
  ...
157
157
 
158
- @typing.overload
159
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
158
+ 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]]]:
160
159
  """
161
- Specifies environment variables to be set prior to the execution of a step.
160
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
161
+
162
+ User code call
163
+ --------------
164
+ @ollama(
165
+ models=[...],
166
+ ...
167
+ )
168
+
169
+ Valid backend options
170
+ ---------------------
171
+ - 'local': Run as a separate process on the local task machine.
172
+ - (TODO) 'managed': Outerbounds hosts and selects compute provider.
173
+ - (TODO) 'remote': Spin up separate instance to serve Ollama models.
174
+
175
+ Valid model options
176
+ -------------------
177
+ Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
162
178
 
163
179
 
164
180
  Parameters
165
181
  ----------
166
- vars : Dict[str, str], default {}
167
- Dictionary of environment variables to set.
182
+ models: list[str]
183
+ List of Ollama containers running models in sidecars.
184
+ backend: str
185
+ Determines where and how to run the Ollama process.
186
+ force_pull: bool
187
+ Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
188
+ cache_update_policy: str
189
+ Cache update policy: "auto", "force", or "never".
190
+ force_cache_update: bool
191
+ Simple override for "force" cache update policy.
192
+ debug: bool
193
+ Whether to turn on verbose debugging logs.
194
+ circuit_breaker_config: dict
195
+ Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
196
+ timeout_config: dict
197
+ Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
168
198
  """
169
199
  ...
170
200
 
171
201
  @typing.overload
172
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
173
- ...
174
-
175
- @typing.overload
176
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
177
- ...
178
-
179
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
202
+ 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]]]:
180
203
  """
181
- Specifies environment variables to be set prior to the execution of a step.
204
+ Specifies that the step will success under all circumstances.
205
+
206
+ The decorator will create an optional artifact, specified by `var`, which
207
+ contains the exception raised. You can use it to detect the presence
208
+ of errors, indicating that all happy-path artifacts produced by the step
209
+ are missing.
182
210
 
183
211
 
184
212
  Parameters
185
213
  ----------
186
- vars : Dict[str, str], default {}
187
- Dictionary of environment variables to set.
214
+ var : str, optional, default None
215
+ Name of the artifact in which to store the caught exception.
216
+ If not specified, the exception is not stored.
217
+ print_exception : bool, default True
218
+ Determines whether or not the exception is printed to
219
+ stdout when caught.
188
220
  """
189
221
  ...
190
222
 
191
223
  @typing.overload
192
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
193
- """
194
- Decorator prototype for all step decorators. This function gets specialized
195
- and imported for all decorators types by _import_plugin_decorators().
196
- """
224
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
197
225
  ...
198
226
 
199
227
  @typing.overload
200
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
201
- ...
202
-
203
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
204
- """
205
- Decorator prototype for all step decorators. This function gets specialized
206
- and imported for all decorators types by _import_plugin_decorators().
207
- """
228
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
208
229
  ...
209
230
 
210
- def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
231
+ 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):
211
232
  """
212
- Specifies that this step is used to deploy an instance of the app.
213
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
233
+ Specifies that the step will success under all circumstances.
234
+
235
+ The decorator will create an optional artifact, specified by `var`, which
236
+ contains the exception raised. You can use it to detect the presence
237
+ of errors, indicating that all happy-path artifacts produced by the step
238
+ are missing.
214
239
 
215
240
 
216
241
  Parameters
217
242
  ----------
218
- app_port : int
219
- Number of GPUs to use.
220
- app_name : str
221
- Name of the app to deploy.
243
+ var : str, optional, default None
244
+ Name of the artifact in which to store the caught exception.
245
+ If not specified, the exception is not stored.
246
+ print_exception : bool, default True
247
+ Determines whether or not the exception is printed to
248
+ stdout when caught.
222
249
  """
223
250
  ...
224
251
 
@@ -280,157 +307,164 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
280
307
  ...
281
308
 
282
309
  @typing.overload
283
- 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]]]:
310
+ 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]]]:
284
311
  """
285
- Enables checkpointing for a step.
286
-
312
+ Specifies environment variables to be set prior to the execution of a step.
287
313
 
288
314
 
289
315
  Parameters
290
316
  ----------
291
- load_policy : str, default: "fresh"
292
- The policy for loading the checkpoint. The following policies are supported:
293
- - "eager": Loads the the latest available checkpoint within the namespace.
294
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
295
- will be loaded at the start of the task.
296
- - "none": Do not load any checkpoint
297
- - "fresh": Loads the lastest checkpoint created within the running Task.
298
- This mode helps loading checkpoints across various retry attempts of the same task.
299
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
300
- created within the task will be loaded when the task is retries execution on failure.
301
-
302
- temp_dir_root : str, default: None
303
- The root directory under which `current.checkpoint.directory` will be created.
317
+ vars : Dict[str, str], default {}
318
+ Dictionary of environment variables to set.
304
319
  """
305
320
  ...
306
321
 
307
322
  @typing.overload
308
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
323
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
309
324
  ...
310
325
 
311
326
  @typing.overload
312
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
327
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
313
328
  ...
314
329
 
315
- 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):
330
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
316
331
  """
317
- Enables checkpointing for a step.
318
-
332
+ Specifies environment variables to be set prior to the execution of a step.
319
333
 
320
334
 
321
335
  Parameters
322
336
  ----------
323
- load_policy : str, default: "fresh"
324
- The policy for loading the checkpoint. The following policies are supported:
325
- - "eager": Loads the the latest available checkpoint within the namespace.
326
- With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
327
- will be loaded at the start of the task.
328
- - "none": Do not load any checkpoint
329
- - "fresh": Loads the lastest checkpoint created within the running Task.
330
- This mode helps loading checkpoints across various retry attempts of the same task.
331
- With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
332
- created within the task will be loaded when the task is retries execution on failure.
333
-
334
- temp_dir_root : str, default: None
335
- The root directory under which `current.checkpoint.directory` will be created.
337
+ vars : Dict[str, str], default {}
338
+ Dictionary of environment variables to set.
336
339
  """
337
340
  ...
338
341
 
339
- 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]]]:
342
+ @typing.overload
343
+ 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]]]:
340
344
  """
341
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
342
-
343
- User code call
344
- --------------
345
- @ollama(
346
- models=[...],
347
- ...
348
- )
349
-
350
- Valid backend options
351
- ---------------------
352
- - 'local': Run as a separate process on the local task machine.
353
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
354
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
345
+ Specifies the PyPI packages for the step.
355
346
 
356
- Valid model options
357
- -------------------
358
- Any model here https://ollama.com/search, e.g. 'llama3.2', 'llama3.3'
347
+ Information in this decorator will augment any
348
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
349
+ you can use `@pypi_base` to set packages required by all
350
+ steps and use `@pypi` to specify step-specific overrides.
359
351
 
360
352
 
361
353
  Parameters
362
354
  ----------
363
- models: list[str]
364
- List of Ollama containers running models in sidecars.
365
- backend: str
366
- Determines where and how to run the Ollama process.
367
- force_pull: bool
368
- Whether to run `ollama pull` no matter what, or first check the remote cache in Metaflow datastore for this model key.
369
- cache_update_policy: str
370
- Cache update policy: "auto", "force", or "never".
371
- force_cache_update: bool
372
- Simple override for "force" cache update policy.
373
- debug: bool
374
- Whether to turn on verbose debugging logs.
375
- circuit_breaker_config: dict
376
- Configuration for circuit breaker protection. Keys: failure_threshold, recovery_timeout, reset_timeout.
377
- timeout_config: dict
378
- Configuration for various operation timeouts. Keys: pull, stop, health_check, install, server_startup.
355
+ packages : Dict[str, str], default: {}
356
+ Packages to use for this step. The key is the name of the package
357
+ and the value is the version to use.
358
+ python : str, optional, default: None
359
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
360
+ that the version used will correspond to the version of the Python interpreter used to start the run.
379
361
  """
380
362
  ...
381
363
 
382
- 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]]]:
364
+ @typing.overload
365
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
366
+ ...
367
+
368
+ @typing.overload
369
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
370
+ ...
371
+
372
+ 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):
383
373
  """
384
- This decorator is used to run vllm APIs as Metaflow task sidecars.
385
-
386
- User code call
387
- --------------
388
- @vllm(
389
- model="...",
390
- ...
391
- )
392
-
393
- Valid backend options
394
- ---------------------
395
- - 'local': Run as a separate process on the local task machine.
396
-
397
- Valid model options
398
- -------------------
399
- Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
374
+ Specifies the PyPI packages for the step.
400
375
 
401
- NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
402
- If you need multiple models, you must create multiple @vllm decorators.
376
+ Information in this decorator will augment any
377
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
378
+ you can use `@pypi_base` to set packages required by all
379
+ steps and use `@pypi` to specify step-specific overrides.
403
380
 
404
381
 
405
382
  Parameters
406
383
  ----------
407
- model: str
408
- HuggingFace model identifier to be served by vLLM.
409
- backend: str
410
- Determines where and how to run the vLLM process.
411
- debug: bool
412
- Whether to turn on verbose debugging logs.
413
- kwargs : Any
414
- Any other keyword arguments are passed directly to the vLLM engine.
415
- This allows for flexible configuration of vLLM server settings.
416
- For example, `tensor_parallel_size=2`.
384
+ packages : Dict[str, str], default: {}
385
+ Packages to use for this step. The key is the name of the package
386
+ and the value is the version to use.
387
+ python : str, optional, default: None
388
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
389
+ that the version used will correspond to the version of the Python interpreter used to start the run.
417
390
  """
418
391
  ...
419
392
 
420
393
  @typing.overload
421
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
422
- """
423
- Internal decorator to support Fast bakery
394
+ 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]]]:
424
395
  """
425
- ...
426
-
427
- @typing.overload
428
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
429
- ...
396
+ Specifies a timeout for your step.
397
+
398
+ This decorator is useful if this step may hang indefinitely.
399
+
400
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
401
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
402
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
403
+
404
+ Note that all the values specified in parameters are added together so if you specify
405
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
406
+
407
+
408
+ Parameters
409
+ ----------
410
+ seconds : int, default 0
411
+ Number of seconds to wait prior to timing out.
412
+ minutes : int, default 0
413
+ Number of minutes to wait prior to timing out.
414
+ hours : int, default 0
415
+ Number of hours to wait prior to timing out.
416
+ """
417
+ ...
430
418
 
431
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
419
+ @typing.overload
420
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
421
+ ...
422
+
423
+ @typing.overload
424
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
425
+ ...
426
+
427
+ 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):
432
428
  """
433
- Internal decorator to support Fast bakery
429
+ Specifies a timeout for your step.
430
+
431
+ This decorator is useful if this step may hang indefinitely.
432
+
433
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
434
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
435
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
436
+
437
+ Note that all the values specified in parameters are added together so if you specify
438
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
439
+
440
+
441
+ Parameters
442
+ ----------
443
+ seconds : int, default 0
444
+ Number of seconds to wait prior to timing out.
445
+ minutes : int, default 0
446
+ Number of minutes to wait prior to timing out.
447
+ hours : int, default 0
448
+ Number of hours to wait prior to timing out.
449
+ """
450
+ ...
451
+
452
+ @typing.overload
453
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
454
+ """
455
+ Decorator prototype for all step decorators. This function gets specialized
456
+ and imported for all decorators types by _import_plugin_decorators().
457
+ """
458
+ ...
459
+
460
+ @typing.overload
461
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
462
+ ...
463
+
464
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
465
+ """
466
+ Decorator prototype for all step decorators. This function gets specialized
467
+ and imported for all decorators types by _import_plugin_decorators().
434
468
  """
435
469
  ...
436
470
 
@@ -470,78 +504,118 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
470
504
  ...
471
505
 
472
506
  @typing.overload
473
- 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]]]:
507
+ 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]]]:
474
508
  """
475
- Specifies that the step will success under all circumstances.
509
+ Specifies the Conda environment for the step.
476
510
 
477
- The decorator will create an optional artifact, specified by `var`, which
478
- contains the exception raised. You can use it to detect the presence
479
- of errors, indicating that all happy-path artifacts produced by the step
480
- are missing.
511
+ Information in this decorator will augment any
512
+ attributes set in the `@conda_base` flow-level decorator. Hence,
513
+ you can use `@conda_base` to set packages required by all
514
+ steps and use `@conda` to specify step-specific overrides.
481
515
 
482
516
 
483
517
  Parameters
484
518
  ----------
485
- var : str, optional, default None
486
- Name of the artifact in which to store the caught exception.
487
- If not specified, the exception is not stored.
488
- print_exception : bool, default True
489
- Determines whether or not the exception is printed to
490
- stdout when caught.
519
+ packages : Dict[str, str], default {}
520
+ Packages to use for this step. The key is the name of the package
521
+ and the value is the version to use.
522
+ libraries : Dict[str, str], default {}
523
+ Supported for backward compatibility. When used with packages, packages will take precedence.
524
+ python : str, optional, default None
525
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
526
+ that the version used will correspond to the version of the Python interpreter used to start the run.
527
+ disabled : bool, default False
528
+ If set to True, disables @conda.
491
529
  """
492
530
  ...
493
531
 
494
532
  @typing.overload
495
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
533
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
496
534
  ...
497
535
 
498
536
  @typing.overload
499
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
537
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
500
538
  ...
501
539
 
502
- 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):
540
+ 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):
503
541
  """
504
- Specifies that the step will success under all circumstances.
542
+ Specifies the Conda environment for the step.
505
543
 
506
- The decorator will create an optional artifact, specified by `var`, which
507
- contains the exception raised. You can use it to detect the presence
508
- of errors, indicating that all happy-path artifacts produced by the step
509
- are missing.
544
+ Information in this decorator will augment any
545
+ attributes set in the `@conda_base` flow-level decorator. Hence,
546
+ you can use `@conda_base` to set packages required by all
547
+ steps and use `@conda` to specify step-specific overrides.
510
548
 
511
549
 
512
550
  Parameters
513
551
  ----------
514
- var : str, optional, default None
515
- Name of the artifact in which to store the caught exception.
516
- If not specified, the exception is not stored.
517
- print_exception : bool, default True
518
- Determines whether or not the exception is printed to
519
- stdout when caught.
552
+ packages : Dict[str, str], default {}
553
+ Packages to use for this step. The key is the name of the package
554
+ and the value is the version to use.
555
+ libraries : Dict[str, str], default {}
556
+ Supported for backward compatibility. When used with packages, packages will take precedence.
557
+ python : str, optional, default None
558
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
559
+ that the version used will correspond to the version of the Python interpreter used to start the run.
560
+ disabled : bool, default False
561
+ If set to True, disables @conda.
520
562
  """
521
563
  ...
522
564
 
523
- 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]]]:
565
+ @typing.overload
566
+ 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]]]:
524
567
  """
525
- Decorator that helps cache, version and store models/datasets from huggingface hub.
568
+ Enables checkpointing for a step.
569
+
526
570
 
527
571
 
528
572
  Parameters
529
573
  ----------
530
- temp_dir_root : str, optional
531
- The root directory that will hold the temporary directory where objects will be downloaded.
574
+ load_policy : str, default: "fresh"
575
+ The policy for loading the checkpoint. The following policies are supported:
576
+ - "eager": Loads the the latest available checkpoint within the namespace.
577
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
578
+ will be loaded at the start of the task.
579
+ - "none": Do not load any checkpoint
580
+ - "fresh": Loads the lastest checkpoint created within the running Task.
581
+ This mode helps loading checkpoints across various retry attempts of the same task.
582
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
583
+ created within the task will be loaded when the task is retries execution on failure.
532
584
 
533
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
534
- The list of repos (models/datasets) to load.
585
+ temp_dir_root : str, default: None
586
+ The root directory under which `current.checkpoint.directory` will be created.
587
+ """
588
+ ...
589
+
590
+ @typing.overload
591
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
592
+ ...
593
+
594
+ @typing.overload
595
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
596
+ ...
597
+
598
+ 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):
599
+ """
600
+ Enables checkpointing for a step.
535
601
 
536
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
537
602
 
538
- - If repo (model/dataset) is not found in the datastore:
539
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
540
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
541
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
542
603
 
543
- - If repo is found in the datastore:
544
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
604
+ Parameters
605
+ ----------
606
+ load_policy : str, default: "fresh"
607
+ The policy for loading the checkpoint. The following policies are supported:
608
+ - "eager": Loads the the latest available checkpoint within the namespace.
609
+ With this mode, the latest checkpoint written by any previous task (can be even a different run) of the step
610
+ will be loaded at the start of the task.
611
+ - "none": Do not load any checkpoint
612
+ - "fresh": Loads the lastest checkpoint created within the running Task.
613
+ This mode helps loading checkpoints across various retry attempts of the same task.
614
+ With this mode, no checkpoint will be loaded at the start of a task but any checkpoints
615
+ created within the task will be loaded when the task is retries execution on failure.
616
+
617
+ temp_dir_root : str, default: None
618
+ The root directory under which `current.checkpoint.directory` will be created.
545
619
  """
546
620
  ...
547
621
 
@@ -600,7 +674,7 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
600
674
  """
601
675
  ...
602
676
 
603
- def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
677
+ 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]]]:
604
678
  """
605
679
  Specifies that this step should execute on DGX cloud.
606
680
 
@@ -611,58 +685,136 @@ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[t
611
685
  Number of GPUs to use.
612
686
  gpu_type : str
613
687
  Type of Nvidia GPU to use.
614
- queue_timeout : int
615
- Time to keep the job in NVCF's queue.
616
688
  """
617
689
  ...
618
690
 
619
- @typing.overload
620
- def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
691
+ 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]]]:
621
692
  """
622
- Specifies the resources needed when executing this step.
623
-
624
- Use `@resources` to specify the resource requirements
625
- independently of the specific compute layer (`@batch`, `@kubernetes`).
626
-
627
- You can choose the compute layer on the command line by executing e.g.
628
- ```
629
- python myflow.py run --with batch
630
- ```
631
- or
632
- ```
633
- python myflow.py run --with kubernetes
634
- ```
635
- which executes the flow on the desired system using the
636
- requirements specified in `@resources`.
693
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
637
694
 
638
695
 
639
696
  Parameters
640
697
  ----------
641
- cpu : int, default 1
642
- Number of CPUs required for this step.
643
- gpu : int, optional, default None
644
- Number of GPUs required for this step.
645
- disk : int, optional, default None
646
- Disk size (in MB) required for this step. Only applies on Kubernetes.
647
- memory : int, default 4096
648
- Memory size (in MB) required for this step.
649
- shared_memory : int, optional, default None
650
- The value for the size (in MiB) of the /dev/shm volume for this step.
651
- This parameter maps to the `--shm-size` option in Docker.
698
+ temp_dir_root : str, optional
699
+ The root directory that will hold the temporary directory where objects will be downloaded.
700
+
701
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
702
+ The list of repos (models/datasets) to load.
703
+
704
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
705
+
706
+ - If repo (model/dataset) is not found in the datastore:
707
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
708
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
709
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
710
+
711
+ - If repo is found in the datastore:
712
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
652
713
  """
653
714
  ...
654
715
 
655
716
  @typing.overload
656
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
717
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
718
+ """
719
+ Internal decorator to support Fast bakery
720
+ """
657
721
  ...
658
722
 
659
723
  @typing.overload
660
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
724
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
661
725
  ...
662
726
 
663
- def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
727
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
664
728
  """
665
- Specifies the resources needed when executing this step.
729
+ Internal decorator to support Fast bakery
730
+ """
731
+ ...
732
+
733
+ 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]]]:
734
+ """
735
+ This decorator is used to run vllm APIs as Metaflow task sidecars.
736
+
737
+ User code call
738
+ --------------
739
+ @vllm(
740
+ model="...",
741
+ ...
742
+ )
743
+
744
+ Valid backend options
745
+ ---------------------
746
+ - 'local': Run as a separate process on the local task machine.
747
+
748
+ Valid model options
749
+ -------------------
750
+ Any HuggingFace model identifier, e.g. 'meta-llama/Llama-3.2-1B'
751
+
752
+ NOTE: vLLM's OpenAI-compatible server serves ONE model per server instance.
753
+ If you need multiple models, you must create multiple @vllm decorators.
754
+
755
+
756
+ Parameters
757
+ ----------
758
+ model: str
759
+ HuggingFace model identifier to be served by vLLM.
760
+ backend: str
761
+ Determines where and how to run the vLLM process.
762
+ debug: bool
763
+ Whether to turn on verbose debugging logs.
764
+ kwargs : Any
765
+ Any other keyword arguments are passed directly to the vLLM engine.
766
+ This allows for flexible configuration of vLLM server settings.
767
+ For example, `tensor_parallel_size=2`.
768
+ """
769
+ ...
770
+
771
+ @typing.overload
772
+ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
773
+ """
774
+ Specifies the resources needed when executing this step.
775
+
776
+ Use `@resources` to specify the resource requirements
777
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
778
+
779
+ You can choose the compute layer on the command line by executing e.g.
780
+ ```
781
+ python myflow.py run --with batch
782
+ ```
783
+ or
784
+ ```
785
+ python myflow.py run --with kubernetes
786
+ ```
787
+ which executes the flow on the desired system using the
788
+ requirements specified in `@resources`.
789
+
790
+
791
+ Parameters
792
+ ----------
793
+ cpu : int, default 1
794
+ Number of CPUs required for this step.
795
+ gpu : int, optional, default None
796
+ Number of GPUs required for this step.
797
+ disk : int, optional, default None
798
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
799
+ memory : int, default 4096
800
+ Memory size (in MB) required for this step.
801
+ shared_memory : int, optional, default None
802
+ The value for the size (in MiB) of the /dev/shm volume for this step.
803
+ This parameter maps to the `--shm-size` option in Docker.
804
+ """
805
+ ...
806
+
807
+ @typing.overload
808
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
809
+ ...
810
+
811
+ @typing.overload
812
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
813
+ ...
814
+
815
+ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Optional[int] = None, memory: int = 4096, shared_memory: typing.Optional[int] = None):
816
+ """
817
+ Specifies the resources needed when executing this step.
666
818
 
667
819
  Use `@resources` to specify the resource requirements
668
820
  independently of the specific compute layer (`@batch`, `@kubernetes`).
@@ -784,54 +936,18 @@ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: ty
784
936
  """
785
937
  ...
786
938
 
787
- @typing.overload
788
- 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]]]:
789
- """
790
- Specifies the PyPI packages for the step.
791
-
792
- Information in this decorator will augment any
793
- attributes set in the `@pyi_base` flow-level decorator. Hence,
794
- you can use `@pypi_base` to set packages required by all
795
- steps and use `@pypi` to specify step-specific overrides.
796
-
797
-
798
- Parameters
799
- ----------
800
- packages : Dict[str, str], default: {}
801
- Packages to use for this step. The key is the name of the package
802
- and the value is the version to use.
803
- python : str, optional, default: None
804
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
805
- that the version used will correspond to the version of the Python interpreter used to start the run.
806
- """
807
- ...
808
-
809
- @typing.overload
810
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
811
- ...
812
-
813
- @typing.overload
814
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
815
- ...
816
-
817
- 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):
939
+ def app_deploy(*, app_port: int, app_name: str) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
818
940
  """
819
- Specifies the PyPI packages for the step.
820
-
821
- Information in this decorator will augment any
822
- attributes set in the `@pyi_base` flow-level decorator. Hence,
823
- you can use `@pypi_base` to set packages required by all
824
- steps and use `@pypi` to specify step-specific overrides.
941
+ Specifies that this step is used to deploy an instance of the app.
942
+ Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
825
943
 
826
944
 
827
945
  Parameters
828
946
  ----------
829
- packages : Dict[str, str], default: {}
830
- Packages to use for this step. The key is the name of the package
831
- and the value is the version to use.
832
- python : str, optional, default: None
833
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
834
- that the version used will correspond to the version of the Python interpreter used to start the run.
947
+ app_port : int
948
+ Number of GPUs to use.
949
+ app_name : str
950
+ Name of the app to deploy.
835
951
  """
836
952
  ...
837
953
 
@@ -884,66 +1000,7 @@ def card(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
884
1000
  """
885
1001
  ...
886
1002
 
887
- @typing.overload
888
- 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]]]:
889
- """
890
- Specifies the Conda environment for the step.
891
-
892
- Information in this decorator will augment any
893
- attributes set in the `@conda_base` flow-level decorator. Hence,
894
- you can use `@conda_base` to set packages required by all
895
- steps and use `@conda` to specify step-specific overrides.
896
-
897
-
898
- Parameters
899
- ----------
900
- packages : Dict[str, str], default {}
901
- Packages to use for this step. The key is the name of the package
902
- and the value is the version to use.
903
- libraries : Dict[str, str], default {}
904
- Supported for backward compatibility. When used with packages, packages will take precedence.
905
- python : str, optional, default None
906
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
907
- that the version used will correspond to the version of the Python interpreter used to start the run.
908
- disabled : bool, default False
909
- If set to True, disables @conda.
910
- """
911
- ...
912
-
913
- @typing.overload
914
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
915
- ...
916
-
917
- @typing.overload
918
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
919
- ...
920
-
921
- 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):
922
- """
923
- Specifies the Conda environment for the step.
924
-
925
- Information in this decorator will augment any
926
- attributes set in the `@conda_base` flow-level decorator. Hence,
927
- you can use `@conda_base` to set packages required by all
928
- steps and use `@conda` to specify step-specific overrides.
929
-
930
-
931
- Parameters
932
- ----------
933
- packages : Dict[str, str], default {}
934
- Packages to use for this step. The key is the name of the package
935
- and the value is the version to use.
936
- libraries : Dict[str, str], default {}
937
- Supported for backward compatibility. When used with packages, packages will take precedence.
938
- python : str, optional, default None
939
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
940
- that the version used will correspond to the version of the Python interpreter used to start the run.
941
- disabled : bool, default False
942
- If set to True, disables @conda.
943
- """
944
- ...
945
-
946
- 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]]]:
1003
+ def nvidia(*, gpu: int, gpu_type: str, queue_timeout: int) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
947
1004
  """
948
1005
  Specifies that this step should execute on DGX cloud.
949
1006
 
@@ -954,65 +1011,8 @@ def nvct(*, gpu: int, gpu_type: str) -> typing.Callable[[typing.Union[typing.Cal
954
1011
  Number of GPUs to use.
955
1012
  gpu_type : str
956
1013
  Type of Nvidia GPU to use.
957
- """
958
- ...
959
-
960
- @typing.overload
961
- 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]]]:
962
- """
963
- Specifies a timeout for your step.
964
-
965
- This decorator is useful if this step may hang indefinitely.
966
-
967
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
968
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
969
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
970
-
971
- Note that all the values specified in parameters are added together so if you specify
972
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
973
-
974
-
975
- Parameters
976
- ----------
977
- seconds : int, default 0
978
- Number of seconds to wait prior to timing out.
979
- minutes : int, default 0
980
- Number of minutes to wait prior to timing out.
981
- hours : int, default 0
982
- Number of hours to wait prior to timing out.
983
- """
984
- ...
985
-
986
- @typing.overload
987
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
988
- ...
989
-
990
- @typing.overload
991
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
992
- ...
993
-
994
- 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):
995
- """
996
- Specifies a timeout for your step.
997
-
998
- This decorator is useful if this step may hang indefinitely.
999
-
1000
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
1001
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
1002
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
1003
-
1004
- Note that all the values specified in parameters are added together so if you specify
1005
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
1006
-
1007
-
1008
- Parameters
1009
- ----------
1010
- seconds : int, default 0
1011
- Number of seconds to wait prior to timing out.
1012
- minutes : int, default 0
1013
- Number of minutes to wait prior to timing out.
1014
- hours : int, default 0
1015
- Number of hours to wait prior to timing out.
1014
+ queue_timeout : int
1015
+ Time to keep the job in NVCF's queue.
1016
1016
  """
1017
1017
  ...
1018
1018
 
@@ -1067,161 +1067,6 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1067
1067
  """
1068
1068
  ...
1069
1069
 
1070
- @typing.overload
1071
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1072
- """
1073
- Specifies the PyPI packages for all steps of the flow.
1074
-
1075
- Use `@pypi_base` to set common packages required by all
1076
- steps and use `@pypi` to specify step-specific overrides.
1077
-
1078
- Parameters
1079
- ----------
1080
- packages : Dict[str, str], default: {}
1081
- Packages to use for this flow. The key is the name of the package
1082
- and the value is the version to use.
1083
- python : str, optional, default: None
1084
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1085
- that the version used will correspond to the version of the Python interpreter used to start the run.
1086
- """
1087
- ...
1088
-
1089
- @typing.overload
1090
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1091
- ...
1092
-
1093
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1094
- """
1095
- Specifies the PyPI packages for all steps of the flow.
1096
-
1097
- Use `@pypi_base` to set common packages required by all
1098
- steps and use `@pypi` to specify step-specific overrides.
1099
-
1100
- Parameters
1101
- ----------
1102
- packages : Dict[str, str], default: {}
1103
- Packages to use for this flow. The key is the name of the package
1104
- and the value is the version to use.
1105
- python : str, optional, default: None
1106
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1107
- that the version used will correspond to the version of the Python interpreter used to start the run.
1108
- """
1109
- ...
1110
-
1111
- def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1112
- """
1113
- Allows setting external datastores to save data for the
1114
- `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1115
-
1116
- This decorator is useful when users wish to save data to a different datastore
1117
- than what is configured in Metaflow. This can be for variety of reasons:
1118
-
1119
- 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1120
- 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1121
- - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1122
- 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1123
- - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1124
-
1125
- Usage:
1126
- ----------
1127
-
1128
- - Using a custom IAM role to access the datastore.
1129
-
1130
- ```python
1131
- @with_artifact_store(
1132
- type="s3",
1133
- config=lambda: {
1134
- "root": "s3://my-bucket-foo/path/to/root",
1135
- "role_arn": ROLE,
1136
- },
1137
- )
1138
- class MyFlow(FlowSpec):
1139
-
1140
- @checkpoint
1141
- @step
1142
- def start(self):
1143
- with open("my_file.txt", "w") as f:
1144
- f.write("Hello, World!")
1145
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1146
- self.next(self.end)
1147
-
1148
- ```
1149
-
1150
- - Using credentials to access the s3-compatible datastore.
1151
-
1152
- ```python
1153
- @with_artifact_store(
1154
- type="s3",
1155
- config=lambda: {
1156
- "root": "s3://my-bucket-foo/path/to/root",
1157
- "client_params": {
1158
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1159
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1160
- },
1161
- },
1162
- )
1163
- class MyFlow(FlowSpec):
1164
-
1165
- @checkpoint
1166
- @step
1167
- def start(self):
1168
- with open("my_file.txt", "w") as f:
1169
- f.write("Hello, World!")
1170
- self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1171
- self.next(self.end)
1172
-
1173
- ```
1174
-
1175
- - Accessing objects stored in external datastores after task execution.
1176
-
1177
- ```python
1178
- run = Run("CheckpointsTestsFlow/8992")
1179
- with artifact_store_from(run=run, config={
1180
- "client_params": {
1181
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1182
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1183
- },
1184
- }):
1185
- with Checkpoint() as cp:
1186
- latest = cp.list(
1187
- task=run["start"].task
1188
- )[0]
1189
- print(latest)
1190
- cp.load(
1191
- latest,
1192
- "test-checkpoints"
1193
- )
1194
-
1195
- task = Task("TorchTuneFlow/8484/train/53673")
1196
- with artifact_store_from(run=run, config={
1197
- "client_params": {
1198
- "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1199
- "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1200
- },
1201
- }):
1202
- load_model(
1203
- task.data.model_ref,
1204
- "test-models"
1205
- )
1206
- ```
1207
- Parameters:
1208
- ----------
1209
-
1210
- type: str
1211
- The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1212
-
1213
- config: dict or Callable
1214
- Dictionary of configuration options for the datastore. The following keys are required:
1215
- - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1216
- - example: 's3://bucket-name/path/to/root'
1217
- - example: 'gs://bucket-name/path/to/root'
1218
- - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1219
- - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1220
- - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1221
- - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1222
- """
1223
- ...
1224
-
1225
1070
  @typing.overload
1226
1071
  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]]:
1227
1072
  """
@@ -1271,98 +1116,47 @@ def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = No
1271
1116
  def trigger(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1272
1117
  ...
1273
1118
 
1274
- 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] = {}):
1275
- """
1276
- Specifies the event(s) that this flow depends on.
1277
-
1278
- ```
1279
- @trigger(event='foo')
1280
- ```
1281
- or
1282
- ```
1283
- @trigger(events=['foo', 'bar'])
1284
- ```
1285
-
1286
- Additionally, you can specify the parameter mappings
1287
- to map event payload to Metaflow parameters for the flow.
1288
- ```
1289
- @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1290
- ```
1291
- or
1292
- ```
1293
- @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1294
- {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1295
- ```
1296
-
1297
- 'parameters' can also be a list of strings and tuples like so:
1298
- ```
1299
- @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1300
- ```
1301
- This is equivalent to:
1302
- ```
1303
- @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1304
- ```
1305
-
1306
-
1307
- Parameters
1308
- ----------
1309
- event : Union[str, Dict[str, Any]], optional, default None
1310
- Event dependency for this flow.
1311
- events : List[Union[str, Dict[str, Any]]], default []
1312
- Events dependency for this flow.
1313
- options : Dict[str, Any], default {}
1314
- Backend-specific configuration for tuning eventing behavior.
1315
- """
1316
- ...
1317
-
1318
- @typing.overload
1319
- 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]]:
1320
- """
1321
- Specifies the times when the flow should be run when running on a
1322
- production scheduler.
1323
-
1324
-
1325
- Parameters
1326
- ----------
1327
- hourly : bool, default False
1328
- Run the workflow hourly.
1329
- daily : bool, default True
1330
- Run the workflow daily.
1331
- weekly : bool, default False
1332
- Run the workflow weekly.
1333
- cron : str, optional, default None
1334
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1335
- specified by this expression.
1336
- timezone : str, optional, default None
1337
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1338
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1339
- """
1340
- ...
1341
-
1342
- @typing.overload
1343
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1344
- ...
1345
-
1346
- def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1119
+ 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] = {}):
1347
1120
  """
1348
- Specifies the times when the flow should be run when running on a
1349
- production scheduler.
1121
+ Specifies the event(s) that this flow depends on.
1122
+
1123
+ ```
1124
+ @trigger(event='foo')
1125
+ ```
1126
+ or
1127
+ ```
1128
+ @trigger(events=['foo', 'bar'])
1129
+ ```
1130
+
1131
+ Additionally, you can specify the parameter mappings
1132
+ to map event payload to Metaflow parameters for the flow.
1133
+ ```
1134
+ @trigger(event={'name':'foo', 'parameters':{'flow_param': 'event_field'}})
1135
+ ```
1136
+ or
1137
+ ```
1138
+ @trigger(events=[{'name':'foo', 'parameters':{'flow_param_1': 'event_field_1'},
1139
+ {'name':'bar', 'parameters':{'flow_param_2': 'event_field_2'}])
1140
+ ```
1141
+
1142
+ 'parameters' can also be a list of strings and tuples like so:
1143
+ ```
1144
+ @trigger(event={'name':'foo', 'parameters':['common_name', ('flow_param', 'event_field')]})
1145
+ ```
1146
+ This is equivalent to:
1147
+ ```
1148
+ @trigger(event={'name':'foo', 'parameters':{'common_name': 'common_name', 'flow_param': 'event_field'}})
1149
+ ```
1350
1150
 
1351
1151
 
1352
1152
  Parameters
1353
1153
  ----------
1354
- hourly : bool, default False
1355
- Run the workflow hourly.
1356
- daily : bool, default True
1357
- Run the workflow daily.
1358
- weekly : bool, default False
1359
- Run the workflow weekly.
1360
- cron : str, optional, default None
1361
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1362
- specified by this expression.
1363
- timezone : str, optional, default None
1364
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1365
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1154
+ event : Union[str, Dict[str, Any]], optional, default None
1155
+ Event dependency for this flow.
1156
+ events : List[Union[str, Dict[str, Any]]], default []
1157
+ Events dependency for this flow.
1158
+ options : Dict[str, Any], default {}
1159
+ Backend-specific configuration for tuning eventing behavior.
1366
1160
  """
1367
1161
  ...
1368
1162
 
@@ -1409,6 +1203,47 @@ def airflow_s3_key_sensor(*, timeout: int, poke_interval: int, mode: str, expone
1409
1203
  """
1410
1204
  ...
1411
1205
 
1206
+ @typing.overload
1207
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1208
+ """
1209
+ Specifies the PyPI packages for all steps of the flow.
1210
+
1211
+ Use `@pypi_base` to set common packages required by all
1212
+ steps and use `@pypi` to specify step-specific overrides.
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
+ python : str, optional, default: None
1220
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1221
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1222
+ """
1223
+ ...
1224
+
1225
+ @typing.overload
1226
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1227
+ ...
1228
+
1229
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1230
+ """
1231
+ Specifies the PyPI packages for all steps of the flow.
1232
+
1233
+ Use `@pypi_base` to set common packages required by all
1234
+ steps and use `@pypi` to specify step-specific overrides.
1235
+
1236
+ Parameters
1237
+ ----------
1238
+ packages : Dict[str, str], default: {}
1239
+ Packages to use for this flow. The key is the name of the package
1240
+ and the value is the version to use.
1241
+ python : str, optional, default: None
1242
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1243
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1244
+ """
1245
+ ...
1246
+
1412
1247
  @typing.overload
1413
1248
  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]]:
1414
1249
  """
@@ -1510,6 +1345,57 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1510
1345
  """
1511
1346
  ...
1512
1347
 
1348
+ @typing.overload
1349
+ 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]]:
1350
+ """
1351
+ Specifies the times when the flow should be run when running on a
1352
+ production scheduler.
1353
+
1354
+
1355
+ Parameters
1356
+ ----------
1357
+ hourly : bool, default False
1358
+ Run the workflow hourly.
1359
+ daily : bool, default True
1360
+ Run the workflow daily.
1361
+ weekly : bool, default False
1362
+ Run the workflow weekly.
1363
+ cron : str, optional, default None
1364
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1365
+ specified by this expression.
1366
+ timezone : str, optional, default None
1367
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1368
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1369
+ """
1370
+ ...
1371
+
1372
+ @typing.overload
1373
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1374
+ ...
1375
+
1376
+ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly: bool = False, daily: bool = True, weekly: bool = False, cron: typing.Optional[str] = None, timezone: typing.Optional[str] = None):
1377
+ """
1378
+ Specifies the times when the flow should be run when running on a
1379
+ production scheduler.
1380
+
1381
+
1382
+ Parameters
1383
+ ----------
1384
+ hourly : bool, default False
1385
+ Run the workflow hourly.
1386
+ daily : bool, default True
1387
+ Run the workflow daily.
1388
+ weekly : bool, default False
1389
+ Run the workflow weekly.
1390
+ cron : str, optional, default None
1391
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1392
+ specified by this expression.
1393
+ timezone : str, optional, default None
1394
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1395
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1396
+ """
1397
+ ...
1398
+
1513
1399
  def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1514
1400
  """
1515
1401
  Specifies what flows belong to the same project.
@@ -1588,5 +1474,119 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1588
1474
  """
1589
1475
  ...
1590
1476
 
1477
+ def with_artifact_store(f: typing.Optional[typing.Type[FlowSpecDerived]] = None):
1478
+ """
1479
+ Allows setting external datastores to save data for the
1480
+ `@checkpoint`/`@model`/`@huggingface_hub` decorators.
1481
+
1482
+ This decorator is useful when users wish to save data to a different datastore
1483
+ than what is configured in Metaflow. This can be for variety of reasons:
1484
+
1485
+ 1. Data security: The objects needs to be stored in a bucket (object storage) that is not accessible by other flows.
1486
+ 2. Data Locality: The location where the task is executing is not located in the same region as the datastore.
1487
+ - Example: Metaflow datastore lives in US East, but the task is executing in Finland datacenters.
1488
+ 3. Data Lifecycle Policies: The objects need to be archived / managed separately from the Metaflow managed objects.
1489
+ - Example: Flow is training very large models that need to be stored separately and will be deleted more aggressively than the Metaflow managed objects.
1490
+
1491
+ Usage:
1492
+ ----------
1493
+
1494
+ - Using a custom IAM role to access the datastore.
1495
+
1496
+ ```python
1497
+ @with_artifact_store(
1498
+ type="s3",
1499
+ config=lambda: {
1500
+ "root": "s3://my-bucket-foo/path/to/root",
1501
+ "role_arn": ROLE,
1502
+ },
1503
+ )
1504
+ class MyFlow(FlowSpec):
1505
+
1506
+ @checkpoint
1507
+ @step
1508
+ def start(self):
1509
+ with open("my_file.txt", "w") as f:
1510
+ f.write("Hello, World!")
1511
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1512
+ self.next(self.end)
1513
+
1514
+ ```
1515
+
1516
+ - Using credentials to access the s3-compatible datastore.
1517
+
1518
+ ```python
1519
+ @with_artifact_store(
1520
+ type="s3",
1521
+ config=lambda: {
1522
+ "root": "s3://my-bucket-foo/path/to/root",
1523
+ "client_params": {
1524
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1525
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1526
+ },
1527
+ },
1528
+ )
1529
+ class MyFlow(FlowSpec):
1530
+
1531
+ @checkpoint
1532
+ @step
1533
+ def start(self):
1534
+ with open("my_file.txt", "w") as f:
1535
+ f.write("Hello, World!")
1536
+ self.external_bucket_checkpoint = current.checkpoint.save("my_file.txt")
1537
+ self.next(self.end)
1538
+
1539
+ ```
1540
+
1541
+ - Accessing objects stored in external datastores after task execution.
1542
+
1543
+ ```python
1544
+ run = Run("CheckpointsTestsFlow/8992")
1545
+ with artifact_store_from(run=run, config={
1546
+ "client_params": {
1547
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1548
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1549
+ },
1550
+ }):
1551
+ with Checkpoint() as cp:
1552
+ latest = cp.list(
1553
+ task=run["start"].task
1554
+ )[0]
1555
+ print(latest)
1556
+ cp.load(
1557
+ latest,
1558
+ "test-checkpoints"
1559
+ )
1560
+
1561
+ task = Task("TorchTuneFlow/8484/train/53673")
1562
+ with artifact_store_from(run=run, config={
1563
+ "client_params": {
1564
+ "aws_access_key_id": os.environ.get("MY_CUSTOM_ACCESS_KEY"),
1565
+ "aws_secret_access_key": os.environ.get("MY_CUSTOM_SECRET_KEY"),
1566
+ },
1567
+ }):
1568
+ load_model(
1569
+ task.data.model_ref,
1570
+ "test-models"
1571
+ )
1572
+ ```
1573
+ Parameters:
1574
+ ----------
1575
+
1576
+ type: str
1577
+ The type of the datastore. Can be one of 's3', 'gcs', 'azure' or any other supported metaflow Datastore.
1578
+
1579
+ config: dict or Callable
1580
+ Dictionary of configuration options for the datastore. The following keys are required:
1581
+ - root: The root path in the datastore where the data will be saved. (needs to be in the format expected by the datastore)
1582
+ - example: 's3://bucket-name/path/to/root'
1583
+ - example: 'gs://bucket-name/path/to/root'
1584
+ - example: 'https://myblockacc.blob.core.windows.net/metaflow/'
1585
+ - role_arn (optional): AWS IAM role to access s3 bucket (only when `type` is 's3')
1586
+ - session_vars (optional): AWS session variables to access s3 bucket (only when `type` is 's3')
1587
+ - client_params (optional): AWS client parameters to access s3 bucket (only when `type` is 's3')
1588
+ """
1589
+ ...
1590
+
1591
1591
  pkg_name: str
1592
1592