ob-metaflow-stubs 6.0.3.177__py2.py3-none-any.whl → 6.0.3.179rc0__py2.py3-none-any.whl

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