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