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