ob-metaflow-stubs 6.0.3.179rc4__py2.py3-none-any.whl → 6.0.3.179rc5__py2.py3-none-any.whl

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