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