metaflow-stubs 2.18.2__py2.py3-none-any.whl → 2.18.4__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 +512 -512
  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 +6 -6
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +3 -3
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  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 +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +6 -2
  20. metaflow-stubs/metaflow_current.pyi +23 -23
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/multicore_utils.pyi +2 -2
  23. metaflow-stubs/packaging_sys/__init__.pyi +4 -4
  24. metaflow-stubs/packaging_sys/backend.pyi +4 -4
  25. metaflow-stubs/packaging_sys/distribution_support.pyi +4 -4
  26. metaflow-stubs/packaging_sys/tar_backend.pyi +5 -5
  27. metaflow-stubs/packaging_sys/utils.pyi +2 -2
  28. metaflow-stubs/packaging_sys/v1.pyi +2 -2
  29. metaflow-stubs/parameters.pyi +3 -3
  30. metaflow-stubs/plugins/__init__.pyi +11 -11
  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 +6 -4
  41. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  42. metaflow-stubs/plugins/argo/argo_workflows.pyi +9 -2
  43. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  44. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +3 -3
  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 +4 -4
  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 +4 -4
  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 +3 -3
  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 +6 -3
  64. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +2 -2
  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 +2 -2
  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 +2 -2
  79. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  80. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  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 +3 -3
  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 +2 -2
  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 +2 -2
  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 +6 -6
  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 +2 -2
  135. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  136. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  137. metaflow-stubs/plugins/uv/uv_environment.pyi +3 -3
  138. metaflow-stubs/pylint_wrapper.pyi +2 -2
  139. metaflow-stubs/runner/__init__.pyi +2 -2
  140. metaflow-stubs/runner/deployer.pyi +33 -33
  141. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  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 +2 -2
  154. metaflow-stubs/user_configs/config_parameters.pyi +6 -6
  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 +5 -5
  161. metaflow-stubs/version.pyi +2 -2
  162. {metaflow_stubs-2.18.2.dist-info → metaflow_stubs-2.18.4.dist-info}/METADATA +2 -2
  163. metaflow_stubs-2.18.4.dist-info/RECORD +166 -0
  164. metaflow_stubs-2.18.2.dist-info/RECORD +0 -166
  165. {metaflow_stubs-2.18.2.dist-info → metaflow_stubs-2.18.4.dist-info}/WHEEL +0 -0
  166. {metaflow_stubs-2.18.2.dist-info → metaflow_stubs-2.18.4.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.2 #
4
- # Generated on 2025-09-03T10:37:45.851135 #
3
+ # MF version: 2.18.4 #
4
+ # Generated on 2025-09-12T00:00:14.759393 #
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 events as events
43
- from . import tuple_util as tuple_util
44
42
  from . import metaflow_git as metaflow_git
43
+ from . import tuple_util as tuple_util
44
+ from . import events as events
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 conda_environment_yml_parser as conda_environment_yml_parser
52
51
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_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,81 +153,53 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
153
153
  ...
154
154
 
155
155
  @typing.overload
156
- 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]]]:
156
+ 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]]]:
157
157
  """
158
- Specifies the resources needed when executing this step.
159
-
160
- Use `@resources` to specify the resource requirements
161
- independently of the specific compute layer (`@batch`, `@kubernetes`).
158
+ Specifies the PyPI packages for the step.
162
159
 
163
- You can choose the compute layer on the command line by executing e.g.
164
- ```
165
- python myflow.py run --with batch
166
- ```
167
- or
168
- ```
169
- python myflow.py run --with kubernetes
170
- ```
171
- which executes the flow on the desired system using the
172
- requirements specified in `@resources`.
160
+ Information in this decorator will augment any
161
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
162
+ you can use `@pypi_base` to set packages required by all
163
+ steps and use `@pypi` to specify step-specific overrides.
173
164
 
174
165
 
175
166
  Parameters
176
167
  ----------
177
- cpu : int, default 1
178
- Number of CPUs required for this step.
179
- gpu : int, optional, default None
180
- Number of GPUs required for this step.
181
- disk : int, optional, default None
182
- Disk size (in MB) required for this step. Only applies on Kubernetes.
183
- memory : int, default 4096
184
- Memory size (in MB) required for this step.
185
- shared_memory : int, optional, default None
186
- The value for the size (in MiB) of the /dev/shm volume for this step.
187
- This parameter maps to the `--shm-size` option in Docker.
168
+ packages : Dict[str, str], default: {}
169
+ Packages to use for this step. The key is the name of the package
170
+ and the value is the version to use.
171
+ python : str, optional, default: None
172
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
173
+ that the version used will correspond to the version of the Python interpreter used to start the run.
188
174
  """
189
175
  ...
190
176
 
191
177
  @typing.overload
192
- def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
178
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
193
179
  ...
194
180
 
195
181
  @typing.overload
196
- def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
182
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
197
183
  ...
198
184
 
199
- 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):
185
+ 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
186
  """
201
- Specifies the resources needed when executing this step.
202
-
203
- Use `@resources` to specify the resource requirements
204
- independently of the specific compute layer (`@batch`, `@kubernetes`).
187
+ Specifies the PyPI packages for the step.
205
188
 
206
- You can choose the compute layer on the command line by executing e.g.
207
- ```
208
- python myflow.py run --with batch
209
- ```
210
- or
211
- ```
212
- python myflow.py run --with kubernetes
213
- ```
214
- which executes the flow on the desired system using the
215
- requirements specified in `@resources`.
189
+ Information in this decorator will augment any
190
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
191
+ you can use `@pypi_base` to set packages required by all
192
+ steps and use `@pypi` to specify step-specific overrides.
216
193
 
217
194
 
218
195
  Parameters
219
196
  ----------
220
- cpu : int, default 1
221
- Number of CPUs required for this step.
222
- gpu : int, optional, default None
223
- Number of GPUs required for this step.
224
- disk : int, optional, default None
225
- Disk size (in MB) required for this step. Only applies on Kubernetes.
226
- memory : int, default 4096
227
- Memory size (in MB) required for this step.
228
- shared_memory : int, optional, default None
229
- The value for the size (in MiB) of the /dev/shm volume for this step.
230
- This parameter maps to the `--shm-size` option in Docker.
197
+ packages : Dict[str, str], default: {}
198
+ Packages to use for this step. The key is the name of the package
199
+ and the value is the version to use.
200
+ python : str, optional, default: None
201
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
202
+ that the version used will correspond to the version of the Python interpreter used to start the run.
231
203
  """
232
204
  ...
233
205
 
@@ -287,9 +259,27 @@ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
287
259
  ...
288
260
 
289
261
  @typing.overload
290
- 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]]]:
262
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
291
263
  """
292
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
264
+ Decorator prototype for all step decorators. This function gets specialized
265
+ and imported for all decorators types by _import_plugin_decorators().
266
+ """
267
+ ...
268
+
269
+ @typing.overload
270
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
271
+ ...
272
+
273
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
274
+ """
275
+ Decorator prototype for all step decorators. This function gets specialized
276
+ and imported for all decorators types by _import_plugin_decorators().
277
+ """
278
+ ...
279
+
280
+ 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]]]:
281
+ """
282
+ Specifies that this step should execute on Kubernetes.
293
283
 
294
284
 
295
285
  Parameters
@@ -297,156 +287,119 @@ def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optio
297
287
  cpu : int, default 1
298
288
  Number of CPUs required for this step. If `@resources` is
299
289
  also present, the maximum value from all decorators is used.
300
- gpu : int, default 0
301
- Number of GPUs required for this step. If `@resources` is
302
- also present, the maximum value from all decorators is used.
303
290
  memory : int, default 4096
304
291
  Memory size (in MB) required for this step. If
305
292
  `@resources` is also present, the maximum value from all decorators is
306
293
  used.
294
+ disk : int, default 10240
295
+ Disk size (in MB) required for this step. If
296
+ `@resources` is also present, the maximum value from all decorators is
297
+ used.
307
298
  image : str, optional, default None
308
- Docker image to use when launching on AWS Batch. If not specified, and
309
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
299
+ Docker image to use when launching on Kubernetes. If not specified, and
300
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
310
301
  not, a default Docker image mapping to the current version of Python is used.
311
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
312
- AWS Batch Job Queue to submit the job to.
313
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
314
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
315
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
316
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
317
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
318
- shared_memory : int, optional, default None
319
- The value for the size (in MiB) of the /dev/shm volume for this step.
320
- This parameter maps to the `--shm-size` option in Docker.
321
- max_swap : int, optional, default None
322
- The total amount of swap memory (in MiB) a container can use for this
323
- step. This parameter is translated to the `--memory-swap` option in
324
- Docker where the value is the sum of the container memory plus the
325
- `max_swap` value.
326
- swappiness : int, optional, default None
327
- This allows you to tune memory swappiness behavior for this step.
328
- A swappiness value of 0 causes swapping not to happen unless absolutely
329
- necessary. A swappiness value of 100 causes pages to be swapped very
330
- aggressively. Accepted values are whole numbers between 0 and 100.
331
- aws_batch_tags: Dict[str, str], optional, default None
332
- Sets arbitrary AWS tags on the AWS Batch compute environment.
333
- Set as string key-value pairs.
302
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
303
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
304
+ image_pull_secrets: List[str], default []
305
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
306
+ Kubernetes image pull secrets to use when pulling container images
307
+ in Kubernetes.
308
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
309
+ Kubernetes service account to use when launching pod in Kubernetes.
310
+ secrets : List[str], optional, default None
311
+ Kubernetes secrets to use when launching pod in Kubernetes. These
312
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
313
+ in Metaflow configuration.
314
+ node_selector: Union[Dict[str,str], str], optional, default None
315
+ Kubernetes node selector(s) to apply to the pod running the task.
316
+ Can be passed in as a comma separated string of values e.g.
317
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
318
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
319
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
320
+ Kubernetes namespace to use when launching pod in Kubernetes.
321
+ gpu : int, optional, default None
322
+ Number of GPUs required for this step. A value of zero implies that
323
+ the scheduled node should not have GPUs.
324
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
325
+ The vendor of the GPUs to be used for this step.
326
+ tolerations : List[Dict[str,str]], default []
327
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
328
+ Kubernetes tolerations to use when launching pod in Kubernetes.
329
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
330
+ Kubernetes labels to use when launching pod in Kubernetes.
331
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
332
+ Kubernetes annotations to use when launching pod in Kubernetes.
334
333
  use_tmpfs : bool, default False
335
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
336
- not available on Fargate compute environments
334
+ This enables an explicit tmpfs mount for this step.
337
335
  tmpfs_tempdir : bool, default True
338
336
  sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
339
- tmpfs_size : int, optional, default None
337
+ tmpfs_size : int, optional, default: None
340
338
  The value for the size (in MiB) of the tmpfs mount for this step.
341
339
  This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
342
340
  memory allocated for this step.
343
- tmpfs_path : str, optional, default None
344
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
345
- inferentia : int, default 0
346
- Number of Inferentia chips required for this step.
347
- trainium : int, default None
348
- Alias for inferentia. Use only one of the two.
349
- efa : int, default 0
350
- Number of elastic fabric adapter network devices to attach to container
351
- ephemeral_storage : int, default None
352
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
353
- This is only relevant for Fargate compute environments
354
- log_driver: str, optional, default None
355
- The log driver to use for the Amazon ECS container.
356
- log_options: List[str], optional, default None
357
- List of strings containing options for the chosen log driver. The configurable values
358
- depend on the `log driver` chosen. Validation of these options is not supported yet.
359
- Example: [`awslogs-group:aws/batch/job`]
341
+ tmpfs_path : str, optional, default /metaflow_temp
342
+ Path to tmpfs mount for this step.
343
+ persistent_volume_claims : Dict[str, str], optional, default None
344
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
345
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
346
+ shared_memory: int, optional
347
+ Shared memory size (in MiB) required for this step
348
+ port: int, optional
349
+ Port number to specify in the Kubernetes job object
350
+ compute_pool : str, optional, default None
351
+ Compute pool to be used for for this step.
352
+ If not specified, any accessible compute pool within the perimeter is used.
353
+ hostname_resolution_timeout: int, default 10 * 60
354
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
355
+ Only applicable when @parallel is used.
356
+ qos: str, default: Burstable
357
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
358
+
359
+ security_context: Dict[str, Any], optional, default None
360
+ Container security context. Applies to the task container. Allows the following keys:
361
+ - privileged: bool, optional, default None
362
+ - allow_privilege_escalation: bool, optional, default None
363
+ - run_as_user: int, optional, default None
364
+ - run_as_group: int, optional, default None
365
+ - run_as_non_root: bool, optional, default None
360
366
  """
361
367
  ...
362
368
 
363
369
  @typing.overload
364
- def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
370
+ 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]]]:
371
+ """
372
+ Creates a human-readable report, a Metaflow Card, after this step completes.
373
+
374
+ Note that you may add multiple `@card` decorators in a step with different parameters.
375
+
376
+
377
+ Parameters
378
+ ----------
379
+ type : str, default 'default'
380
+ Card type.
381
+ id : str, optional, default None
382
+ If multiple cards are present, use this id to identify this card.
383
+ options : Dict[str, Any], default {}
384
+ Options passed to the card. The contents depend on the card type.
385
+ timeout : int, default 45
386
+ Interrupt reporting if it takes more than this many seconds.
387
+ """
365
388
  ...
366
389
 
367
390
  @typing.overload
368
- def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
391
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
369
392
  ...
370
393
 
371
- 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):
394
+ @typing.overload
395
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
396
+ ...
397
+
398
+ 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):
372
399
  """
373
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
374
-
375
-
376
- Parameters
377
- ----------
378
- cpu : int, default 1
379
- Number of CPUs required for this step. If `@resources` is
380
- also present, the maximum value from all decorators is used.
381
- gpu : int, default 0
382
- Number of GPUs required for this step. If `@resources` is
383
- also present, the maximum value from all decorators is used.
384
- memory : int, default 4096
385
- Memory size (in MB) required for this step. If
386
- `@resources` is also present, the maximum value from all decorators is
387
- used.
388
- image : str, optional, default None
389
- Docker image to use when launching on AWS Batch. If not specified, and
390
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
391
- not, a default Docker image mapping to the current version of Python is used.
392
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
393
- AWS Batch Job Queue to submit the job to.
394
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
395
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
396
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
397
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
398
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
399
- shared_memory : int, optional, default None
400
- The value for the size (in MiB) of the /dev/shm volume for this step.
401
- This parameter maps to the `--shm-size` option in Docker.
402
- max_swap : int, optional, default None
403
- The total amount of swap memory (in MiB) a container can use for this
404
- step. This parameter is translated to the `--memory-swap` option in
405
- Docker where the value is the sum of the container memory plus the
406
- `max_swap` value.
407
- swappiness : int, optional, default None
408
- This allows you to tune memory swappiness behavior for this step.
409
- A swappiness value of 0 causes swapping not to happen unless absolutely
410
- necessary. A swappiness value of 100 causes pages to be swapped very
411
- aggressively. Accepted values are whole numbers between 0 and 100.
412
- aws_batch_tags: Dict[str, str], optional, default None
413
- Sets arbitrary AWS tags on the AWS Batch compute environment.
414
- Set as string key-value pairs.
415
- use_tmpfs : bool, default False
416
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
417
- not available on Fargate compute environments
418
- tmpfs_tempdir : bool, default True
419
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
420
- tmpfs_size : int, optional, default None
421
- The value for the size (in MiB) of the tmpfs mount for this step.
422
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
423
- memory allocated for this step.
424
- tmpfs_path : str, optional, default None
425
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
426
- inferentia : int, default 0
427
- Number of Inferentia chips required for this step.
428
- trainium : int, default None
429
- Alias for inferentia. Use only one of the two.
430
- efa : int, default 0
431
- Number of elastic fabric adapter network devices to attach to container
432
- ephemeral_storage : int, default None
433
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
434
- This is only relevant for Fargate compute environments
435
- log_driver: str, optional, default None
436
- The log driver to use for the Amazon ECS container.
437
- log_options: List[str], optional, default None
438
- List of strings containing options for the chosen log driver. The configurable values
439
- depend on the `log driver` chosen. Validation of these options is not supported yet.
440
- Example: [`awslogs-group:aws/batch/job`]
441
- """
442
- ...
443
-
444
- @typing.overload
445
- 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]]]:
446
- """
447
- Creates a human-readable report, a Metaflow Card, after this step completes.
448
-
449
- Note that you may add multiple `@card` decorators in a step with different parameters.
400
+ Creates a human-readable report, a Metaflow Card, after this step completes.
401
+
402
+ Note that you may add multiple `@card` decorators in a step with different parameters.
450
403
 
451
404
 
452
405
  Parameters
@@ -462,85 +415,6 @@ def card(*, type: str = 'default', id: typing.Optional[str] = None, options: typ
462
415
  """
463
416
  ...
464
417
 
465
- @typing.overload
466
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
467
- ...
468
-
469
- @typing.overload
470
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
471
- ...
472
-
473
- 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):
474
- """
475
- Creates a human-readable report, a Metaflow Card, after this step completes.
476
-
477
- Note that you may add multiple `@card` decorators in a step with different parameters.
478
-
479
-
480
- Parameters
481
- ----------
482
- type : str, default 'default'
483
- Card type.
484
- id : str, optional, default None
485
- If multiple cards are present, use this id to identify this card.
486
- options : Dict[str, Any], default {}
487
- Options passed to the card. The contents depend on the card type.
488
- timeout : int, default 45
489
- Interrupt reporting if it takes more than this many seconds.
490
- """
491
- ...
492
-
493
- @typing.overload
494
- 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]]]:
495
- """
496
- Specifies the PyPI packages for the step.
497
-
498
- Information in this decorator will augment any
499
- attributes set in the `@pyi_base` flow-level decorator. Hence,
500
- you can use `@pypi_base` to set packages required by all
501
- steps and use `@pypi` to specify step-specific overrides.
502
-
503
-
504
- Parameters
505
- ----------
506
- packages : Dict[str, str], default: {}
507
- Packages to use for this step. The key is the name of the package
508
- and the value is the version to use.
509
- python : str, optional, default: None
510
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
511
- that the version used will correspond to the version of the Python interpreter used to start the run.
512
- """
513
- ...
514
-
515
- @typing.overload
516
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
517
- ...
518
-
519
- @typing.overload
520
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
521
- ...
522
-
523
- 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):
524
- """
525
- Specifies the PyPI packages for the step.
526
-
527
- Information in this decorator will augment any
528
- attributes set in the `@pyi_base` flow-level decorator. Hence,
529
- you can use `@pypi_base` to set packages required by all
530
- steps and use `@pypi` to specify step-specific overrides.
531
-
532
-
533
- Parameters
534
- ----------
535
- packages : Dict[str, str], default: {}
536
- Packages to use for this step. The key is the name of the package
537
- and the value is the version to use.
538
- python : str, optional, default: None
539
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
540
- that the version used will correspond to the version of the Python interpreter used to start the run.
541
- """
542
- ...
543
-
544
418
  @typing.overload
545
419
  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]]]:
546
420
  """
@@ -601,253 +475,379 @@ def conda(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
601
475
  ...
602
476
 
603
477
  @typing.overload
604
- 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]]]:
478
+ 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]]]:
605
479
  """
606
- Specifies environment variables to be set prior to the execution of a step.
480
+ Specifies a timeout for your step.
481
+
482
+ This decorator is useful if this step may hang indefinitely.
483
+
484
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
485
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
486
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
487
+
488
+ Note that all the values specified in parameters are added together so if you specify
489
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
607
490
 
608
491
 
609
492
  Parameters
610
493
  ----------
611
- vars : Dict[str, str], default {}
612
- Dictionary of environment variables to set.
494
+ seconds : int, default 0
495
+ Number of seconds to wait prior to timing out.
496
+ minutes : int, default 0
497
+ Number of minutes to wait prior to timing out.
498
+ hours : int, default 0
499
+ Number of hours to wait prior to timing out.
613
500
  """
614
501
  ...
615
502
 
616
503
  @typing.overload
617
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
504
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
618
505
  ...
619
506
 
620
507
  @typing.overload
621
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
508
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
622
509
  ...
623
510
 
624
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
511
+ 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):
625
512
  """
626
- Specifies environment variables to be set prior to the execution of a step.
513
+ Specifies a timeout for your step.
514
+
515
+ This decorator is useful if this step may hang indefinitely.
516
+
517
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
518
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
519
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
520
+
521
+ Note that all the values specified in parameters are added together so if you specify
522
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
627
523
 
628
524
 
629
525
  Parameters
630
526
  ----------
631
- vars : Dict[str, str], default {}
632
- Dictionary of environment variables to set.
527
+ seconds : int, default 0
528
+ Number of seconds to wait prior to timing out.
529
+ minutes : int, default 0
530
+ Number of minutes to wait prior to timing out.
531
+ hours : int, default 0
532
+ Number of hours to wait prior to timing out.
633
533
  """
634
534
  ...
635
535
 
636
536
  @typing.overload
637
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
537
+ 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]]]:
638
538
  """
639
- Decorator prototype for all step decorators. This function gets specialized
640
- and imported for all decorators types by _import_plugin_decorators().
539
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
540
+
541
+
542
+ Parameters
543
+ ----------
544
+ cpu : int, default 1
545
+ Number of CPUs required for this step. If `@resources` is
546
+ also present, the maximum value from all decorators is used.
547
+ gpu : int, default 0
548
+ Number of GPUs required for this step. If `@resources` is
549
+ also present, the maximum value from all decorators is used.
550
+ memory : int, default 4096
551
+ Memory size (in MB) required for this step. If
552
+ `@resources` is also present, the maximum value from all decorators is
553
+ used.
554
+ image : str, optional, default None
555
+ Docker image to use when launching on AWS Batch. If not specified, and
556
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
557
+ not, a default Docker image mapping to the current version of Python is used.
558
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
559
+ AWS Batch Job Queue to submit the job to.
560
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
561
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
562
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
563
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
564
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
565
+ shared_memory : int, optional, default None
566
+ The value for the size (in MiB) of the /dev/shm volume for this step.
567
+ This parameter maps to the `--shm-size` option in Docker.
568
+ max_swap : int, optional, default None
569
+ The total amount of swap memory (in MiB) a container can use for this
570
+ step. This parameter is translated to the `--memory-swap` option in
571
+ Docker where the value is the sum of the container memory plus the
572
+ `max_swap` value.
573
+ swappiness : int, optional, default None
574
+ This allows you to tune memory swappiness behavior for this step.
575
+ A swappiness value of 0 causes swapping not to happen unless absolutely
576
+ necessary. A swappiness value of 100 causes pages to be swapped very
577
+ aggressively. Accepted values are whole numbers between 0 and 100.
578
+ aws_batch_tags: Dict[str, str], optional, default None
579
+ Sets arbitrary AWS tags on the AWS Batch compute environment.
580
+ Set as string key-value pairs.
581
+ use_tmpfs : bool, default False
582
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
583
+ not available on Fargate compute environments
584
+ tmpfs_tempdir : bool, default True
585
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
586
+ tmpfs_size : int, optional, default None
587
+ The value for the size (in MiB) of the tmpfs mount for this step.
588
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
589
+ memory allocated for this step.
590
+ tmpfs_path : str, optional, default None
591
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
592
+ inferentia : int, default 0
593
+ Number of Inferentia chips required for this step.
594
+ trainium : int, default None
595
+ Alias for inferentia. Use only one of the two.
596
+ efa : int, default 0
597
+ Number of elastic fabric adapter network devices to attach to container
598
+ ephemeral_storage : int, default None
599
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
600
+ This is only relevant for Fargate compute environments
601
+ log_driver: str, optional, default None
602
+ The log driver to use for the Amazon ECS container.
603
+ log_options: List[str], optional, default None
604
+ List of strings containing options for the chosen log driver. The configurable values
605
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
606
+ Example: [`awslogs-group:aws/batch/job`]
641
607
  """
642
608
  ...
643
609
 
644
610
  @typing.overload
645
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
611
+ def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
646
612
  ...
647
613
 
648
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
614
+ @typing.overload
615
+ def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
616
+ ...
617
+
618
+ 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):
649
619
  """
650
- Decorator prototype for all step decorators. This function gets specialized
651
- and imported for all decorators types by _import_plugin_decorators().
620
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
621
+
622
+
623
+ Parameters
624
+ ----------
625
+ cpu : int, default 1
626
+ Number of CPUs required for this step. If `@resources` is
627
+ also present, the maximum value from all decorators is used.
628
+ gpu : int, default 0
629
+ Number of GPUs required for this step. If `@resources` is
630
+ also present, the maximum value from all decorators is used.
631
+ memory : int, default 4096
632
+ Memory size (in MB) required for this step. If
633
+ `@resources` is also present, the maximum value from all decorators is
634
+ used.
635
+ image : str, optional, default None
636
+ Docker image to use when launching on AWS Batch. If not specified, and
637
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
638
+ not, a default Docker image mapping to the current version of Python is used.
639
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
640
+ AWS Batch Job Queue to submit the job to.
641
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
642
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
643
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
644
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
645
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
646
+ shared_memory : int, optional, default None
647
+ The value for the size (in MiB) of the /dev/shm volume for this step.
648
+ This parameter maps to the `--shm-size` option in Docker.
649
+ max_swap : int, optional, default None
650
+ The total amount of swap memory (in MiB) a container can use for this
651
+ step. This parameter is translated to the `--memory-swap` option in
652
+ Docker where the value is the sum of the container memory plus the
653
+ `max_swap` value.
654
+ swappiness : int, optional, default None
655
+ This allows you to tune memory swappiness behavior for this step.
656
+ A swappiness value of 0 causes swapping not to happen unless absolutely
657
+ necessary. A swappiness value of 100 causes pages to be swapped very
658
+ aggressively. Accepted values are whole numbers between 0 and 100.
659
+ aws_batch_tags: Dict[str, str], optional, default None
660
+ Sets arbitrary AWS tags on the AWS Batch compute environment.
661
+ Set as string key-value pairs.
662
+ use_tmpfs : bool, default False
663
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
664
+ not available on Fargate compute environments
665
+ tmpfs_tempdir : bool, default True
666
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
667
+ tmpfs_size : int, optional, default None
668
+ The value for the size (in MiB) of the tmpfs mount for this step.
669
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
670
+ memory allocated for this step.
671
+ tmpfs_path : str, optional, default None
672
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
673
+ inferentia : int, default 0
674
+ Number of Inferentia chips required for this step.
675
+ trainium : int, default None
676
+ Alias for inferentia. Use only one of the two.
677
+ efa : int, default 0
678
+ Number of elastic fabric adapter network devices to attach to container
679
+ ephemeral_storage : int, default None
680
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
681
+ This is only relevant for Fargate compute environments
682
+ log_driver: str, optional, default None
683
+ The log driver to use for the Amazon ECS container.
684
+ log_options: List[str], optional, default None
685
+ List of strings containing options for the chosen log driver. The configurable values
686
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
687
+ Example: [`awslogs-group:aws/batch/job`]
652
688
  """
653
689
  ...
654
690
 
655
691
  @typing.overload
656
- 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]]]:
692
+ 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]]]:
657
693
  """
658
- Specifies that the step will success under all circumstances.
694
+ Specifies the resources needed when executing this step.
659
695
 
660
- The decorator will create an optional artifact, specified by `var`, which
661
- contains the exception raised. You can use it to detect the presence
662
- of errors, indicating that all happy-path artifacts produced by the step
663
- are missing.
696
+ Use `@resources` to specify the resource requirements
697
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
698
+
699
+ You can choose the compute layer on the command line by executing e.g.
700
+ ```
701
+ python myflow.py run --with batch
702
+ ```
703
+ or
704
+ ```
705
+ python myflow.py run --with kubernetes
706
+ ```
707
+ which executes the flow on the desired system using the
708
+ requirements specified in `@resources`.
664
709
 
665
710
 
666
711
  Parameters
667
712
  ----------
668
- var : str, optional, default None
669
- Name of the artifact in which to store the caught exception.
670
- If not specified, the exception is not stored.
671
- print_exception : bool, default True
672
- Determines whether or not the exception is printed to
673
- stdout when caught.
713
+ cpu : int, default 1
714
+ Number of CPUs required for this step.
715
+ gpu : int, optional, default None
716
+ Number of GPUs required for this step.
717
+ disk : int, optional, default None
718
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
719
+ memory : int, default 4096
720
+ Memory size (in MB) required for this step.
721
+ shared_memory : int, optional, default None
722
+ The value for the size (in MiB) of the /dev/shm volume for this step.
723
+ This parameter maps to the `--shm-size` option in Docker.
674
724
  """
675
725
  ...
676
726
 
677
727
  @typing.overload
678
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
728
+ def resources(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
679
729
  ...
680
730
 
681
731
  @typing.overload
682
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
732
+ def resources(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
683
733
  ...
684
734
 
685
- 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):
735
+ 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):
686
736
  """
687
- Specifies that the step will success under all circumstances.
737
+ Specifies the resources needed when executing this step.
688
738
 
689
- The decorator will create an optional artifact, specified by `var`, which
690
- contains the exception raised. You can use it to detect the presence
691
- of errors, indicating that all happy-path artifacts produced by the step
692
- are missing.
739
+ Use `@resources` to specify the resource requirements
740
+ independently of the specific compute layer (`@batch`, `@kubernetes`).
741
+
742
+ You can choose the compute layer on the command line by executing e.g.
743
+ ```
744
+ python myflow.py run --with batch
745
+ ```
746
+ or
747
+ ```
748
+ python myflow.py run --with kubernetes
749
+ ```
750
+ which executes the flow on the desired system using the
751
+ requirements specified in `@resources`.
693
752
 
694
753
 
695
754
  Parameters
696
755
  ----------
697
- var : str, optional, default None
698
- Name of the artifact in which to store the caught exception.
699
- If not specified, the exception is not stored.
700
- print_exception : bool, default True
701
- Determines whether or not the exception is printed to
702
- stdout when caught.
756
+ cpu : int, default 1
757
+ Number of CPUs required for this step.
758
+ gpu : int, optional, default None
759
+ Number of GPUs required for this step.
760
+ disk : int, optional, default None
761
+ Disk size (in MB) required for this step. Only applies on Kubernetes.
762
+ memory : int, default 4096
763
+ Memory size (in MB) required for this step.
764
+ shared_memory : int, optional, default None
765
+ The value for the size (in MiB) of the /dev/shm volume for this step.
766
+ This parameter maps to the `--shm-size` option in Docker.
703
767
  """
704
768
  ...
705
769
 
706
- 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]]]:
770
+ @typing.overload
771
+ 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]]]:
707
772
  """
708
- Specifies that this step should execute on Kubernetes.
773
+ Specifies that the step will success under all circumstances.
709
774
 
775
+ The decorator will create an optional artifact, specified by `var`, which
776
+ contains the exception raised. You can use it to detect the presence
777
+ of errors, indicating that all happy-path artifacts produced by the step
778
+ are missing.
710
779
 
711
- Parameters
712
- ----------
713
- cpu : int, default 1
714
- Number of CPUs required for this step. If `@resources` is
715
- also present, the maximum value from all decorators is used.
716
- memory : int, default 4096
717
- Memory size (in MB) required for this step. If
718
- `@resources` is also present, the maximum value from all decorators is
719
- used.
720
- disk : int, default 10240
721
- Disk size (in MB) required for this step. If
722
- `@resources` is also present, the maximum value from all decorators is
723
- used.
724
- image : str, optional, default None
725
- Docker image to use when launching on Kubernetes. If not specified, and
726
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
727
- not, a default Docker image mapping to the current version of Python is used.
728
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
729
- If given, the imagePullPolicy to be applied to the Docker image of the step.
730
- image_pull_secrets: List[str], default []
731
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
732
- Kubernetes image pull secrets to use when pulling container images
733
- in Kubernetes.
734
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
735
- Kubernetes service account to use when launching pod in Kubernetes.
736
- secrets : List[str], optional, default None
737
- Kubernetes secrets to use when launching pod in Kubernetes. These
738
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
739
- in Metaflow configuration.
740
- node_selector: Union[Dict[str,str], str], optional, default None
741
- Kubernetes node selector(s) to apply to the pod running the task.
742
- Can be passed in as a comma separated string of values e.g.
743
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
744
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
745
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
746
- Kubernetes namespace to use when launching pod in Kubernetes.
747
- gpu : int, optional, default None
748
- Number of GPUs required for this step. A value of zero implies that
749
- the scheduled node should not have GPUs.
750
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
751
- The vendor of the GPUs to be used for this step.
752
- tolerations : List[Dict[str,str]], default []
753
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
754
- Kubernetes tolerations to use when launching pod in Kubernetes.
755
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
756
- Kubernetes labels to use when launching pod in Kubernetes.
757
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
758
- Kubernetes annotations to use when launching pod in Kubernetes.
759
- use_tmpfs : bool, default False
760
- This enables an explicit tmpfs mount for this step.
761
- tmpfs_tempdir : bool, default True
762
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
763
- tmpfs_size : int, optional, default: None
764
- The value for the size (in MiB) of the tmpfs mount for this step.
765
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
766
- memory allocated for this step.
767
- tmpfs_path : str, optional, default /metaflow_temp
768
- Path to tmpfs mount for this step.
769
- persistent_volume_claims : Dict[str, str], optional, default None
770
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
771
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
772
- shared_memory: int, optional
773
- Shared memory size (in MiB) required for this step
774
- port: int, optional
775
- Port number to specify in the Kubernetes job object
776
- compute_pool : str, optional, default None
777
- Compute pool to be used for for this step.
778
- If not specified, any accessible compute pool within the perimeter is used.
779
- hostname_resolution_timeout: int, default 10 * 60
780
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
781
- Only applicable when @parallel is used.
782
- qos: str, default: Burstable
783
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
784
780
 
785
- security_context: Dict[str, Any], optional, default None
786
- Container security context. Applies to the task container. Allows the following keys:
787
- - privileged: bool, optional, default None
788
- - allow_privilege_escalation: bool, optional, default None
789
- - run_as_user: int, optional, default None
790
- - run_as_group: int, optional, default None
791
- - run_as_non_root: bool, optional, default None
781
+ Parameters
782
+ ----------
783
+ var : str, optional, default None
784
+ Name of the artifact in which to store the caught exception.
785
+ If not specified, the exception is not stored.
786
+ print_exception : bool, default True
787
+ Determines whether or not the exception is printed to
788
+ stdout when caught.
792
789
  """
793
790
  ...
794
791
 
795
792
  @typing.overload
796
- 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]]]:
793
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
794
+ ...
795
+
796
+ @typing.overload
797
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
798
+ ...
799
+
800
+ 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):
797
801
  """
798
- Specifies a timeout for your step.
802
+ Specifies that the step will success under all circumstances.
799
803
 
800
- This decorator is useful if this step may hang indefinitely.
804
+ The decorator will create an optional artifact, specified by `var`, which
805
+ contains the exception raised. You can use it to detect the presence
806
+ of errors, indicating that all happy-path artifacts produced by the step
807
+ are missing.
801
808
 
802
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
803
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
804
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
805
809
 
806
- Note that all the values specified in parameters are added together so if you specify
807
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
810
+ Parameters
811
+ ----------
812
+ var : str, optional, default None
813
+ Name of the artifact in which to store the caught exception.
814
+ If not specified, the exception is not stored.
815
+ print_exception : bool, default True
816
+ Determines whether or not the exception is printed to
817
+ stdout when caught.
818
+ """
819
+ ...
820
+
821
+ @typing.overload
822
+ 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]]]:
823
+ """
824
+ Specifies environment variables to be set prior to the execution of a step.
808
825
 
809
826
 
810
827
  Parameters
811
828
  ----------
812
- seconds : int, default 0
813
- Number of seconds to wait prior to timing out.
814
- minutes : int, default 0
815
- Number of minutes to wait prior to timing out.
816
- hours : int, default 0
817
- Number of hours to wait prior to timing out.
829
+ vars : Dict[str, str], default {}
830
+ Dictionary of environment variables to set.
818
831
  """
819
832
  ...
820
833
 
821
834
  @typing.overload
822
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
835
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
823
836
  ...
824
837
 
825
838
  @typing.overload
826
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
839
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
827
840
  ...
828
841
 
829
- 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):
842
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
830
843
  """
831
- Specifies a timeout for your step.
832
-
833
- This decorator is useful if this step may hang indefinitely.
834
-
835
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
836
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
837
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
838
-
839
- Note that all the values specified in parameters are added together so if you specify
840
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
844
+ Specifies environment variables to be set prior to the execution of a step.
841
845
 
842
846
 
843
847
  Parameters
844
848
  ----------
845
- seconds : int, default 0
846
- Number of seconds to wait prior to timing out.
847
- minutes : int, default 0
848
- Number of minutes to wait prior to timing out.
849
- hours : int, default 0
850
- Number of hours to wait prior to timing out.
849
+ vars : Dict[str, str], default {}
850
+ Dictionary of environment variables to set.
851
851
  """
852
852
  ...
853
853
 
@@ -890,57 +890,6 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
890
890
  """
891
891
  ...
892
892
 
893
- @typing.overload
894
- 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]]:
895
- """
896
- Specifies the times when the flow should be run when running on a
897
- production scheduler.
898
-
899
-
900
- Parameters
901
- ----------
902
- hourly : bool, default False
903
- Run the workflow hourly.
904
- daily : bool, default True
905
- Run the workflow daily.
906
- weekly : bool, default False
907
- Run the workflow weekly.
908
- cron : str, optional, default None
909
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
910
- specified by this expression.
911
- timezone : str, optional, default None
912
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
913
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
914
- """
915
- ...
916
-
917
- @typing.overload
918
- def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
919
- ...
920
-
921
- 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):
922
- """
923
- Specifies the times when the flow should be run when running on a
924
- production scheduler.
925
-
926
-
927
- Parameters
928
- ----------
929
- hourly : bool, default False
930
- Run the workflow hourly.
931
- daily : bool, default True
932
- Run the workflow daily.
933
- weekly : bool, default False
934
- Run the workflow weekly.
935
- cron : str, optional, default None
936
- Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
937
- specified by this expression.
938
- timezone : str, optional, default None
939
- Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
940
- which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
941
- """
942
- ...
943
-
944
893
  @typing.overload
945
894
  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]]:
946
895
  """
@@ -1085,41 +1034,6 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
1085
1034
  """
1086
1035
  ...
1087
1036
 
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
1037
  @typing.overload
1124
1038
  def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1125
1039
  """
@@ -1204,6 +1118,57 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1204
1118
  """
1205
1119
  ...
1206
1120
 
1121
+ @typing.overload
1122
+ 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]]:
1123
+ """
1124
+ Specifies the times when the flow should be run when running on a
1125
+ production scheduler.
1126
+
1127
+
1128
+ Parameters
1129
+ ----------
1130
+ hourly : bool, default False
1131
+ Run the workflow hourly.
1132
+ daily : bool, default True
1133
+ Run the workflow daily.
1134
+ weekly : bool, default False
1135
+ Run the workflow weekly.
1136
+ cron : str, optional, default None
1137
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1138
+ specified by this expression.
1139
+ timezone : str, optional, default None
1140
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1141
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1142
+ """
1143
+ ...
1144
+
1145
+ @typing.overload
1146
+ def schedule(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1147
+ ...
1148
+
1149
+ 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):
1150
+ """
1151
+ Specifies the times when the flow should be run when running on a
1152
+ production scheduler.
1153
+
1154
+
1155
+ Parameters
1156
+ ----------
1157
+ hourly : bool, default False
1158
+ Run the workflow hourly.
1159
+ daily : bool, default True
1160
+ Run the workflow daily.
1161
+ weekly : bool, default False
1162
+ Run the workflow weekly.
1163
+ cron : str, optional, default None
1164
+ Run the workflow at [a custom Cron schedule](https://docs.aws.amazon.com/eventbridge/latest/userguide/scheduled-events.html#cron-expressions)
1165
+ specified by this expression.
1166
+ timezone : str, optional, default None
1167
+ Timezone on which the schedule runs (default: None). Currently supported only for Argo workflows,
1168
+ which accepts timezones in [IANA format](https://nodatime.org/TimeZones).
1169
+ """
1170
+ ...
1171
+
1207
1172
  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]]:
1208
1173
  """
1209
1174
  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)
@@ -1348,3 +1313,38 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1348
1313
  """
1349
1314
  ...
1350
1315
 
1316
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1317
+ """
1318
+ Specifies what flows belong to the same project.
1319
+
1320
+ A project-specific namespace is created for all flows that
1321
+ use the same `@project(name)`.
1322
+
1323
+
1324
+ Parameters
1325
+ ----------
1326
+ name : str
1327
+ Project name. Make sure that the name is unique amongst all
1328
+ projects that use the same production scheduler. The name may
1329
+ contain only lowercase alphanumeric characters and underscores.
1330
+
1331
+ branch : Optional[str], default None
1332
+ The branch to use. If not specified, the branch is set to
1333
+ `user.<username>` unless `production` is set to `True`. This can
1334
+ also be set on the command line using `--branch` as a top-level option.
1335
+ It is an error to specify `branch` in the decorator and on the command line.
1336
+
1337
+ production : bool, default False
1338
+ Whether or not the branch is the production branch. This can also be set on the
1339
+ command line using `--production` as a top-level option. It is an error to specify
1340
+ `production` in the decorator and on the command line.
1341
+ The project branch name will be:
1342
+ - if `branch` is specified:
1343
+ - if `production` is True: `prod.<branch>`
1344
+ - if `production` is False: `test.<branch>`
1345
+ - if `branch` is not specified:
1346
+ - if `production` is True: `prod`
1347
+ - if `production` is False: `user.<username>`
1348
+ """
1349
+ ...
1350
+