ob-metaflow-stubs 6.0.3.141__py2.py3-none-any.whl → 6.0.3.142__py2.py3-none-any.whl

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