ob-metaflow-stubs 6.0.3.176rc5__py2.py3-none-any.whl → 6.0.3.176rc6__py2.py3-none-any.whl

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