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