metaflow-stubs 2.18.5__py2.py3-none-any.whl → 2.18.7__py2.py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of metaflow-stubs might be problematic. Click here for more details.

Files changed (166) hide show
  1. metaflow-stubs/__init__.pyi +544 -544
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +3 -3
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +6 -6
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/meta_files.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +3 -3
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +2 -2
  20. metaflow-stubs/metaflow_current.pyi +36 -36
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/multicore_utils.pyi +2 -2
  23. metaflow-stubs/packaging_sys/__init__.pyi +6 -6
  24. metaflow-stubs/packaging_sys/backend.pyi +4 -4
  25. metaflow-stubs/packaging_sys/distribution_support.pyi +3 -3
  26. metaflow-stubs/packaging_sys/tar_backend.pyi +7 -7
  27. metaflow-stubs/packaging_sys/utils.pyi +2 -2
  28. metaflow-stubs/packaging_sys/v1.pyi +3 -3
  29. metaflow-stubs/parameters.pyi +3 -3
  30. metaflow-stubs/plugins/__init__.pyi +13 -13
  31. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  32. metaflow-stubs/plugins/airflow/airflow_decorator.pyi +2 -2
  33. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  34. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  35. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  36. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  37. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  38. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  39. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  40. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  41. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  42. metaflow-stubs/plugins/argo/argo_workflows.pyi +3 -3
  43. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +2 -2
  44. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  45. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +3 -3
  46. metaflow-stubs/plugins/argo/exit_hooks.pyi +2 -2
  47. metaflow-stubs/plugins/aws/__init__.pyi +3 -3
  48. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  49. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  50. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  51. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  52. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  53. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  54. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  55. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  56. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  57. metaflow-stubs/plugins/aws/step_functions/dynamo_db_client.pyi +2 -2
  58. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  59. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  60. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  61. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  62. metaflow-stubs/plugins/aws/step_functions/step_functions_decorator.pyi +2 -2
  63. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +3 -3
  64. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +4 -4
  65. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  66. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  67. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  68. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  69. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  70. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  71. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  72. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  73. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  74. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  75. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  76. metaflow-stubs/plugins/cards/card_decorator.pyi +2 -2
  77. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  78. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  79. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  80. metaflow-stubs/plugins/cards/card_modules/components.pyi +5 -5
  81. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  82. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  83. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  84. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  85. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  86. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  87. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  88. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  89. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  90. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  91. metaflow-stubs/plugins/datatools/s3/s3.pyi +3 -3
  92. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  93. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  94. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  95. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  96. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  97. metaflow-stubs/plugins/events_decorator.pyi +45 -4
  98. metaflow-stubs/plugins/exit_hook/__init__.pyi +2 -2
  99. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +2 -2
  100. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  101. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  102. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  103. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  104. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  105. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  106. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  107. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  108. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  109. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  110. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  111. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  112. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  113. metaflow-stubs/plugins/kubernetes/kubernetes_job.pyi +2 -2
  114. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  115. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  116. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  117. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  118. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  119. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/pypi/conda_environment.pyi +3 -3
  121. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  122. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  123. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  124. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  125. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  126. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  127. metaflow-stubs/plugins/secrets/__init__.pyi +2 -2
  128. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  129. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  130. metaflow-stubs/plugins/secrets/secrets_func.pyi +2 -2
  131. metaflow-stubs/plugins/secrets/secrets_spec.pyi +2 -2
  132. metaflow-stubs/plugins/secrets/utils.pyi +2 -2
  133. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  134. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +3 -3
  135. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  136. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  137. metaflow-stubs/plugins/uv/uv_environment.pyi +2 -2
  138. metaflow-stubs/pylint_wrapper.pyi +2 -2
  139. metaflow-stubs/runner/__init__.pyi +2 -2
  140. metaflow-stubs/runner/deployer.pyi +5 -5
  141. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  142. metaflow-stubs/runner/metaflow_runner.pyi +3 -3
  143. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  144. metaflow-stubs/runner/nbrun.pyi +2 -2
  145. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  146. metaflow-stubs/runner/utils.pyi +4 -4
  147. metaflow-stubs/system/__init__.pyi +2 -2
  148. metaflow-stubs/system/system_logger.pyi +2 -2
  149. metaflow-stubs/system/system_monitor.pyi +2 -2
  150. metaflow-stubs/tagging_util.pyi +2 -2
  151. metaflow-stubs/tuple_util.pyi +2 -2
  152. metaflow-stubs/user_configs/__init__.pyi +2 -2
  153. metaflow-stubs/user_configs/config_options.pyi +3 -3
  154. metaflow-stubs/user_configs/config_parameters.pyi +4 -4
  155. metaflow-stubs/user_decorators/__init__.pyi +2 -2
  156. metaflow-stubs/user_decorators/common.pyi +2 -2
  157. metaflow-stubs/user_decorators/mutable_flow.pyi +5 -5
  158. metaflow-stubs/user_decorators/mutable_step.pyi +4 -4
  159. metaflow-stubs/user_decorators/user_flow_decorator.pyi +4 -4
  160. metaflow-stubs/user_decorators/user_step_decorator.pyi +4 -4
  161. metaflow-stubs/version.pyi +2 -2
  162. {metaflow_stubs-2.18.5.dist-info → metaflow_stubs-2.18.7.dist-info}/METADATA +2 -2
  163. metaflow_stubs-2.18.7.dist-info/RECORD +166 -0
  164. metaflow_stubs-2.18.5.dist-info/RECORD +0 -166
  165. {metaflow_stubs-2.18.5.dist-info → metaflow_stubs-2.18.7.dist-info}/WHEEL +0 -0
  166. {metaflow_stubs-2.18.5.dist-info → metaflow_stubs-2.18.7.dist-info}/top_level.txt +0 -0
@@ -1,7 +1,7 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.18.5 #
4
- # Generated on 2025-09-16T00:24:37.068070 #
3
+ # MF version: 2.18.7 #
4
+ # Generated on 2025-09-18T19:24:16.691351 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
@@ -39,17 +39,17 @@ from .user_decorators.user_step_decorator import UserStepDecorator as UserStepDe
39
39
  from .user_decorators.user_step_decorator import StepMutator as StepMutator
40
40
  from .user_decorators.user_step_decorator import user_step_decorator as user_step_decorator
41
41
  from .user_decorators.user_flow_decorator import FlowMutator as FlowMutator
42
- from . import metaflow_git as metaflow_git
43
42
  from . import events as events
43
+ from . import metaflow_git as metaflow_git
44
44
  from . import tuple_util as tuple_util
45
45
  from . import runner as runner
46
46
  from . import plugins as plugins
47
47
  from .plugins.datatools.s3.s3 import S3 as S3
48
48
  from . import includefile as includefile
49
49
  from .includefile import IncludeFile as IncludeFile
50
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
51
50
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
52
51
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
52
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
53
53
  from . import cards as cards
54
54
  from . import client as client
55
55
  from .client.core import namespace as namespace
@@ -153,36 +153,66 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
153
153
  ...
154
154
 
155
155
  @typing.overload
156
- def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
156
+ 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]]]:
157
157
  """
158
- Specifies that the step will success under all circumstances.
158
+ Specifies a timeout for your step.
159
159
 
160
- The decorator will create an optional artifact, specified by `var`, which
161
- contains the exception raised. You can use it to detect the presence
162
- of errors, indicating that all happy-path artifacts produced by the step
163
- are missing.
160
+ This decorator is useful if this step may hang indefinitely.
161
+
162
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
163
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
164
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
165
+
166
+ Note that all the values specified in parameters are added together so if you specify
167
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
164
168
 
165
169
 
166
170
  Parameters
167
171
  ----------
168
- var : str, optional, default None
169
- Name of the artifact in which to store the caught exception.
170
- If not specified, the exception is not stored.
171
- print_exception : bool, default True
172
- Determines whether or not the exception is printed to
173
- stdout when caught.
172
+ seconds : int, default 0
173
+ Number of seconds to wait prior to timing out.
174
+ minutes : int, default 0
175
+ Number of minutes to wait prior to timing out.
176
+ hours : int, default 0
177
+ Number of hours to wait prior to timing out.
174
178
  """
175
179
  ...
176
180
 
177
181
  @typing.overload
178
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
182
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
179
183
  ...
180
184
 
181
185
  @typing.overload
182
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
186
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
183
187
  ...
184
188
 
185
- def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
189
+ 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):
190
+ """
191
+ Specifies a timeout for your step.
192
+
193
+ This decorator is useful if this step may hang indefinitely.
194
+
195
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
196
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
197
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
198
+
199
+ Note that all the values specified in parameters are added together so if you specify
200
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
201
+
202
+
203
+ Parameters
204
+ ----------
205
+ seconds : int, default 0
206
+ Number of seconds to wait prior to timing out.
207
+ minutes : int, default 0
208
+ Number of minutes to wait prior to timing out.
209
+ hours : int, default 0
210
+ Number of hours to wait prior to timing out.
211
+ """
212
+ ...
213
+
214
+ @typing.overload
215
+ def catch(*, var: typing.Optional[str] = None, print_exception: bool = True) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
186
216
  """
187
217
  Specifies that the step will success under all circumstances.
188
218
 
@@ -204,61 +234,31 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
204
234
  ...
205
235
 
206
236
  @typing.overload
207
- 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]]]:
208
- """
209
- Specifies the Conda environment for the step.
210
-
211
- Information in this decorator will augment any
212
- attributes set in the `@conda_base` flow-level decorator. Hence,
213
- you can use `@conda_base` to set packages required by all
214
- steps and use `@conda` to specify step-specific overrides.
215
-
216
-
217
- Parameters
218
- ----------
219
- packages : Dict[str, str], default {}
220
- Packages to use for this step. The key is the name of the package
221
- and the value is the version to use.
222
- libraries : Dict[str, str], default {}
223
- Supported for backward compatibility. When used with packages, packages will take precedence.
224
- python : str, optional, default None
225
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
226
- that the version used will correspond to the version of the Python interpreter used to start the run.
227
- disabled : bool, default False
228
- If set to True, disables @conda.
229
- """
230
- ...
231
-
232
- @typing.overload
233
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
237
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
234
238
  ...
235
239
 
236
240
  @typing.overload
237
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
241
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
238
242
  ...
239
243
 
240
- 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):
244
+ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, var: typing.Optional[str] = None, print_exception: bool = True):
241
245
  """
242
- Specifies the Conda environment for the step.
246
+ Specifies that the step will success under all circumstances.
243
247
 
244
- Information in this decorator will augment any
245
- attributes set in the `@conda_base` flow-level decorator. Hence,
246
- you can use `@conda_base` to set packages required by all
247
- steps and use `@conda` to specify step-specific overrides.
248
+ The decorator will create an optional artifact, specified by `var`, which
249
+ contains the exception raised. You can use it to detect the presence
250
+ of errors, indicating that all happy-path artifacts produced by the step
251
+ are missing.
248
252
 
249
253
 
250
254
  Parameters
251
255
  ----------
252
- packages : Dict[str, str], default {}
253
- Packages to use for this step. The key is the name of the package
254
- and the value is the version to use.
255
- libraries : Dict[str, str], default {}
256
- Supported for backward compatibility. When used with packages, packages will take precedence.
257
- python : str, optional, default None
258
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
259
- that the version used will correspond to the version of the Python interpreter used to start the run.
260
- disabled : bool, default False
261
- If set to True, disables @conda.
256
+ var : str, optional, default None
257
+ Name of the artifact in which to store the caught exception.
258
+ If not specified, the exception is not stored.
259
+ print_exception : bool, default True
260
+ Determines whether or not the exception is printed to
261
+ stdout when caught.
262
262
  """
263
263
  ...
264
264
 
@@ -301,25 +301,6 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
301
301
  """
302
302
  ...
303
303
 
304
- @typing.overload
305
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
306
- """
307
- Decorator prototype for all step decorators. This function gets specialized
308
- and imported for all decorators types by _import_plugin_decorators().
309
- """
310
- ...
311
-
312
- @typing.overload
313
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
314
- ...
315
-
316
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
317
- """
318
- Decorator prototype for all step decorators. This function gets specialized
319
- and imported for all decorators types by _import_plugin_decorators().
320
- """
321
- ...
322
-
323
304
  @typing.overload
324
305
  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]]]:
325
306
  """
@@ -375,10 +356,9 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
375
356
  """
376
357
  ...
377
358
 
378
- @typing.overload
379
- def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optional[str] = None, queue: str = 'METAFLOW_BATCH_JOB_QUEUE', iam_role: str = 'METAFLOW_ECS_S3_ACCESS_IAM_ROLE', execution_role: str = 'METAFLOW_ECS_FARGATE_EXECUTION_ROLE', shared_memory: typing.Optional[int] = None, max_swap: typing.Optional[int] = None, swappiness: typing.Optional[int] = None, aws_batch_tags: typing.Optional[typing.Dict[str, str]] = None, use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = None, inferentia: int = 0, trainium: int = None, efa: int = 0, ephemeral_storage: int = None, log_driver: typing.Optional[str] = None, log_options: typing.Optional[typing.List[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]]]:
359
+ def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
380
360
  """
381
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
361
+ Specifies that this step should execute on Kubernetes.
382
362
 
383
363
 
384
364
  Parameters
@@ -386,89 +366,190 @@ def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optio
386
366
  cpu : int, default 1
387
367
  Number of CPUs required for this step. If `@resources` is
388
368
  also present, the maximum value from all decorators is used.
389
- gpu : int, default 0
390
- Number of GPUs required for this step. If `@resources` is
391
- also present, the maximum value from all decorators is used.
392
369
  memory : int, default 4096
393
370
  Memory size (in MB) required for this step. If
394
371
  `@resources` is also present, the maximum value from all decorators is
395
372
  used.
373
+ disk : int, default 10240
374
+ Disk size (in MB) required for this step. If
375
+ `@resources` is also present, the maximum value from all decorators is
376
+ used.
396
377
  image : str, optional, default None
397
- Docker image to use when launching on AWS Batch. If not specified, and
398
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
378
+ Docker image to use when launching on Kubernetes. If not specified, and
379
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
399
380
  not, a default Docker image mapping to the current version of Python is used.
400
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
401
- AWS Batch Job Queue to submit the job to.
402
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
403
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
404
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
405
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
406
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
407
- shared_memory : int, optional, default None
408
- The value for the size (in MiB) of the /dev/shm volume for this step.
409
- This parameter maps to the `--shm-size` option in Docker.
410
- max_swap : int, optional, default None
411
- The total amount of swap memory (in MiB) a container can use for this
412
- step. This parameter is translated to the `--memory-swap` option in
413
- Docker where the value is the sum of the container memory plus the
414
- `max_swap` value.
415
- swappiness : int, optional, default None
416
- This allows you to tune memory swappiness behavior for this step.
417
- A swappiness value of 0 causes swapping not to happen unless absolutely
418
- necessary. A swappiness value of 100 causes pages to be swapped very
419
- aggressively. Accepted values are whole numbers between 0 and 100.
420
- aws_batch_tags: Dict[str, str], optional, default None
421
- Sets arbitrary AWS tags on the AWS Batch compute environment.
422
- Set as string key-value pairs.
381
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
382
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
383
+ image_pull_secrets: List[str], default []
384
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
385
+ Kubernetes image pull secrets to use when pulling container images
386
+ in Kubernetes.
387
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
388
+ Kubernetes service account to use when launching pod in Kubernetes.
389
+ secrets : List[str], optional, default None
390
+ Kubernetes secrets to use when launching pod in Kubernetes. These
391
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
392
+ in Metaflow configuration.
393
+ node_selector: Union[Dict[str,str], str], optional, default None
394
+ Kubernetes node selector(s) to apply to the pod running the task.
395
+ Can be passed in as a comma separated string of values e.g.
396
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
397
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
398
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
399
+ Kubernetes namespace to use when launching pod in Kubernetes.
400
+ gpu : int, optional, default None
401
+ Number of GPUs required for this step. A value of zero implies that
402
+ the scheduled node should not have GPUs.
403
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
404
+ The vendor of the GPUs to be used for this step.
405
+ tolerations : List[Dict[str,str]], default []
406
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
407
+ Kubernetes tolerations to use when launching pod in Kubernetes.
408
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
409
+ Kubernetes labels to use when launching pod in Kubernetes.
410
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
411
+ Kubernetes annotations to use when launching pod in Kubernetes.
423
412
  use_tmpfs : bool, default False
424
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
425
- not available on Fargate compute environments
413
+ This enables an explicit tmpfs mount for this step.
426
414
  tmpfs_tempdir : bool, default True
427
415
  sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
428
- tmpfs_size : int, optional, default None
416
+ tmpfs_size : int, optional, default: None
429
417
  The value for the size (in MiB) of the tmpfs mount for this step.
430
418
  This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
431
419
  memory allocated for this step.
432
- tmpfs_path : str, optional, default None
433
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
434
- inferentia : int, default 0
435
- Number of Inferentia chips required for this step.
436
- trainium : int, default None
437
- Alias for inferentia. Use only one of the two.
438
- efa : int, default 0
439
- Number of elastic fabric adapter network devices to attach to container
440
- ephemeral_storage : int, default None
441
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
442
- This is only relevant for Fargate compute environments
443
- log_driver: str, optional, default None
444
- The log driver to use for the Amazon ECS container.
445
- log_options: List[str], optional, default None
446
- List of strings containing options for the chosen log driver. The configurable values
447
- depend on the `log driver` chosen. Validation of these options is not supported yet.
448
- Example: [`awslogs-group:aws/batch/job`]
449
- """
450
- ...
451
-
452
- @typing.overload
453
- def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
454
- ...
455
-
456
- @typing.overload
457
- def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
458
- ...
459
-
460
- def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optional[str] = None, queue: str = 'METAFLOW_BATCH_JOB_QUEUE', iam_role: str = 'METAFLOW_ECS_S3_ACCESS_IAM_ROLE', execution_role: str = 'METAFLOW_ECS_FARGATE_EXECUTION_ROLE', shared_memory: typing.Optional[int] = None, max_swap: typing.Optional[int] = None, swappiness: typing.Optional[int] = None, aws_batch_tags: typing.Optional[typing.Dict[str, str]] = None, use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = None, inferentia: int = 0, trainium: int = None, efa: int = 0, ephemeral_storage: int = None, log_driver: typing.Optional[str] = None, log_options: typing.Optional[typing.List[str]] = None):
461
- """
462
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
463
-
464
-
465
- Parameters
466
- ----------
467
- cpu : int, default 1
468
- Number of CPUs required for this step. If `@resources` is
469
- also present, the maximum value from all decorators is used.
470
- gpu : int, default 0
471
- Number of GPUs required for this step. If `@resources` is
420
+ tmpfs_path : str, optional, default /metaflow_temp
421
+ Path to tmpfs mount for this step.
422
+ persistent_volume_claims : Dict[str, str], optional, default None
423
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
424
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
425
+ shared_memory: int, optional
426
+ Shared memory size (in MiB) required for this step
427
+ port: int, optional
428
+ Port number to specify in the Kubernetes job object
429
+ compute_pool : str, optional, default None
430
+ Compute pool to be used for for this step.
431
+ If not specified, any accessible compute pool within the perimeter is used.
432
+ hostname_resolution_timeout: int, default 10 * 60
433
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
434
+ Only applicable when @parallel is used.
435
+ qos: str, default: Burstable
436
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
437
+
438
+ security_context: Dict[str, Any], optional, default None
439
+ Container security context. Applies to the task container. Allows the following keys:
440
+ - privileged: bool, optional, default None
441
+ - allow_privilege_escalation: bool, optional, default None
442
+ - run_as_user: int, optional, default None
443
+ - run_as_group: int, optional, default None
444
+ - run_as_non_root: bool, optional, default None
445
+ """
446
+ ...
447
+
448
+ @typing.overload
449
+ 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]]]:
450
+ """
451
+ Specifies the Conda environment for the step.
452
+
453
+ Information in this decorator will augment any
454
+ attributes set in the `@conda_base` flow-level decorator. Hence,
455
+ you can use `@conda_base` to set packages required by all
456
+ steps and use `@conda` to specify step-specific overrides.
457
+
458
+
459
+ Parameters
460
+ ----------
461
+ packages : Dict[str, str], default {}
462
+ Packages to use for this step. The key is the name of the package
463
+ and the value is the version to use.
464
+ libraries : Dict[str, str], default {}
465
+ Supported for backward compatibility. When used with packages, packages will take precedence.
466
+ python : str, optional, default None
467
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
468
+ that the version used will correspond to the version of the Python interpreter used to start the run.
469
+ disabled : bool, default False
470
+ If set to True, disables @conda.
471
+ """
472
+ ...
473
+
474
+ @typing.overload
475
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
476
+ ...
477
+
478
+ @typing.overload
479
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
480
+ ...
481
+
482
+ 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):
483
+ """
484
+ Specifies the Conda environment for the step.
485
+
486
+ Information in this decorator will augment any
487
+ attributes set in the `@conda_base` flow-level decorator. Hence,
488
+ you can use `@conda_base` to set packages required by all
489
+ steps and use `@conda` to specify step-specific overrides.
490
+
491
+
492
+ Parameters
493
+ ----------
494
+ packages : Dict[str, str], default {}
495
+ Packages to use for this step. The key is the name of the package
496
+ and the value is the version to use.
497
+ libraries : Dict[str, str], default {}
498
+ Supported for backward compatibility. When used with packages, packages will take precedence.
499
+ python : str, optional, default None
500
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
501
+ that the version used will correspond to the version of the Python interpreter used to start the run.
502
+ disabled : bool, default False
503
+ If set to True, disables @conda.
504
+ """
505
+ ...
506
+
507
+ @typing.overload
508
+ def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
509
+ """
510
+ Specifies environment variables to be set prior to the execution of a step.
511
+
512
+
513
+ Parameters
514
+ ----------
515
+ vars : Dict[str, str], default {}
516
+ Dictionary of environment variables to set.
517
+ """
518
+ ...
519
+
520
+ @typing.overload
521
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
522
+ ...
523
+
524
+ @typing.overload
525
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
526
+ ...
527
+
528
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
529
+ """
530
+ Specifies environment variables to be set prior to the execution of a step.
531
+
532
+
533
+ Parameters
534
+ ----------
535
+ vars : Dict[str, str], default {}
536
+ Dictionary of environment variables to set.
537
+ """
538
+ ...
539
+
540
+ @typing.overload
541
+ def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optional[str] = None, queue: str = 'METAFLOW_BATCH_JOB_QUEUE', iam_role: str = 'METAFLOW_ECS_S3_ACCESS_IAM_ROLE', execution_role: str = 'METAFLOW_ECS_FARGATE_EXECUTION_ROLE', shared_memory: typing.Optional[int] = None, max_swap: typing.Optional[int] = None, swappiness: typing.Optional[int] = None, aws_batch_tags: typing.Optional[typing.Dict[str, str]] = None, use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = None, inferentia: int = 0, trainium: int = None, efa: int = 0, ephemeral_storage: int = None, log_driver: typing.Optional[str] = None, log_options: typing.Optional[typing.List[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]]]:
542
+ """
543
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
544
+
545
+
546
+ Parameters
547
+ ----------
548
+ cpu : int, default 1
549
+ Number of CPUs required for this step. If `@resources` is
550
+ also present, the maximum value from all decorators is used.
551
+ gpu : int, default 0
552
+ Number of GPUs required for this step. If `@resources` is
472
553
  also present, the maximum value from all decorators is used.
473
554
  memory : int, default 4096
474
555
  Memory size (in MB) required for this step. If
@@ -531,61 +612,83 @@ def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
531
612
  ...
532
613
 
533
614
  @typing.overload
534
- 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]]]:
535
- """
536
- Specifies a timeout for your step.
537
-
538
- This decorator is useful if this step may hang indefinitely.
539
-
540
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
541
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
542
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
543
-
544
- Note that all the values specified in parameters are added together so if you specify
545
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
546
-
547
-
548
- Parameters
549
- ----------
550
- seconds : int, default 0
551
- Number of seconds to wait prior to timing out.
552
- minutes : int, default 0
553
- Number of minutes to wait prior to timing out.
554
- hours : int, default 0
555
- Number of hours to wait prior to timing out.
556
- """
557
- ...
558
-
559
- @typing.overload
560
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
615
+ def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
561
616
  ...
562
617
 
563
618
  @typing.overload
564
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
619
+ def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
565
620
  ...
566
621
 
567
- 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):
622
+ def batch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optional[str] = None, queue: str = 'METAFLOW_BATCH_JOB_QUEUE', iam_role: str = 'METAFLOW_ECS_S3_ACCESS_IAM_ROLE', execution_role: str = 'METAFLOW_ECS_FARGATE_EXECUTION_ROLE', shared_memory: typing.Optional[int] = None, max_swap: typing.Optional[int] = None, swappiness: typing.Optional[int] = None, aws_batch_tags: typing.Optional[typing.Dict[str, str]] = None, use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = None, inferentia: int = 0, trainium: int = None, efa: int = 0, ephemeral_storage: int = None, log_driver: typing.Optional[str] = None, log_options: typing.Optional[typing.List[str]] = None):
568
623
  """
569
- Specifies a timeout for your step.
570
-
571
- This decorator is useful if this step may hang indefinitely.
572
-
573
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
574
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
575
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
576
-
577
- Note that all the values specified in parameters are added together so if you specify
578
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
624
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
579
625
 
580
626
 
581
627
  Parameters
582
628
  ----------
583
- seconds : int, default 0
584
- Number of seconds to wait prior to timing out.
585
- minutes : int, default 0
586
- Number of minutes to wait prior to timing out.
587
- hours : int, default 0
588
- Number of hours to wait prior to timing out.
629
+ cpu : int, default 1
630
+ Number of CPUs required for this step. If `@resources` is
631
+ also present, the maximum value from all decorators is used.
632
+ gpu : int, default 0
633
+ Number of GPUs required for this step. If `@resources` is
634
+ also present, the maximum value from all decorators is used.
635
+ memory : int, default 4096
636
+ Memory size (in MB) required for this step. If
637
+ `@resources` is also present, the maximum value from all decorators is
638
+ used.
639
+ image : str, optional, default None
640
+ Docker image to use when launching on AWS Batch. If not specified, and
641
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
642
+ not, a default Docker image mapping to the current version of Python is used.
643
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
644
+ AWS Batch Job Queue to submit the job to.
645
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
646
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
647
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
648
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
649
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
650
+ shared_memory : int, optional, default None
651
+ The value for the size (in MiB) of the /dev/shm volume for this step.
652
+ This parameter maps to the `--shm-size` option in Docker.
653
+ max_swap : int, optional, default None
654
+ The total amount of swap memory (in MiB) a container can use for this
655
+ step. This parameter is translated to the `--memory-swap` option in
656
+ Docker where the value is the sum of the container memory plus the
657
+ `max_swap` value.
658
+ swappiness : int, optional, default None
659
+ This allows you to tune memory swappiness behavior for this step.
660
+ A swappiness value of 0 causes swapping not to happen unless absolutely
661
+ necessary. A swappiness value of 100 causes pages to be swapped very
662
+ aggressively. Accepted values are whole numbers between 0 and 100.
663
+ aws_batch_tags: Dict[str, str], optional, default None
664
+ Sets arbitrary AWS tags on the AWS Batch compute environment.
665
+ Set as string key-value pairs.
666
+ use_tmpfs : bool, default False
667
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
668
+ not available on Fargate compute environments
669
+ tmpfs_tempdir : bool, default True
670
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
671
+ tmpfs_size : int, optional, default None
672
+ The value for the size (in MiB) of the tmpfs mount for this step.
673
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
674
+ memory allocated for this step.
675
+ tmpfs_path : str, optional, default None
676
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
677
+ inferentia : int, default 0
678
+ Number of Inferentia chips required for this step.
679
+ trainium : int, default None
680
+ Alias for inferentia. Use only one of the two.
681
+ efa : int, default 0
682
+ Number of elastic fabric adapter network devices to attach to container
683
+ ephemeral_storage : int, default None
684
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
685
+ This is only relevant for Fargate compute environments
686
+ log_driver: str, optional, default None
687
+ The log driver to use for the Amazon ECS container.
688
+ log_options: List[str], optional, default None
689
+ List of strings containing options for the chosen log driver. The configurable values
690
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
691
+ Example: [`awslogs-group:aws/batch/job`]
589
692
  """
590
693
  ...
591
694
 
@@ -690,129 +793,69 @@ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typ
690
793
  ...
691
794
 
692
795
  @typing.overload
693
- def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
796
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
694
797
  """
695
- Specifies environment variables to be set prior to the execution of a step.
696
-
697
-
698
- Parameters
699
- ----------
700
- vars : Dict[str, str], default {}
701
- Dictionary of environment variables to set.
798
+ Decorator prototype for all step decorators. This function gets specialized
799
+ and imported for all decorators types by _import_plugin_decorators().
702
800
  """
703
801
  ...
704
802
 
705
803
  @typing.overload
706
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
707
- ...
708
-
709
- @typing.overload
710
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
804
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
711
805
  ...
712
806
 
713
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
807
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
714
808
  """
715
- Specifies environment variables to be set prior to the execution of a step.
716
-
717
-
718
- Parameters
719
- ----------
720
- vars : Dict[str, str], default {}
721
- Dictionary of environment variables to set.
809
+ Decorator prototype for all step decorators. This function gets specialized
810
+ and imported for all decorators types by _import_plugin_decorators().
722
811
  """
723
812
  ...
724
813
 
725
- def kubernetes(*, cpu: int = 1, memory: int = 4096, disk: int = 10240, image: typing.Optional[str] = None, image_pull_policy: str = 'KUBERNETES_IMAGE_PULL_POLICY', image_pull_secrets: typing.List[str] = [], service_account: str = 'METAFLOW_KUBERNETES_SERVICE_ACCOUNT', secrets: typing.Optional[typing.List[str]] = None, node_selector: typing.Union[typing.Dict[str, str], str, None] = None, namespace: str = 'METAFLOW_KUBERNETES_NAMESPACE', gpu: typing.Optional[int] = None, gpu_vendor: str = 'KUBERNETES_GPU_VENDOR', tolerations: typing.List[typing.Dict[str, str]] = [], labels: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_LABELS', annotations: typing.Dict[str, str] = 'METAFLOW_KUBERNETES_ANNOTATIONS', use_tmpfs: bool = False, tmpfs_tempdir: bool = True, tmpfs_size: typing.Optional[int] = None, tmpfs_path: typing.Optional[str] = '/metaflow_temp', persistent_volume_claims: typing.Optional[typing.Dict[str, str]] = None, shared_memory: typing.Optional[int] = None, port: typing.Optional[int] = None, compute_pool: typing.Optional[str] = None, hostname_resolution_timeout: int = 600, qos: str = 'Burstable', security_context: typing.Optional[typing.Dict[str, typing.Any]] = None) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
814
+ @typing.overload
815
+ 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]]]:
726
816
  """
727
- Specifies that this step should execute on Kubernetes.
817
+ Specifies the resources needed when executing this step.
818
+
819
+ Use `@resources` to specify the resource requirements
820
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
821
+
822
+ You can choose the compute layer on the command line by executing e.g.
823
+ ```
824
+ python myflow.py run --with batch
825
+ ```
826
+ or
827
+ ```
828
+ python myflow.py run --with kubernetes
829
+ ```
830
+ which executes the flow on the desired system using the
831
+ requirements specified in `@resources`.
728
832
 
729
833
 
730
834
  Parameters
731
835
  ----------
732
836
  cpu : int, default 1
733
- Number of CPUs required for this step. If `@resources` is
734
- also present, the maximum value from all decorators is used.
735
- memory : int, default 4096
736
- Memory size (in MB) required for this step. If
737
- `@resources` is also present, the maximum value from all decorators is
738
- used.
739
- disk : int, default 10240
740
- Disk size (in MB) required for this step. If
741
- `@resources` is also present, the maximum value from all decorators is
742
- used.
743
- image : str, optional, default None
744
- Docker image to use when launching on Kubernetes. If not specified, and
745
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
746
- not, a default Docker image mapping to the current version of Python is used.
747
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
748
- If given, the imagePullPolicy to be applied to the Docker image of the step.
749
- image_pull_secrets: List[str], default []
750
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
751
- Kubernetes image pull secrets to use when pulling container images
752
- in Kubernetes.
753
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
754
- Kubernetes service account to use when launching pod in Kubernetes.
755
- secrets : List[str], optional, default None
756
- Kubernetes secrets to use when launching pod in Kubernetes. These
757
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
758
- in Metaflow configuration.
759
- node_selector: Union[Dict[str,str], str], optional, default None
760
- Kubernetes node selector(s) to apply to the pod running the task.
761
- Can be passed in as a comma separated string of values e.g.
762
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
763
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
764
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
765
- Kubernetes namespace to use when launching pod in Kubernetes.
837
+ Number of CPUs required for this step.
766
838
  gpu : int, optional, default None
767
- Number of GPUs required for this step. A value of zero implies that
768
- the scheduled node should not have GPUs.
769
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
770
- The vendor of the GPUs to be used for this step.
771
- tolerations : List[Dict[str,str]], default []
772
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
773
- Kubernetes tolerations to use when launching pod in Kubernetes.
774
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
775
- Kubernetes labels to use when launching pod in Kubernetes.
776
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
777
- Kubernetes annotations to use when launching pod in Kubernetes.
778
- use_tmpfs : bool, default False
779
- This enables an explicit tmpfs mount for this step.
780
- tmpfs_tempdir : bool, default True
781
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
782
- tmpfs_size : int, optional, default: None
783
- The value for the size (in MiB) of the tmpfs mount for this step.
784
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
785
- memory allocated for this step.
786
- tmpfs_path : str, optional, default /metaflow_temp
787
- Path to tmpfs mount for this step.
788
- persistent_volume_claims : Dict[str, str], optional, default None
789
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
790
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
791
- shared_memory: int, optional
792
- Shared memory size (in MiB) required for this step
793
- port: int, optional
794
- Port number to specify in the Kubernetes job object
795
- compute_pool : str, optional, default None
796
- Compute pool to be used for for this step.
797
- If not specified, any accessible compute pool within the perimeter is used.
798
- hostname_resolution_timeout: int, default 10 * 60
799
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
800
- Only applicable when @parallel is used.
801
- qos: str, default: Burstable
802
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
803
-
804
- security_context: Dict[str, Any], optional, default None
805
- Container security context. Applies to the task container. Allows the following keys:
806
- - privileged: bool, optional, default None
807
- - allow_privilege_escalation: bool, optional, default None
808
- - run_as_user: int, optional, default None
809
- - run_as_group: int, optional, default None
810
- - run_as_non_root: bool, optional, default None
839
+ Number of GPUs required for this step.
840
+ disk : int, optional, default None
841
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
842
+ memory : int, default 4096
843
+ Memory size (in MB) required for this step.
844
+ shared_memory : int, optional, default None
845
+ The value for the size (in MiB) of the /dev/shm volume for this step.
846
+ This parameter maps to the `--shm-size` option in Docker.
811
847
  """
812
848
  ...
813
849
 
814
850
  @typing.overload
815
- 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]]]:
851
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
852
+ ...
853
+
854
+ @typing.overload
855
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
856
+ ...
857
+
858
+ 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):
816
859
  """
817
860
  Specifies the resources needed when executing this step.
818
861
 
@@ -848,45 +891,215 @@ def resources(*, cpu: int = 1, gpu: typing.Optional[int] = None, disk: typing.Op
848
891
  ...
849
892
 
850
893
  @typing.overload
851
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
894
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
895
+ """
896
+ Specifies the PyPI packages for all steps of the flow.
897
+
898
+ Use `@pypi_base` to set common packages required by all
899
+ steps and use `@pypi` to specify step-specific overrides.
900
+
901
+ Parameters
902
+ ----------
903
+ packages : Dict[str, str], default: {}
904
+ Packages to use for this flow. The key is the name of the package
905
+ and the value is the version to use.
906
+ python : str, optional, default: None
907
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
908
+ that the version used will correspond to the version of the Python interpreter used to start the run.
909
+ """
852
910
  ...
853
911
 
854
912
  @typing.overload
855
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
913
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
914
+ ...
915
+
916
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
917
+ """
918
+ Specifies the PyPI packages for all steps of the flow.
919
+
920
+ Use `@pypi_base` to set common packages required by all
921
+ steps and use `@pypi` to specify step-specific overrides.
922
+
923
+ Parameters
924
+ ----------
925
+ packages : Dict[str, str], default: {}
926
+ Packages to use for this flow. The key is the name of the package
927
+ and the value is the version to use.
928
+ python : str, optional, default: None
929
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
930
+ that the version used will correspond to the version of the Python interpreter used to start the run.
931
+ """
932
+ ...
933
+
934
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
935
+ """
936
+ Specifies what flows belong to the same project.
937
+
938
+ A project-specific namespace is created for all flows that
939
+ use the same `@project(name)`.
940
+
941
+
942
+ Parameters
943
+ ----------
944
+ name : str
945
+ Project name. Make sure that the name is unique amongst all
946
+ projects that use the same production scheduler. The name may
947
+ contain only lowercase alphanumeric characters and underscores.
948
+
949
+ branch : Optional[str], default None
950
+ The branch to use. If not specified, the branch is set to
951
+ `user.<username>` unless `production` is set to `True`. This can
952
+ also be set on the command line using `--branch` as a top-level option.
953
+ It is an error to specify `branch` in the decorator and on the command line.
954
+
955
+ production : bool, default False
956
+ Whether or not the branch is the production branch. This can also be set on the
957
+ command line using `--production` as a top-level option. It is an error to specify
958
+ `production` in the decorator and on the command line.
959
+ The project branch name will be:
960
+ - if `branch` is specified:
961
+ - if `production` is True: `prod.<branch>`
962
+ - if `production` is False: `test.<branch>`
963
+ - if `branch` is not specified:
964
+ - if `production` is True: `prod`
965
+ - if `production` is False: `user.<username>`
966
+ """
967
+ ...
968
+
969
+ 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]]:
970
+ """
971
+ 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.
972
+ 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.
973
+
974
+
975
+ Parameters
976
+ ----------
977
+ timeout : int
978
+ Time, in seconds before the task times out and fails. (Default: 3600)
979
+ poke_interval : int
980
+ Time in seconds that the job should wait in between each try. (Default: 60)
981
+ mode : str
982
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
983
+ exponential_backoff : bool
984
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
985
+ pool : str
986
+ the slot pool this task should run in,
987
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
988
+ soft_fail : bool
989
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
990
+ name : str
991
+ Name of the sensor on Airflow
992
+ description : str
993
+ Description of sensor in the Airflow UI
994
+ external_dag_id : str
995
+ The dag_id that contains the task you want to wait for.
996
+ external_task_ids : List[str]
997
+ The list of task_ids that you want to wait for.
998
+ If None (default value) the sensor waits for the DAG. (Default: None)
999
+ allowed_states : List[str]
1000
+ Iterable of allowed states, (Default: ['success'])
1001
+ failed_states : List[str]
1002
+ Iterable of failed or dis-allowed states. (Default: None)
1003
+ execution_delta : datetime.timedelta
1004
+ time difference with the previous execution to look at,
1005
+ the default is the same logical date as the current task or DAG. (Default: None)
1006
+ check_existence: bool
1007
+ Set to True to check if the external task exists or check if
1008
+ the DAG to wait for exists. (Default: True)
1009
+ """
1010
+ ...
1011
+
1012
+ 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]]:
1013
+ """
1014
+ 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)
1015
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1016
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1017
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1018
+ starts only after all sensors finish.
1019
+
1020
+
1021
+ Parameters
1022
+ ----------
1023
+ timeout : int
1024
+ Time, in seconds before the task times out and fails. (Default: 3600)
1025
+ poke_interval : int
1026
+ Time in seconds that the job should wait in between each try. (Default: 60)
1027
+ mode : str
1028
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1029
+ exponential_backoff : bool
1030
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1031
+ pool : str
1032
+ the slot pool this task should run in,
1033
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1034
+ soft_fail : bool
1035
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1036
+ name : str
1037
+ Name of the sensor on Airflow
1038
+ description : str
1039
+ Description of sensor in the Airflow UI
1040
+ bucket_key : Union[str, List[str]]
1041
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1042
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1043
+ bucket_name : str
1044
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1045
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1046
+ wildcard_match : bool
1047
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1048
+ aws_conn_id : str
1049
+ a reference to the s3 connection on Airflow. (Default: None)
1050
+ verify : bool
1051
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1052
+ """
1053
+ ...
1054
+
1055
+ @typing.overload
1056
+ 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]]:
1057
+ """
1058
+ Specifies the times when the flow should be run when running on a
1059
+ production scheduler.
1060
+
1061
+
1062
+ Parameters
1063
+ ----------
1064
+ hourly : bool, default False
1065
+ Run the workflow hourly.
1066
+ daily : bool, default True
1067
+ Run the workflow daily.
1068
+ weekly : bool, default False
1069
+ Run the workflow weekly.
1070
+ cron : str, optional, default None
1071
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1072
+ specified by this expression.
1073
+ timezone : str, optional, default None
1074
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1075
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1076
+ """
1077
+ ...
1078
+
1079
+ @typing.overload
1080
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
856
1081
  ...
857
1082
 
858
- 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):
1083
+ 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):
859
1084
  """
860
- Specifies the resources needed when executing this step.
861
-
862
- Use `@resources` to specify the resource requirements
863
- independently of the specific compute layer (`@batch`, `@kubernetes`).
864
-
865
- You can choose the compute layer on the command line by executing e.g.
866
- ```
867
- python myflow.py run --with batch
868
- ```
869
- or
870
- ```
871
- python myflow.py run --with kubernetes
872
- ```
873
- which executes the flow on the desired system using the
874
- requirements specified in `@resources`.
1085
+ Specifies the times when the flow should be run when running on a
1086
+ production scheduler.
875
1087
 
876
1088
 
877
1089
  Parameters
878
1090
  ----------
879
- cpu : int, default 1
880
- Number of CPUs required for this step.
881
- gpu : int, optional, default None
882
- Number of GPUs required for this step.
883
- disk : int, optional, default None
884
- Disk size (in MB) required for this step. Only applies on Kubernetes.
885
- memory : int, default 4096
886
- Memory size (in MB) required for this step.
887
- shared_memory : int, optional, default None
888
- The value for the size (in MiB) of the /dev/shm volume for this step.
889
- This parameter maps to the `--shm-size` option in Docker.
1091
+ hourly : bool, default False
1092
+ Run the workflow hourly.
1093
+ daily : bool, default True
1094
+ Run the workflow daily.
1095
+ weekly : bool, default False
1096
+ Run the workflow weekly.
1097
+ cron : str, optional, default None
1098
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1099
+ specified by this expression.
1100
+ timezone : str, optional, default None
1101
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1102
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
890
1103
  """
891
1104
  ...
892
1105
 
@@ -1042,84 +1255,6 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1042
1255
  """
1043
1256
  ...
1044
1257
 
1045
- 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]]:
1046
- """
1047
- 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)
1048
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1049
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1050
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1051
- starts only after all sensors finish.
1052
-
1053
-
1054
- Parameters
1055
- ----------
1056
- timeout : int
1057
- Time, in seconds before the task times out and fails. (Default: 3600)
1058
- poke_interval : int
1059
- Time in seconds that the job should wait in between each try. (Default: 60)
1060
- mode : str
1061
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1062
- exponential_backoff : bool
1063
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1064
- pool : str
1065
- the slot pool this task should run in,
1066
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1067
- soft_fail : bool
1068
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1069
- name : str
1070
- Name of the sensor on Airflow
1071
- description : str
1072
- Description of sensor in the Airflow UI
1073
- bucket_key : Union[str, List[str]]
1074
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1075
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1076
- bucket_name : str
1077
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1078
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1079
- wildcard_match : bool
1080
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1081
- aws_conn_id : str
1082
- a reference to the s3 connection on Airflow. (Default: None)
1083
- verify : bool
1084
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1085
- """
1086
- ...
1087
-
1088
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1089
- """
1090
- Specifies what flows belong to the same project.
1091
-
1092
- A project-specific namespace is created for all flows that
1093
- use the same `@project(name)`.
1094
-
1095
-
1096
- Parameters
1097
- ----------
1098
- name : str
1099
- Project name. Make sure that the name is unique amongst all
1100
- projects that use the same production scheduler. The name may
1101
- contain only lowercase alphanumeric characters and underscores.
1102
-
1103
- branch : Optional[str], default None
1104
- The branch to use. If not specified, the branch is set to
1105
- `user.<username>` unless `production` is set to `True`. This can
1106
- also be set on the command line using `--branch` as a top-level option.
1107
- It is an error to specify `branch` in the decorator and on the command line.
1108
-
1109
- production : bool, default False
1110
- Whether or not the branch is the production branch. This can also be set on the
1111
- command line using `--production` as a top-level option. It is an error to specify
1112
- `production` in the decorator and on the command line.
1113
- The project branch name will be:
1114
- - if `branch` is specified:
1115
- - if `production` is True: `prod.<branch>`
1116
- - if `production` is False: `test.<branch>`
1117
- - if `branch` is not specified:
1118
- - if `production` is True: `prod`
1119
- - if `production` is False: `user.<username>`
1120
- """
1121
- ...
1122
-
1123
1258
  @typing.overload
1124
1259
  def trigger(*, event: typing.Union[str, typing.Dict[str, typing.Any], None] = None, events: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1125
1260
  """
@@ -1213,138 +1348,3 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
1213
1348
  """
1214
1349
  ...
1215
1350
 
1216
- @typing.overload
1217
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1218
- """
1219
- Specifies the PyPI packages for all steps of the flow.
1220
-
1221
- Use `@pypi_base` to set common packages required by all
1222
- steps and use `@pypi` to specify step-specific overrides.
1223
-
1224
- Parameters
1225
- ----------
1226
- packages : Dict[str, str], default: {}
1227
- Packages to use for this flow. The key is the name of the package
1228
- and the value is the version to use.
1229
- python : str, optional, default: None
1230
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1231
- that the version used will correspond to the version of the Python interpreter used to start the run.
1232
- """
1233
- ...
1234
-
1235
- @typing.overload
1236
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1237
- ...
1238
-
1239
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
1240
- """
1241
- Specifies the PyPI packages for all steps of the flow.
1242
-
1243
- Use `@pypi_base` to set common packages required by all
1244
- steps and use `@pypi` to specify step-specific overrides.
1245
-
1246
- Parameters
1247
- ----------
1248
- packages : Dict[str, str], default: {}
1249
- Packages to use for this flow. The key is the name of the package
1250
- and the value is the version to use.
1251
- python : str, optional, default: None
1252
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
1253
- that the version used will correspond to the version of the Python interpreter used to start the run.
1254
- """
1255
- ...
1256
-
1257
- 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]]:
1258
- """
1259
- 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.
1260
- 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.
1261
-
1262
-
1263
- Parameters
1264
- ----------
1265
- timeout : int
1266
- Time, in seconds before the task times out and fails. (Default: 3600)
1267
- poke_interval : int
1268
- Time in seconds that the job should wait in between each try. (Default: 60)
1269
- mode : str
1270
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1271
- exponential_backoff : bool
1272
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1273
- pool : str
1274
- the slot pool this task should run in,
1275
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1276
- soft_fail : bool
1277
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1278
- name : str
1279
- Name of the sensor on Airflow
1280
- description : str
1281
- Description of sensor in the Airflow UI
1282
- external_dag_id : str
1283
- The dag_id that contains the task you want to wait for.
1284
- external_task_ids : List[str]
1285
- The list of task_ids that you want to wait for.
1286
- If None (default value) the sensor waits for the DAG. (Default: None)
1287
- allowed_states : List[str]
1288
- Iterable of allowed states, (Default: ['success'])
1289
- failed_states : List[str]
1290
- Iterable of failed or dis-allowed states. (Default: None)
1291
- execution_delta : datetime.timedelta
1292
- time difference with the previous execution to look at,
1293
- the default is the same logical date as the current task or DAG. (Default: None)
1294
- check_existence: bool
1295
- Set to True to check if the external task exists or check if
1296
- the DAG to wait for exists. (Default: True)
1297
- """
1298
- ...
1299
-
1300
- @typing.overload
1301
- 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]]:
1302
- """
1303
- Specifies the times when the flow should be run when running on a
1304
- production scheduler.
1305
-
1306
-
1307
- Parameters
1308
- ----------
1309
- hourly : bool, default False
1310
- Run the workflow hourly.
1311
- daily : bool, default True
1312
- Run the workflow daily.
1313
- weekly : bool, default False
1314
- Run the workflow weekly.
1315
- cron : str, optional, default None
1316
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1317
- specified by this expression.
1318
- timezone : str, optional, default None
1319
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1320
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1321
- """
1322
- ...
1323
-
1324
- @typing.overload
1325
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1326
- ...
1327
-
1328
- 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):
1329
- """
1330
- Specifies the times when the flow should be run when running on a
1331
- production scheduler.
1332
-
1333
-
1334
- Parameters
1335
- ----------
1336
- hourly : bool, default False
1337
- Run the workflow hourly.
1338
- daily : bool, default True
1339
- Run the workflow daily.
1340
- weekly : bool, default False
1341
- Run the workflow weekly.
1342
- cron : str, optional, default None
1343
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1344
- specified by this expression.
1345
- timezone : str, optional, default None
1346
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1347
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1348
- """
1349
- ...
1350
-