ob-metaflow-stubs 6.0.3.152__py2.py3-none-any.whl → 6.0.3.154__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 +697 -697
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +4 -4
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +59 -59
  21. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  22. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  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 +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +3 -3
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  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 +2 -2
  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 +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +4 -4
  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 +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +2 -2
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  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 +2 -2
  43. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/task_utils.pyi +3 -3
  44. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/utils.pyi +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +2 -2
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  48. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/decorator.pyi +3 -3
  49. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/__init__.pyi +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  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 +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  56. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/identity_utils.pyi +3 -3
  57. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/__init__.pyi +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +2 -2
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +2 -2
  61. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  62. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  63. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/card_utilities/injector.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/plugins/ollama/ollama.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  70. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  71. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  72. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  73. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  74. metaflow-stubs/multicore_utils.pyi +2 -2
  75. metaflow-stubs/parameters.pyi +3 -3
  76. metaflow-stubs/plugins/__init__.pyi +13 -13
  77. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  78. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  79. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  80. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  81. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  82. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  83. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  84. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  85. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  86. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  87. metaflow-stubs/plugins/argo/argo_workflows.pyi +10 -4
  88. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  89. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  90. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  91. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  92. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  93. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  94. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  96. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  97. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  98. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +3 -3
  100. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  101. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  102. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  103. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  104. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  105. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  106. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  107. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  108. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  109. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  110. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  111. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  112. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  113. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  114. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  115. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  116. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  117. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  118. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  119. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  121. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  123. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  124. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  125. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  126. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  127. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  128. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  129. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  130. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  132. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  133. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  134. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  135. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  136. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  137. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  138. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  139. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  140. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  141. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  142. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  143. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  144. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  145. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  146. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  147. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  148. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  149. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  150. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  151. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  152. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  153. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  154. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  155. metaflow-stubs/plugins/ollama/__init__.pyi +2 -2
  156. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  157. metaflow-stubs/plugins/perimeters.pyi +2 -2
  158. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  159. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  160. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  162. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  163. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  164. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  165. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  166. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  168. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  169. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  170. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  171. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  172. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  173. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  174. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  175. metaflow-stubs/profilers/__init__.pyi +2 -2
  176. metaflow-stubs/pylint_wrapper.pyi +2 -2
  177. metaflow-stubs/runner/__init__.pyi +2 -2
  178. metaflow-stubs/runner/deployer.pyi +29 -29
  179. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  180. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  181. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  182. metaflow-stubs/runner/nbrun.pyi +2 -2
  183. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  184. metaflow-stubs/runner/utils.pyi +4 -4
  185. metaflow-stubs/system/__init__.pyi +2 -2
  186. metaflow-stubs/system/system_logger.pyi +2 -2
  187. metaflow-stubs/system/system_monitor.pyi +2 -2
  188. metaflow-stubs/tagging_util.pyi +2 -2
  189. metaflow-stubs/tuple_util.pyi +2 -2
  190. metaflow-stubs/user_configs/__init__.pyi +2 -2
  191. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  192. metaflow-stubs/user_configs/config_options.pyi +3 -3
  193. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  194. {ob_metaflow_stubs-6.0.3.152.dist-info → ob_metaflow_stubs-6.0.3.154.dist-info}/METADATA +1 -1
  195. ob_metaflow_stubs-6.0.3.154.dist-info/RECORD +198 -0
  196. ob_metaflow_stubs-6.0.3.152.dist-info/RECORD +0 -198
  197. {ob_metaflow_stubs-6.0.3.152.dist-info → ob_metaflow_stubs-6.0.3.154.dist-info}/WHEEL +0 -0
  198. {ob_metaflow_stubs-6.0.3.152.dist-info → ob_metaflow_stubs-6.0.3.154.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.7.1+obcheckpoint(0.1.9);ob(v1) #
4
- # Generated on 2025-03-24T18:51:07.966816 #
3
+ # MF version: 2.15.7.2+obcheckpoint(0.1.9);ob(v1) #
4
+ # Generated on 2025-04-10T22:13:15.011907 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -35,17 +35,17 @@ 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
38
39
  from . import tuple_util as tuple_util
39
40
  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 conda_environment_yml_parser as conda_environment_yml_parser
47
46
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
47
  from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
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
51
51
  from .client.core import get_namespace as get_namespace
@@ -150,19 +150,56 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
150
150
  ...
151
151
 
152
152
  @typing.overload
153
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
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]]]:
154
154
  """
155
- Internal decorator to support Fast bakery
155
+ Specifies secrets to be retrieved and injected as environment variables prior to
156
+ the execution of a step.
157
+
158
+
159
+ Parameters
160
+ ----------
161
+ sources : List[Union[str, Dict[str, Any]]], default: []
162
+ List of secret specs, defining how the secrets are to be retrieved
156
163
  """
157
164
  ...
158
165
 
159
166
  @typing.overload
160
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
167
+ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
161
168
  ...
162
169
 
163
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
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]]] = []):
164
175
  """
165
- Internal decorator to support Fast bakery
176
+ Specifies secrets to be retrieved and injected as environment variables prior to
177
+ the execution of a step.
178
+
179
+
180
+ Parameters
181
+ ----------
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().
166
203
  """
167
204
  ...
168
205
 
@@ -223,6 +260,125 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
223
260
  """
224
261
  ...
225
262
 
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]]]:
264
+ """
265
+ Specifies that this step should execute on DGX cloud.
266
+
267
+
268
+ Parameters
269
+ ----------
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.
276
+ """
277
+ ...
278
+
279
+ @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]]]:
281
+ """
282
+ Specifies the number of times the task corresponding
283
+ to a step needs to be retried.
284
+
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
+
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.
292
+
293
+
294
+ Parameters
295
+ ----------
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.
300
+ """
301
+ ...
302
+
303
+ @typing.overload
304
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
305
+ ...
306
+
307
+ @typing.overload
308
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
309
+ ...
310
+
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):
312
+ """
313
+ Specifies the number of times the task corresponding
314
+ to a step needs to be retried.
315
+
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)`.
319
+
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.
323
+
324
+
325
+ Parameters
326
+ ----------
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.
331
+ """
332
+ ...
333
+
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]]]:
335
+ """
336
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
337
+
338
+ User code call
339
+ -----------
340
+ @nim(
341
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
342
+ backend='managed'
343
+ )
344
+
345
+ Valid backend options
346
+ ---------------------
347
+ - 'managed': Outerbounds selects a compute provider based on the model.
348
+
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
354
+
355
+
356
+ Parameters
357
+ ----------
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.
364
+ """
365
+ ...
366
+
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]]]:
368
+ """
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.
371
+
372
+
373
+ Parameters
374
+ ----------
375
+ app_port : int
376
+ Number of GPUs to use.
377
+ app_name : str
378
+ Name of the app to deploy.
379
+ """
380
+ ...
381
+
226
382
  @typing.overload
227
383
  def conda(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
228
384
  """
@@ -283,233 +439,262 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
283
439
  ...
284
440
 
285
441
  @typing.overload
286
- 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]]]:
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]]]:
287
443
  """
288
- Specifies that the step will success under all circumstances.
444
+ Specifies a timeout for your step.
289
445
 
290
- The decorator will create an optional artifact, specified by `var`, which
291
- contains the exception raised. You can use it to detect the presence
292
- of errors, indicating that all happy-path artifacts produced by the step
293
- are missing.
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.
294
454
 
295
455
 
296
456
  Parameters
297
457
  ----------
298
- var : str, optional, default None
299
- Name of the artifact in which to store the caught exception.
300
- If not specified, the exception is not stored.
301
- print_exception : bool, default True
302
- Determines whether or not the exception is printed to
303
- stdout when caught.
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.
304
464
  """
305
465
  ...
306
466
 
307
467
  @typing.overload
308
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
468
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
309
469
  ...
310
470
 
311
471
  @typing.overload
312
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
472
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
313
473
  ...
314
474
 
315
- 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):
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):
316
476
  """
317
- Specifies that the step will success under all circumstances.
477
+ Specifies a timeout for your step.
318
478
 
319
- The decorator will create an optional artifact, specified by `var`, which
320
- contains the exception raised. You can use it to detect the presence
321
- of errors, indicating that all happy-path artifacts produced by the step
322
- are missing.
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.
323
487
 
324
488
 
325
489
  Parameters
326
490
  ----------
327
- var : str, optional, default None
328
- Name of the artifact in which to store the caught exception.
329
- If not specified, the exception is not stored.
330
- print_exception : bool, default True
331
- Determines whether or not the exception is printed to
332
- stdout when caught.
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.
333
497
  """
334
498
  ...
335
499
 
336
- def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
500
+ @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]]]:
337
502
  """
338
- Decorator that helps cache, version and store models/datasets from huggingface hub.
503
+ Specifies environment variables to be set prior to the execution of a step.
339
504
 
340
505
 
341
506
  Parameters
342
507
  ----------
343
- temp_dir_root : str, optional
344
- The root directory that will hold the temporary directory where objects will be downloaded.
345
-
346
- load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
347
- The list of repos (models/datasets) to load.
348
-
349
- Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
350
-
351
- - If repo (model/dataset) is not found in the datastore:
352
- - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
353
- - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
354
- - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
355
-
356
- - If repo is found in the datastore:
357
- - Loads it directly from datastore to local path (can be temporary directory or specified path)
508
+ vars : Dict[str, str], default {}
509
+ Dictionary of environment variables to set.
358
510
  """
359
511
  ...
360
512
 
361
513
  @typing.overload
362
- def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
363
- """
364
- Specifies the PyPI packages for the step.
365
-
366
- Information in this decorator will augment any
367
- attributes set in the `@pyi_base` flow-level decorator. Hence,
368
- you can use `@pypi_base` to set packages required by all
369
- steps and use `@pypi` to specify step-specific overrides.
370
-
371
-
372
- Parameters
373
- ----------
374
- packages : Dict[str, str], default: {}
375
- Packages to use for this step. The key is the name of the package
376
- and the value is the version to use.
377
- python : str, optional, default: None
378
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
379
- that the version used will correspond to the version of the Python interpreter used to start the run.
380
- """
381
- ...
382
-
383
- @typing.overload
384
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
514
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
385
515
  ...
386
516
 
387
517
  @typing.overload
388
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
518
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
389
519
  ...
390
520
 
391
- def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
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] = {}):
392
522
  """
393
- Specifies the PyPI packages for the step.
394
-
395
- Information in this decorator will augment any
396
- attributes set in the `@pyi_base` flow-level decorator. Hence,
397
- you can use `@pypi_base` to set packages required by all
398
- steps and use `@pypi` to specify step-specific overrides.
523
+ Specifies environment variables to be set prior to the execution of a step.
399
524
 
400
525
 
401
526
  Parameters
402
527
  ----------
403
- packages : Dict[str, str], default: {}
404
- Packages to use for this step. The key is the name of the package
405
- and the value is the version to use.
406
- python : str, optional, default: None
407
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
408
- that the version used will correspond to the version of the Python interpreter used to start the run.
528
+ vars : Dict[str, str], default {}
529
+ Dictionary of environment variables to set.
409
530
  """
410
531
  ...
411
532
 
412
- 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]]]:
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]]]:
413
534
  """
414
- Specifies that this step should execute on DGX cloud.
535
+ Specifies that this step should execute on Kubernetes.
415
536
 
416
537
 
417
538
  Parameters
418
539
  ----------
419
- gpu : int
420
- Number of GPUs to use.
421
- gpu_type : str
422
- Type of Nvidia GPU to use.
423
- queue_timeout : int
424
- Time to keep the job in NVCF's queue.
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
425
607
  """
426
608
  ...
427
609
 
428
- 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]]]:
610
+ @typing.overload
611
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
429
612
  """
430
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
431
-
432
- User code call
433
- -----------
434
- @nim(
435
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
436
- backend='managed'
437
- )
438
-
439
- Valid backend options
440
- ---------------------
441
- - 'managed': Outerbounds selects a compute provider based on the model.
442
-
443
- Valid model options
444
- ----------------
445
- - 'meta/llama3-8b-instruct': 8B parameter model
446
- - 'meta/llama3-70b-instruct': 70B parameter model
447
- - any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
448
-
449
-
450
- Parameters
451
- ----------
452
- models: list[NIM]
453
- List of NIM containers running models in sidecars.
454
- backend: str
455
- Compute provider to run the NIM container.
456
- queue_timeout : int
457
- Time to keep the job in NVCF's queue.
613
+ Internal decorator to support Fast bakery
458
614
  """
459
615
  ...
460
616
 
461
617
  @typing.overload
462
- 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]]]:
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):
463
622
  """
464
- Specifies the number of times the task corresponding
465
- to a step needs to be retried.
466
-
467
- This decorator is useful for handling transient errors, such as networking issues.
468
- If your task contains operations that can't be retried safely, e.g. database updates,
469
- it is advisable to annotate it with `@retry(times=0)`.
623
+ Internal decorator to support Fast bakery
624
+ """
625
+ ...
626
+
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]]]:
629
+ """
630
+ Creates a human-readable report, a Metaflow Card, after this step completes.
470
631
 
471
- This can be used in conjunction with the `@catch` decorator. The `@catch`
472
- decorator will execute a no-op task after all retries have been exhausted,
473
- ensuring that the flow execution can continue.
632
+ Note that you may add multiple `@card` decorators in a step with different parameters.
474
633
 
475
634
 
476
635
  Parameters
477
636
  ----------
478
- times : int, default 3
479
- Number of times to retry this task.
480
- minutes_between_retries : int, default 2
481
- Number of minutes between retries.
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.
482
645
  """
483
646
  ...
484
647
 
485
648
  @typing.overload
486
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
649
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
487
650
  ...
488
651
 
489
652
  @typing.overload
490
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
653
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
491
654
  ...
492
655
 
493
- 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):
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):
494
657
  """
495
- Specifies the number of times the task corresponding
496
- to a step needs to be retried.
658
+ Creates a human-readable report, a Metaflow Card, after this step completes.
497
659
 
498
- This decorator is useful for handling transient errors, such as networking issues.
499
- If your task contains operations that can't be retried safely, e.g. database updates,
500
- it is advisable to annotate it with `@retry(times=0)`.
660
+ Note that you may add multiple `@card` decorators in a step with different parameters.
501
661
 
502
- This can be used in conjunction with the `@catch` decorator. The `@catch`
503
- decorator will execute a no-op task after all retries have been exhausted,
504
- ensuring that the flow execution can continue.
662
+
663
+ Parameters
664
+ ----------
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.
673
+ """
674
+ ...
675
+
676
+ def huggingface_hub(*, temp_dir_root: typing.Optional[str] = None, load: typing.Union[typing.List[str], typing.List[typing.Tuple[typing.Dict, str]], typing.List[typing.Tuple[str, str]], typing.List[typing.Dict], None]) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
677
+ """
678
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
505
679
 
506
680
 
507
681
  Parameters
508
682
  ----------
509
- times : int, default 3
510
- Number of times to retry this task.
511
- minutes_between_retries : int, default 2
512
- Number of minutes between retries.
683
+ temp_dir_root : str, optional
684
+ The root directory that will hold the temporary directory where objects will be downloaded.
685
+
686
+ load: Union[List[str], List[Tuple[Dict, str]], List[Tuple[str, str]], List[Dict], None]
687
+ The list of repos (models/datasets) to load.
688
+
689
+ Loaded repos can be accessed via `current.huggingface_hub.loaded`. If load is set, then the following happens:
690
+
691
+ - If repo (model/dataset) is not found in the datastore:
692
+ - Downloads the repo from Hugging Face Hub to a temporary directory (or uses specified path) for local access
693
+ - Stores it in Metaflow's datastore (s3/gcs/azure etc.) with a unique name based on repo_type/repo_id
694
+ - All HF models loaded for a `@step` will be cached separately under flow/step/namespace.
695
+
696
+ - If repo is found in the datastore:
697
+ - Loads it directly from datastore to local path (can be temporary directory or specified path)
513
698
  """
514
699
  ...
515
700
 
@@ -592,139 +777,36 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
592
777
  """
593
778
  ...
594
779
 
595
- 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]]]:
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]]]:
596
781
  """
597
- Specifies that this step should execute on Kubernetes.
782
+ This decorator is used to run Ollama APIs as Metaflow task sidecars.
598
783
 
784
+ User code call
785
+ -----------
786
+ @ollama(
787
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
788
+ backend='local'
789
+ )
599
790
 
600
- Parameters
601
- ----------
602
- cpu : int, default 1
603
- Number of CPUs required for this step. If `@resources` is
604
- also present, the maximum value from all decorators is used.
605
- memory : int, default 4096
606
- Memory size (in MB) required for this step. If
607
- `@resources` is also present, the maximum value from all decorators is
608
- used.
609
- disk : int, default 10240
610
- Disk size (in MB) required for this step. If
611
- `@resources` is also present, the maximum value from all decorators is
612
- used.
613
- image : str, optional, default None
614
- Docker image to use when launching on Kubernetes. If not specified, and
615
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
616
- not, a default Docker image mapping to the current version of Python is used.
617
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
618
- If given, the imagePullPolicy to be applied to the Docker image of the step.
619
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
620
- Kubernetes service account to use when launching pod in Kubernetes.
621
- secrets : List[str], optional, default None
622
- Kubernetes secrets to use when launching pod in Kubernetes. These
623
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
624
- in Metaflow configuration.
625
- node_selector: Union[Dict[str,str], str], optional, default None
626
- Kubernetes node selector(s) to apply to the pod running the task.
627
- Can be passed in as a comma separated string of values e.g.
628
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
629
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
630
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
631
- Kubernetes namespace to use when launching pod in Kubernetes.
632
- gpu : int, optional, default None
633
- Number of GPUs required for this step. A value of zero implies that
634
- the scheduled node should not have GPUs.
635
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
636
- The vendor of the GPUs to be used for this step.
637
- tolerations : List[str], default []
638
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
639
- Kubernetes tolerations to use when launching pod in Kubernetes.
640
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
641
- Kubernetes labels to use when launching pod in Kubernetes.
642
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
643
- Kubernetes annotations to use when launching pod in Kubernetes.
644
- use_tmpfs : bool, default False
645
- This enables an explicit tmpfs mount for this step.
646
- tmpfs_tempdir : bool, default True
647
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
648
- tmpfs_size : int, optional, default: None
649
- The value for the size (in MiB) of the tmpfs mount for this step.
650
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
651
- memory allocated for this step.
652
- tmpfs_path : str, optional, default /metaflow_temp
653
- Path to tmpfs mount for this step.
654
- persistent_volume_claims : Dict[str, str], optional, default None
655
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
656
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
657
- shared_memory: int, optional
658
- Shared memory size (in MiB) required for this step
659
- port: int, optional
660
- Port number to specify in the Kubernetes job object
661
- compute_pool : str, optional, default None
662
- Compute pool to be used for for this step.
663
- If not specified, any accessible compute pool within the perimeter is used.
664
- hostname_resolution_timeout: int, default 10 * 60
665
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
666
- Only applicable when @parallel is used.
667
- qos: str, default: Burstable
668
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
669
- """
670
- ...
671
-
672
- @typing.overload
673
- 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]]]:
674
- """
675
- Specifies a timeout for your step.
676
-
677
- This decorator is useful if this step may hang indefinitely.
678
-
679
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
680
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
681
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
682
-
683
- Note that all the values specified in parameters are added together so if you specify
684
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
685
-
686
-
687
- Parameters
688
- ----------
689
- seconds : int, default 0
690
- Number of seconds to wait prior to timing out.
691
- minutes : int, default 0
692
- Number of minutes to wait prior to timing out.
693
- hours : int, default 0
694
- Number of hours to wait prior to timing out.
695
- """
696
- ...
697
-
698
- @typing.overload
699
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
700
- ...
701
-
702
- @typing.overload
703
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
704
- ...
705
-
706
- 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):
707
- """
708
- Specifies a timeout for your step.
709
-
710
- This decorator is useful if this step may hang indefinitely.
711
-
712
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
713
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
714
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
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.
715
796
 
716
- Note that all the values specified in parameters are added together so if you specify
717
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
797
+ Valid model options
798
+ ----------------
799
+ - 'llama3.2'
800
+ - 'llama3.3'
801
+ - any model here https://ollama.com/search
718
802
 
719
803
 
720
804
  Parameters
721
805
  ----------
722
- seconds : int, default 0
723
- Number of seconds to wait prior to timing out.
724
- minutes : int, default 0
725
- Number of minutes to wait prior to timing out.
726
- hours : int, default 0
727
- Number of hours to wait prior to timing out.
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.
728
810
  """
729
811
  ...
730
812
 
@@ -786,464 +868,131 @@ def checkpoint(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None
786
868
  ...
787
869
 
788
870
  @typing.overload
789
- 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]]]:
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]]]:
790
872
  """
791
- Creates a human-readable report, a Metaflow Card, after this step completes.
873
+ Specifies that the step will success under all circumstances.
792
874
 
793
- Note that you may add multiple `@card` decorators in a step with different parameters.
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.
794
879
 
795
880
 
796
881
  Parameters
797
882
  ----------
798
- type : str, default 'default'
799
- Card type.
800
- id : str, optional, default None
801
- If multiple cards are present, use this id to identify this card.
802
- options : Dict[str, Any], default {}
803
- Options passed to the card. The contents depend on the card type.
804
- timeout : int, default 45
805
- Interrupt reporting if it takes more than this many seconds.
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.
806
889
  """
807
890
  ...
808
891
 
809
892
  @typing.overload
810
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
893
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
811
894
  ...
812
895
 
813
896
  @typing.overload
814
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
897
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
815
898
  ...
816
899
 
817
- 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):
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):
818
901
  """
819
- Creates a human-readable report, a Metaflow Card, after this step completes.
902
+ Specifies that the step will success under all circumstances.
820
903
 
821
- Note that you may add multiple `@card` decorators in a step with different parameters.
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.
822
908
 
823
909
 
824
910
  Parameters
825
911
  ----------
826
- type : str, default 'default'
827
- Card type.
828
- id : str, optional, default None
829
- If multiple cards are present, use this id to identify this card.
830
- options : Dict[str, Any], default {}
831
- Options passed to the card. The contents depend on the card type.
832
- timeout : int, default 45
833
- Interrupt reporting if it takes more than this many seconds.
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.
834
918
  """
835
919
  ...
836
920
 
837
- 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]]]:
921
+ @typing.overload
922
+ def pypi(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
838
923
  """
839
- This decorator is used to run Ollama APIs as Metaflow task sidecars.
840
-
841
- User code call
842
- -----------
843
- @ollama(
844
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
845
- backend='local'
846
- )
847
-
848
- Valid backend options
849
- ---------------------
850
- - 'local': Run as a separate process on the local task machine.
851
- - (TODO) 'managed': Outerbounds hosts and selects compute provider.
852
- - (TODO) 'remote': Spin up separate instance to serve Ollama models.
924
+ Specifies the PyPI packages for the step.
853
925
 
854
- Valid model options
855
- ----------------
856
- - 'llama3.2'
857
- - 'llama3.3'
858
- - any model here https://ollama.com/search
926
+ Information in this decorator will augment any
927
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
928
+ you can use `@pypi_base` to set packages required by all
929
+ steps and use `@pypi` to specify step-specific overrides.
859
930
 
860
931
 
861
932
  Parameters
862
933
  ----------
863
- models: list[Ollama]
864
- List of Ollama containers running models in sidecars.
865
- backend: str
866
- Determines where and how to run the Ollama process.
934
+ packages : Dict[str, str], default: {}
935
+ Packages to use for this step. The key is the name of the package
936
+ and the value is the version to use.
937
+ python : str, optional, default: None
938
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
939
+ that the version used will correspond to the version of the Python interpreter used to start the run.
867
940
  """
868
941
  ...
869
942
 
870
943
  @typing.overload
871
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
872
- """
873
- Decorator prototype for all step decorators. This function gets specialized
874
- and imported for all decorators types by _import_plugin_decorators().
875
- """
944
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
876
945
  ...
877
946
 
878
947
  @typing.overload
879
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
948
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
880
949
  ...
881
950
 
882
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
951
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
883
952
  """
884
- Decorator prototype for all step decorators. This function gets specialized
885
- and imported for all decorators types by _import_plugin_decorators().
953
+ Specifies the PyPI packages for the step.
954
+
955
+ Information in this decorator will augment any
956
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
957
+ you can use `@pypi_base` to set packages required by all
958
+ steps and use `@pypi` to specify step-specific overrides.
959
+
960
+
961
+ Parameters
962
+ ----------
963
+ packages : Dict[str, str], default: {}
964
+ Packages to use for this step. The key is the name of the package
965
+ and the value is the version to use.
966
+ python : str, optional, default: None
967
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
968
+ that the version used will correspond to the version of the Python interpreter used to start the run.
886
969
  """
887
970
  ...
888
971
 
889
972
  @typing.overload
890
- 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]]]:
973
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
891
974
  """
892
- Specifies secrets to be retrieved and injected as environment variables prior to
893
- the execution of a step.
975
+ Specifies the PyPI packages for all steps of the flow.
894
976
 
977
+ Use `@pypi_base` to set common packages required by all
978
+ steps and use `@pypi` to specify step-specific overrides.
895
979
 
896
980
  Parameters
897
981
  ----------
898
- sources : List[Union[str, Dict[str, Any]]], default: []
899
- List of secret specs, defining how the secrets are to be retrieved
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.
900
988
  """
901
989
  ...
902
990
 
903
991
  @typing.overload
904
- def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
905
- ...
906
-
907
- @typing.overload
908
- def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
909
- ...
910
-
911
- 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]]] = []):
912
- """
913
- Specifies secrets to be retrieved and injected as environment variables prior to
914
- the execution of a step.
915
-
916
-
917
- Parameters
918
- ----------
919
- sources : List[Union[str, Dict[str, Any]]], default: []
920
- List of secret specs, defining how the secrets are to be retrieved
921
- """
922
- ...
923
-
924
- @typing.overload
925
- 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]]]:
926
- """
927
- Specifies environment variables to be set prior to the execution of a step.
928
-
929
-
930
- Parameters
931
- ----------
932
- vars : Dict[str, str], default {}
933
- Dictionary of environment variables to set.
934
- """
935
- ...
936
-
937
- @typing.overload
938
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
939
- ...
940
-
941
- @typing.overload
942
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
943
- ...
944
-
945
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
946
- """
947
- Specifies environment variables to be set prior to the execution of a step.
948
-
949
-
950
- Parameters
951
- ----------
952
- vars : Dict[str, str], default {}
953
- Dictionary of environment variables to set.
954
- """
955
- ...
956
-
957
- 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]]]:
958
- """
959
- Specifies that this step is used to deploy an instance of the app.
960
- Requires that self.app_name, self.app_port, self.entrypoint and self.deployDir is set.
961
-
962
-
963
- Parameters
964
- ----------
965
- app_port : int
966
- Number of GPUs to use.
967
- app_name : str
968
- Name of the app to deploy.
969
- """
970
- ...
971
-
972
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
973
- """
974
- Specifies what flows belong to the same project.
975
-
976
- A project-specific namespace is created for all flows that
977
- use the same `@project(name)`.
978
-
979
-
980
- Parameters
981
- ----------
982
- name : str
983
- Project name. Make sure that the name is unique amongst all
984
- projects that use the same production scheduler. The name may
985
- contain only lowercase alphanumeric characters and underscores.
986
-
987
- branch : Optional[str], default None
988
- The branch to use. If not specified, the branch is set to
989
- `user.<username>` unless `production` is set to `True`. This can
990
- also be set on the command line using `--branch` as a top-level option.
991
- It is an error to specify `branch` in the decorator and on the command line.
992
-
993
- production : bool, default False
994
- Whether or not the branch is the production branch. This can also be set on the
995
- command line using `--production` as a top-level option. It is an error to specify
996
- `production` in the decorator and on the command line.
997
- The project branch name will be:
998
- - if `branch` is specified:
999
- - if `production` is True: `prod.<branch>`
1000
- - if `production` is False: `test.<branch>`
1001
- - if `branch` is not specified:
1002
- - if `production` is True: `prod`
1003
- - if `production` is False: `user.<username>`
1004
- """
1005
- ...
1006
-
1007
- @typing.overload
1008
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1009
- """
1010
- Specifies the flow(s) that this flow depends on.
1011
-
1012
- ```
1013
- @trigger_on_finish(flow='FooFlow')
1014
- ```
1015
- or
1016
- ```
1017
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1018
- ```
1019
- This decorator respects the @project decorator and triggers the flow
1020
- when upstream runs within the same namespace complete successfully
1021
-
1022
- Additionally, you can specify project aware upstream flow dependencies
1023
- by specifying the fully qualified project_flow_name.
1024
- ```
1025
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1026
- ```
1027
- or
1028
- ```
1029
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1030
- ```
1031
-
1032
- You can also specify just the project or project branch (other values will be
1033
- inferred from the current project or project branch):
1034
- ```
1035
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1036
- ```
1037
-
1038
- Note that `branch` is typically one of:
1039
- - `prod`
1040
- - `user.bob`
1041
- - `test.my_experiment`
1042
- - `prod.staging`
1043
-
1044
-
1045
- Parameters
1046
- ----------
1047
- flow : Union[str, Dict[str, str]], optional, default None
1048
- Upstream flow dependency for this flow.
1049
- flows : List[Union[str, Dict[str, str]]], default []
1050
- Upstream flow dependencies for this flow.
1051
- options : Dict[str, Any], default {}
1052
- Backend-specific configuration for tuning eventing behavior.
1053
- """
1054
- ...
1055
-
1056
- @typing.overload
1057
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1058
- ...
1059
-
1060
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1061
- """
1062
- Specifies the flow(s) that this flow depends on.
1063
-
1064
- ```
1065
- @trigger_on_finish(flow='FooFlow')
1066
- ```
1067
- or
1068
- ```
1069
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1070
- ```
1071
- This decorator respects the @project decorator and triggers the flow
1072
- when upstream runs within the same namespace complete successfully
1073
-
1074
- Additionally, you can specify project aware upstream flow dependencies
1075
- by specifying the fully qualified project_flow_name.
1076
- ```
1077
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1078
- ```
1079
- or
1080
- ```
1081
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1082
- ```
1083
-
1084
- You can also specify just the project or project branch (other values will be
1085
- inferred from the current project or project branch):
1086
- ```
1087
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1088
- ```
1089
-
1090
- Note that `branch` is typically one of:
1091
- - `prod`
1092
- - `user.bob`
1093
- - `test.my_experiment`
1094
- - `prod.staging`
1095
-
1096
-
1097
- Parameters
1098
- ----------
1099
- flow : Union[str, Dict[str, str]], optional, default None
1100
- Upstream flow dependency for this flow.
1101
- flows : List[Union[str, Dict[str, str]]], default []
1102
- Upstream flow dependencies for this flow.
1103
- options : Dict[str, Any], default {}
1104
- Backend-specific configuration for tuning eventing behavior.
1105
- """
1106
- ...
1107
-
1108
- 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]]:
1109
- """
1110
- 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)
1111
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1112
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1113
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1114
- starts only after all sensors finish.
1115
-
1116
-
1117
- Parameters
1118
- ----------
1119
- timeout : int
1120
- Time, in seconds before the task times out and fails. (Default: 3600)
1121
- poke_interval : int
1122
- Time in seconds that the job should wait in between each try. (Default: 60)
1123
- mode : str
1124
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1125
- exponential_backoff : bool
1126
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1127
- pool : str
1128
- the slot pool this task should run in,
1129
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1130
- soft_fail : bool
1131
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1132
- name : str
1133
- Name of the sensor on Airflow
1134
- description : str
1135
- Description of sensor in the Airflow UI
1136
- bucket_key : Union[str, List[str]]
1137
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1138
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1139
- bucket_name : str
1140
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1141
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1142
- wildcard_match : bool
1143
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1144
- aws_conn_id : str
1145
- a reference to the s3 connection on Airflow. (Default: None)
1146
- verify : bool
1147
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1148
- """
1149
- ...
1150
-
1151
- @typing.overload
1152
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1153
- """
1154
- Specifies the Conda environment for all steps of the flow.
1155
-
1156
- Use `@conda_base` to set common libraries required by all
1157
- steps and use `@conda` to specify step-specific additions.
1158
-
1159
-
1160
- Parameters
1161
- ----------
1162
- packages : Dict[str, str], default {}
1163
- Packages to use for this flow. The key is the name of the package
1164
- and the value is the version to use.
1165
- libraries : Dict[str, str], default {}
1166
- Supported for backward compatibility. When used with packages, packages will take precedence.
1167
- python : str, optional, default None
1168
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1169
- that the version used will correspond to the version of the Python interpreter used to start the run.
1170
- disabled : bool, default False
1171
- If set to True, disables Conda.
1172
- """
1173
- ...
1174
-
1175
- @typing.overload
1176
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1177
- ...
1178
-
1179
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1180
- """
1181
- Specifies the Conda environment for all steps of the flow.
1182
-
1183
- Use `@conda_base` to set common libraries required by all
1184
- steps and use `@conda` to specify step-specific additions.
1185
-
1186
-
1187
- Parameters
1188
- ----------
1189
- packages : Dict[str, str], default {}
1190
- Packages to use for this flow. The key is the name of the package
1191
- and the value is the version to use.
1192
- libraries : Dict[str, str], default {}
1193
- Supported for backward compatibility. When used with packages, packages will take precedence.
1194
- python : str, optional, default None
1195
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1196
- that the version used will correspond to the version of the Python interpreter used to start the run.
1197
- disabled : bool, default False
1198
- If set to True, disables Conda.
1199
- """
1200
- ...
1201
-
1202
- 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]]:
1203
- """
1204
- 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.
1205
- This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1206
-
1207
-
1208
- Parameters
1209
- ----------
1210
- timeout : int
1211
- Time, in seconds before the task times out and fails. (Default: 3600)
1212
- poke_interval : int
1213
- Time in seconds that the job should wait in between each try. (Default: 60)
1214
- mode : str
1215
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1216
- exponential_backoff : bool
1217
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1218
- pool : str
1219
- the slot pool this task should run in,
1220
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1221
- soft_fail : bool
1222
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1223
- name : str
1224
- Name of the sensor on Airflow
1225
- description : str
1226
- Description of sensor in the Airflow UI
1227
- external_dag_id : str
1228
- The dag_id that contains the task you want to wait for.
1229
- external_task_ids : List[str]
1230
- The list of task_ids that you want to wait for.
1231
- If None (default value) the sensor waits for the DAG. (Default: None)
1232
- allowed_states : List[str]
1233
- Iterable of allowed states, (Default: ['success'])
1234
- failed_states : List[str]
1235
- Iterable of failed or dis-allowed states. (Default: None)
1236
- execution_delta : datetime.timedelta
1237
- time difference with the previous execution to look at,
1238
- the default is the same logical date as the current task or DAG. (Default: None)
1239
- check_existence: bool
1240
- Set to True to check if the external task exists or check if
1241
- the DAG to wait for exists. (Default: True)
1242
- """
992
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1243
993
  ...
1244
994
 
1245
- @typing.overload
1246
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
995
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1247
996
  """
1248
997
  Specifies the PyPI packages for all steps of the flow.
1249
998
 
@@ -1261,25 +1010,46 @@ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[s
1261
1010
  """
1262
1011
  ...
1263
1012
 
1264
- @typing.overload
1265
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1266
- ...
1267
-
1268
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
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]]:
1269
1014
  """
1270
- Specifies the PyPI packages for all steps of the flow.
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.
1271
1020
 
1272
- Use `@pypi_base` to set common packages required by all
1273
- steps and use `@pypi` to specify step-specific overrides.
1274
1021
 
1275
1022
  Parameters
1276
1023
  ----------
1277
- packages : Dict[str, str], default: {}
1278
- Packages to use for this flow. The key is the name of the package
1279
- and the value is the version to use.
1280
- python : str, optional, default: None
1281
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1282
- that the version used will correspond to the version of the Python interpreter used to start the run.
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)
1283
1053
  """
1284
1054
  ...
1285
1055
 
@@ -1376,6 +1146,57 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1376
1146
  """
1377
1147
  ...
1378
1148
 
1149
+ @typing.overload
1150
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1151
+ """
1152
+ Specifies the Conda environment for all steps of the flow.
1153
+
1154
+ Use `@conda_base` to set common libraries required by all
1155
+ steps and use `@conda` to specify step-specific additions.
1156
+
1157
+
1158
+ Parameters
1159
+ ----------
1160
+ packages : Dict[str, str], default {}
1161
+ Packages to use for this flow. The key is the name of the package
1162
+ and the value is the version to use.
1163
+ libraries : Dict[str, str], default {}
1164
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1165
+ python : str, optional, default None
1166
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1167
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1168
+ disabled : bool, default False
1169
+ If set to True, disables Conda.
1170
+ """
1171
+ ...
1172
+
1173
+ @typing.overload
1174
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1175
+ ...
1176
+
1177
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1178
+ """
1179
+ Specifies the Conda environment for all steps of the flow.
1180
+
1181
+ Use `@conda_base` to set common libraries required by all
1182
+ steps and use `@conda` to specify step-specific additions.
1183
+
1184
+
1185
+ Parameters
1186
+ ----------
1187
+ packages : Dict[str, str], default {}
1188
+ Packages to use for this flow. The key is the name of the package
1189
+ and the value is the version to use.
1190
+ libraries : Dict[str, str], default {}
1191
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1192
+ python : str, optional, default None
1193
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1194
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1195
+ disabled : bool, default False
1196
+ If set to True, disables Conda.
1197
+ """
1198
+ ...
1199
+
1379
1200
  @typing.overload
1380
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]]:
1381
1202
  """
@@ -1427,5 +1248,184 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1427
1248
  """
1428
1249
  ...
1429
1250
 
1251
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1252
+ """
1253
+ Specifies what flows belong to the same project.
1254
+
1255
+ A project-specific namespace is created for all flows that
1256
+ use the same `@project(name)`.
1257
+
1258
+
1259
+ Parameters
1260
+ ----------
1261
+ name : str
1262
+ Project name. Make sure that the name is unique amongst all
1263
+ projects that use the same production scheduler. The name may
1264
+ contain only lowercase alphanumeric characters and underscores.
1265
+
1266
+ branch : Optional[str], default None
1267
+ The branch to use. If not specified, the branch is set to
1268
+ `user.<username>` unless `production` is set to `True`. This can
1269
+ also be set on the command line using `--branch` as a top-level option.
1270
+ It is an error to specify `branch` in the decorator and on the command line.
1271
+
1272
+ production : bool, default False
1273
+ Whether or not the branch is the production branch. This can also be set on the
1274
+ command line using `--production` as a top-level option. It is an error to specify
1275
+ `production` in the decorator and on the command line.
1276
+ The project branch name will be:
1277
+ - if `branch` is specified:
1278
+ - if `production` is True: `prod.<branch>`
1279
+ - if `production` is False: `test.<branch>`
1280
+ - if `branch` is not specified:
1281
+ - if `production` is True: `prod`
1282
+ - if `production` is False: `user.<username>`
1283
+ """
1284
+ ...
1285
+
1286
+ @typing.overload
1287
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1288
+ """
1289
+ Specifies the flow(s) that this flow depends on.
1290
+
1291
+ ```
1292
+ @trigger_on_finish(flow='FooFlow')
1293
+ ```
1294
+ or
1295
+ ```
1296
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1297
+ ```
1298
+ This decorator respects the @project decorator and triggers the flow
1299
+ when upstream runs within the same namespace complete successfully
1300
+
1301
+ Additionally, you can specify project aware upstream flow dependencies
1302
+ by specifying the fully qualified project_flow_name.
1303
+ ```
1304
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1305
+ ```
1306
+ or
1307
+ ```
1308
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1309
+ ```
1310
+
1311
+ You can also specify just the project or project branch (other values will be
1312
+ inferred from the current project or project branch):
1313
+ ```
1314
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1315
+ ```
1316
+
1317
+ Note that `branch` is typically one of:
1318
+ - `prod`
1319
+ - `user.bob`
1320
+ - `test.my_experiment`
1321
+ - `prod.staging`
1322
+
1323
+
1324
+ Parameters
1325
+ ----------
1326
+ flow : Union[str, Dict[str, str]], optional, default None
1327
+ Upstream flow dependency for this flow.
1328
+ flows : List[Union[str, Dict[str, str]]], default []
1329
+ Upstream flow dependencies for this flow.
1330
+ options : Dict[str, Any], default {}
1331
+ Backend-specific configuration for tuning eventing behavior.
1332
+ """
1333
+ ...
1334
+
1335
+ @typing.overload
1336
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1337
+ ...
1338
+
1339
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
1340
+ """
1341
+ Specifies the flow(s) that this flow depends on.
1342
+
1343
+ ```
1344
+ @trigger_on_finish(flow='FooFlow')
1345
+ ```
1346
+ or
1347
+ ```
1348
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
1349
+ ```
1350
+ This decorator respects the @project decorator and triggers the flow
1351
+ when upstream runs within the same namespace complete successfully
1352
+
1353
+ Additionally, you can specify project aware upstream flow dependencies
1354
+ by specifying the fully qualified project_flow_name.
1355
+ ```
1356
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
1357
+ ```
1358
+ or
1359
+ ```
1360
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
1361
+ ```
1362
+
1363
+ You can also specify just the project or project branch (other values will be
1364
+ inferred from the current project or project branch):
1365
+ ```
1366
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
1367
+ ```
1368
+
1369
+ Note that `branch` is typically one of:
1370
+ - `prod`
1371
+ - `user.bob`
1372
+ - `test.my_experiment`
1373
+ - `prod.staging`
1374
+
1375
+
1376
+ Parameters
1377
+ ----------
1378
+ flow : Union[str, Dict[str, str]], optional, default None
1379
+ Upstream flow dependency for this flow.
1380
+ flows : List[Union[str, Dict[str, str]]], default []
1381
+ Upstream flow dependencies for this flow.
1382
+ options : Dict[str, Any], default {}
1383
+ Backend-specific configuration for tuning eventing behavior.
1384
+ """
1385
+ ...
1386
+
1387
+ 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
+ """
1389
+ 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.
1390
+ This decorator only works when a flow is scheduled on Airflow and is compiled using `airflow create`. More than one `@airflow_external_task_sensor` can be added as a flow decorators. Adding more than one decorator will ensure that `start` step starts only after all sensors finish.
1391
+
1392
+
1393
+ Parameters
1394
+ ----------
1395
+ timeout : int
1396
+ Time, in seconds before the task times out and fails. (Default: 3600)
1397
+ poke_interval : int
1398
+ Time in seconds that the job should wait in between each try. (Default: 60)
1399
+ mode : str
1400
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1401
+ exponential_backoff : bool
1402
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1403
+ pool : str
1404
+ the slot pool this task should run in,
1405
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1406
+ soft_fail : bool
1407
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1408
+ name : str
1409
+ Name of the sensor on Airflow
1410
+ description : str
1411
+ Description of sensor in the Airflow UI
1412
+ external_dag_id : str
1413
+ The dag_id that contains the task you want to wait for.
1414
+ external_task_ids : List[str]
1415
+ The list of task_ids that you want to wait for.
1416
+ If None (default value) the sensor waits for the DAG. (Default: None)
1417
+ allowed_states : List[str]
1418
+ Iterable of allowed states, (Default: ['success'])
1419
+ failed_states : List[str]
1420
+ Iterable of failed or dis-allowed states. (Default: None)
1421
+ execution_delta : datetime.timedelta
1422
+ time difference with the previous execution to look at,
1423
+ the default is the same logical date as the current task or DAG. (Default: None)
1424
+ check_existence: bool
1425
+ Set to True to check if the external task exists or check if
1426
+ the DAG to wait for exists. (Default: True)
1427
+ """
1428
+ ...
1429
+
1430
1430
  pkg_name: str
1431
1431