metaflow-stubs 2.15.18__py2.py3-none-any.whl → 2.15.20__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 (155) hide show
  1. metaflow-stubs/__init__.pyi +512 -508
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +5 -5
  8. metaflow-stubs/client/filecache.pyi +3 -3
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +2 -2
  11. metaflow-stubs/flowspec.pyi +5 -5
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +3 -3
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +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 +21 -21
  21. metaflow-stubs/metaflow_git.pyi +2 -2
  22. metaflow-stubs/multicore_utils.pyi +2 -2
  23. metaflow-stubs/parameters.pyi +3 -3
  24. metaflow-stubs/plugins/__init__.pyi +14 -12
  25. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  26. metaflow-stubs/plugins/airflow/airflow_decorator.pyi +2 -2
  27. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  28. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  29. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  30. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  31. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  32. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  33. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  34. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  35. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  36. metaflow-stubs/plugins/argo/argo_workflows.pyi +7 -34
  37. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +4 -4
  38. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  39. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +4 -4
  40. metaflow-stubs/plugins/argo/exit_hooks.pyi +48 -0
  41. metaflow-stubs/plugins/aws/__init__.pyi +3 -3
  42. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  43. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  44. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  45. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  46. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  47. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  48. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  49. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +4 -4
  50. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  51. metaflow-stubs/plugins/aws/step_functions/dynamo_db_client.pyi +2 -2
  52. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  53. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  54. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  55. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  56. metaflow-stubs/plugins/aws/step_functions/step_functions_decorator.pyi +2 -2
  57. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +5 -5
  58. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +4 -4
  59. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  60. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  61. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  62. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +4 -4
  63. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  64. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  65. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  66. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  67. metaflow-stubs/plugins/cards/card_client.pyi +3 -3
  68. metaflow-stubs/plugins/cards/card_creator.pyi +4 -3
  69. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  70. metaflow-stubs/plugins/cards/card_decorator.pyi +13 -2
  71. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  72. metaflow-stubs/plugins/cards/card_modules/basic.pyi +2 -2
  73. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  74. metaflow-stubs/plugins/cards/card_modules/components.pyi +3 -3
  75. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  76. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  77. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +10 -2
  78. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  79. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  80. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  81. metaflow-stubs/plugins/catch_decorator.pyi +2 -2
  82. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  83. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  84. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  85. metaflow-stubs/plugins/datatools/s3/s3.pyi +5 -5
  86. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  87. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  88. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  89. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  90. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  91. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  92. metaflow-stubs/plugins/exit_hook/__init__.pyi +11 -0
  93. metaflow-stubs/plugins/exit_hook/exit_hook_decorator.pyi +20 -0
  94. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  96. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  97. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +4 -4
  98. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  99. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  100. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  101. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  102. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  103. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +3 -3
  104. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +2 -2
  105. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  106. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  107. metaflow-stubs/plugins/kubernetes/kubernetes_job.pyi +2 -2
  108. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  109. metaflow-stubs/plugins/kubernetes/spot_monitor_sidecar.pyi +2 -2
  110. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  111. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  112. metaflow-stubs/plugins/pypi/__init__.pyi +3 -3
  113. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  114. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  115. metaflow-stubs/plugins/pypi/parsers.pyi +2 -2
  116. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  117. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  118. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  119. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  120. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  121. metaflow-stubs/plugins/secrets/__init__.pyi +6 -2
  122. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  123. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +9 -49
  124. metaflow-stubs/plugins/secrets/secrets_func.pyi +31 -0
  125. metaflow-stubs/plugins/secrets/secrets_spec.pyi +42 -0
  126. metaflow-stubs/plugins/secrets/utils.pyi +28 -0
  127. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  128. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  129. metaflow-stubs/plugins/timeout_decorator.pyi +2 -2
  130. metaflow-stubs/plugins/uv/__init__.pyi +2 -2
  131. metaflow-stubs/plugins/uv/uv_environment.pyi +3 -3
  132. metaflow-stubs/pylint_wrapper.pyi +2 -2
  133. metaflow-stubs/runner/__init__.pyi +2 -2
  134. metaflow-stubs/runner/deployer.pyi +29 -29
  135. metaflow-stubs/runner/deployer_impl.pyi +3 -3
  136. metaflow-stubs/runner/metaflow_runner.pyi +4 -4
  137. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  138. metaflow-stubs/runner/nbrun.pyi +2 -2
  139. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  140. metaflow-stubs/runner/utils.pyi +4 -4
  141. metaflow-stubs/system/__init__.pyi +2 -2
  142. metaflow-stubs/system/system_logger.pyi +2 -2
  143. metaflow-stubs/system/system_monitor.pyi +2 -2
  144. metaflow-stubs/tagging_util.pyi +2 -2
  145. metaflow-stubs/tuple_util.pyi +2 -2
  146. metaflow-stubs/user_configs/__init__.pyi +2 -2
  147. metaflow-stubs/user_configs/config_decorators.pyi +4 -4
  148. metaflow-stubs/user_configs/config_options.pyi +3 -3
  149. metaflow-stubs/user_configs/config_parameters.pyi +7 -7
  150. metaflow-stubs/version.pyi +2 -2
  151. {metaflow_stubs-2.15.18.dist-info → metaflow_stubs-2.15.20.dist-info}/METADATA +2 -2
  152. metaflow_stubs-2.15.20.dist-info/RECORD +155 -0
  153. metaflow_stubs-2.15.18.dist-info/RECORD +0 -149
  154. {metaflow_stubs-2.15.18.dist-info → metaflow_stubs-2.15.20.dist-info}/WHEEL +0 -0
  155. {metaflow_stubs-2.15.18.dist-info → metaflow_stubs-2.15.20.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.15.18 #
4
- # Generated on 2025-06-19T23:01:30.833215 #
3
+ # MF version: 2.15.20 #
4
+ # Generated on 2025-07-10T18:07:14.479436 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -35,17 +35,17 @@ from .user_configs.config_parameters import ConfigValue as ConfigValue
35
35
  from .user_configs.config_parameters import config_expr as config_expr
36
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
37
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
38
- from . import metaflow_git as metaflow_git
39
- from . import events as events
40
38
  from . import tuple_util as tuple_util
39
+ from . import events as events
40
+ from . import metaflow_git as metaflow_git
41
41
  from . import runner as runner
42
42
  from . import plugins as plugins
43
43
  from .plugins.datatools.s3.s3 import S3 as S3
44
44
  from . import includefile as includefile
45
45
  from .includefile import IncludeFile as IncludeFile
46
- from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
47
46
  from .plugins.pypi.parsers import requirements_txt_parser as requirements_txt_parser
48
47
  from .plugins.pypi.parsers import conda_environment_yml_parser as conda_environment_yml_parser
48
+ from .plugins.pypi.parsers import pyproject_toml_parser as pyproject_toml_parser
49
49
  from . import cards as cards
50
50
  from . import client as client
51
51
  from .client.core import namespace as namespace
@@ -147,141 +147,200 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
147
147
  ...
148
148
 
149
149
  @typing.overload
150
- 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]]]:
151
- """
152
- Specifies environment variables to be set prior to the execution of a step.
153
-
154
-
155
- Parameters
156
- ----------
157
- vars : Dict[str, str], default {}
158
- Dictionary of environment variables to set.
159
- """
160
- ...
161
-
162
- @typing.overload
163
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
164
- ...
165
-
166
- @typing.overload
167
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
168
- ...
169
-
170
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
171
- """
172
- Specifies environment variables to be set prior to the execution of a step.
173
-
174
-
175
- Parameters
176
- ----------
177
- vars : Dict[str, str], default {}
178
- Dictionary of environment variables to set.
179
- """
180
- ...
181
-
182
- @typing.overload
183
- 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]]]:
150
+ 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, 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]]]:
184
151
  """
185
- Specifies the number of times the task corresponding
186
- to a step needs to be retried.
187
-
188
- This decorator is useful for handling transient errors, such as networking issues.
189
- If your task contains operations that can't be retried safely, e.g. database updates,
190
- it is advisable to annotate it with `@retry(times=0)`.
191
-
192
- This can be used in conjunction with the `@catch` decorator. The `@catch`
193
- decorator will execute a no-op task after all retries have been exhausted,
194
- ensuring that the flow execution can continue.
152
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
195
153
 
196
154
 
197
155
  Parameters
198
156
  ----------
199
- times : int, default 3
200
- Number of times to retry this task.
201
- minutes_between_retries : int, default 2
202
- Number of minutes between retries.
157
+ cpu : int, default 1
158
+ Number of CPUs required for this step. If `@resources` is
159
+ also present, the maximum value from all decorators is used.
160
+ gpu : int, default 0
161
+ Number of GPUs required for this step. If `@resources` is
162
+ also present, the maximum value from all decorators is used.
163
+ memory : int, default 4096
164
+ Memory size (in MB) required for this step. If
165
+ `@resources` is also present, the maximum value from all decorators is
166
+ used.
167
+ image : str, optional, default None
168
+ Docker image to use when launching on AWS Batch. If not specified, and
169
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
170
+ not, a default Docker image mapping to the current version of Python is used.
171
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
172
+ AWS Batch Job Queue to submit the job to.
173
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
174
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
175
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
176
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
177
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
178
+ shared_memory : int, optional, default None
179
+ The value for the size (in MiB) of the /dev/shm volume for this step.
180
+ This parameter maps to the `--shm-size` option in Docker.
181
+ max_swap : int, optional, default None
182
+ The total amount of swap memory (in MiB) a container can use for this
183
+ step. This parameter is translated to the `--memory-swap` option in
184
+ Docker where the value is the sum of the container memory plus the
185
+ `max_swap` value.
186
+ swappiness : int, optional, default None
187
+ This allows you to tune memory swappiness behavior for this step.
188
+ A swappiness value of 0 causes swapping not to happen unless absolutely
189
+ necessary. A swappiness value of 100 causes pages to be swapped very
190
+ aggressively. Accepted values are whole numbers between 0 and 100.
191
+ use_tmpfs : bool, default False
192
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
193
+ not available on Fargate compute environments
194
+ tmpfs_tempdir : bool, default True
195
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
196
+ tmpfs_size : int, optional, default None
197
+ The value for the size (in MiB) of the tmpfs mount for this step.
198
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
199
+ memory allocated for this step.
200
+ tmpfs_path : str, optional, default None
201
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
202
+ inferentia : int, default 0
203
+ Number of Inferentia chips required for this step.
204
+ trainium : int, default None
205
+ Alias for inferentia. Use only one of the two.
206
+ efa : int, default 0
207
+ Number of elastic fabric adapter network devices to attach to container
208
+ ephemeral_storage : int, default None
209
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
210
+ This is only relevant for Fargate compute environments
211
+ log_driver: str, optional, default None
212
+ The log driver to use for the Amazon ECS container.
213
+ log_options: List[str], optional, default None
214
+ List of strings containing options for the chosen log driver. The configurable values
215
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
216
+ Example: [`awslogs-group:aws/batch/job`]
203
217
  """
204
218
  ...
205
219
 
206
220
  @typing.overload
207
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
221
+ def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
208
222
  ...
209
223
 
210
224
  @typing.overload
211
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
225
+ def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
212
226
  ...
213
227
 
214
- def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
228
+ 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, 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):
215
229
  """
216
- Specifies the number of times the task corresponding
217
- to a step needs to be retried.
218
-
219
- This decorator is useful for handling transient errors, such as networking issues.
220
- If your task contains operations that can't be retried safely, e.g. database updates,
221
- it is advisable to annotate it with `@retry(times=0)`.
222
-
223
- This can be used in conjunction with the `@catch` decorator. The `@catch`
224
- decorator will execute a no-op task after all retries have been exhausted,
225
- ensuring that the flow execution can continue.
230
+ Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
226
231
 
227
232
 
228
233
  Parameters
229
234
  ----------
230
- times : int, default 3
231
- Number of times to retry this task.
232
- minutes_between_retries : int, default 2
233
- Number of minutes between retries.
235
+ cpu : int, default 1
236
+ Number of CPUs required for this step. If `@resources` is
237
+ also present, the maximum value from all decorators is used.
238
+ gpu : int, default 0
239
+ Number of GPUs required for this step. If `@resources` is
240
+ also present, the maximum value from all decorators is used.
241
+ memory : int, default 4096
242
+ Memory size (in MB) required for this step. If
243
+ `@resources` is also present, the maximum value from all decorators is
244
+ used.
245
+ image : str, optional, default None
246
+ Docker image to use when launching on AWS Batch. If not specified, and
247
+ METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
248
+ not, a default Docker image mapping to the current version of Python is used.
249
+ queue : str, default METAFLOW_BATCH_JOB_QUEUE
250
+ AWS Batch Job Queue to submit the job to.
251
+ iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
252
+ AWS IAM role that AWS Batch container uses to access AWS cloud resources.
253
+ execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
254
+ AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
255
+ (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
256
+ shared_memory : int, optional, default None
257
+ The value for the size (in MiB) of the /dev/shm volume for this step.
258
+ This parameter maps to the `--shm-size` option in Docker.
259
+ max_swap : int, optional, default None
260
+ The total amount of swap memory (in MiB) a container can use for this
261
+ step. This parameter is translated to the `--memory-swap` option in
262
+ Docker where the value is the sum of the container memory plus the
263
+ `max_swap` value.
264
+ swappiness : int, optional, default None
265
+ This allows you to tune memory swappiness behavior for this step.
266
+ A swappiness value of 0 causes swapping not to happen unless absolutely
267
+ necessary. A swappiness value of 100 causes pages to be swapped very
268
+ aggressively. Accepted values are whole numbers between 0 and 100.
269
+ use_tmpfs : bool, default False
270
+ This enables an explicit tmpfs mount for this step. Note that tmpfs is
271
+ not available on Fargate compute environments
272
+ tmpfs_tempdir : bool, default True
273
+ sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
274
+ tmpfs_size : int, optional, default None
275
+ The value for the size (in MiB) of the tmpfs mount for this step.
276
+ This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
277
+ memory allocated for this step.
278
+ tmpfs_path : str, optional, default None
279
+ Path to tmpfs mount for this step. Defaults to /metaflow_temp.
280
+ inferentia : int, default 0
281
+ Number of Inferentia chips required for this step.
282
+ trainium : int, default None
283
+ Alias for inferentia. Use only one of the two.
284
+ efa : int, default 0
285
+ Number of elastic fabric adapter network devices to attach to container
286
+ ephemeral_storage : int, default None
287
+ The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
288
+ This is only relevant for Fargate compute environments
289
+ log_driver: str, optional, default None
290
+ The log driver to use for the Amazon ECS container.
291
+ log_options: List[str], optional, default None
292
+ List of strings containing options for the chosen log driver. The configurable values
293
+ depend on the `log driver` chosen. Validation of these options is not supported yet.
294
+ Example: [`awslogs-group:aws/batch/job`]
234
295
  """
235
296
  ...
236
297
 
237
298
  @typing.overload
238
- 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]]]:
299
+ 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]]]:
239
300
  """
240
- Specifies the PyPI packages for the step.
301
+ Creates a human-readable report, a Metaflow Card, after this step completes.
241
302
 
242
- Information in this decorator will augment any
243
- attributes set in the `@pyi_base` flow-level decorator. Hence,
244
- you can use `@pypi_base` to set packages required by all
245
- steps and use `@pypi` to specify step-specific overrides.
303
+ Note that you may add multiple `@card` decorators in a step with different parameters.
246
304
 
247
305
 
248
306
  Parameters
249
307
  ----------
250
- packages : Dict[str, str], default: {}
251
- Packages to use for this step. The key is the name of the package
252
- and the value is the version to use.
253
- python : str, optional, default: None
254
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
255
- that the version used will correspond to the version of the Python interpreter used to start the run.
308
+ type : str, default 'default'
309
+ Card type.
310
+ id : str, optional, default None
311
+ If multiple cards are present, use this id to identify this card.
312
+ options : Dict[str, Any], default {}
313
+ Options passed to the card. The contents depend on the card type.
314
+ timeout : int, default 45
315
+ Interrupt reporting if it takes more than this many seconds.
256
316
  """
257
317
  ...
258
318
 
259
319
  @typing.overload
260
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
320
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
261
321
  ...
262
322
 
263
323
  @typing.overload
264
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
324
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
265
325
  ...
266
326
 
267
- 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):
327
+ 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):
268
328
  """
269
- Specifies the PyPI packages for the step.
329
+ Creates a human-readable report, a Metaflow Card, after this step completes.
270
330
 
271
- Information in this decorator will augment any
272
- attributes set in the `@pyi_base` flow-level decorator. Hence,
273
- you can use `@pypi_base` to set packages required by all
274
- steps and use `@pypi` to specify step-specific overrides.
331
+ Note that you may add multiple `@card` decorators in a step with different parameters.
275
332
 
276
333
 
277
334
  Parameters
278
335
  ----------
279
- packages : Dict[str, str], default: {}
280
- Packages to use for this step. The key is the name of the package
281
- and the value is the version to use.
282
- python : str, optional, default: None
283
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
284
- that the version used will correspond to the version of the Python interpreter used to start the run.
336
+ type : str, default 'default'
337
+ Card type.
338
+ id : str, optional, default None
339
+ If multiple cards are present, use this id to identify this card.
340
+ options : Dict[str, Any], default {}
341
+ Options passed to the card. The contents depend on the card type.
342
+ timeout : int, default 45
343
+ Interrupt reporting if it takes more than this many seconds.
285
344
  """
286
345
  ...
287
346
 
@@ -336,151 +395,117 @@ def catch(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
336
395
  """
337
396
  ...
338
397
 
339
- 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[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]]]:
398
+ @typing.overload
399
+ 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]]]:
340
400
  """
341
- Specifies that this step should execute on Kubernetes.
342
-
401
+ Specifies the number of times the task corresponding
402
+ to a step needs to be retried.
403
+
404
+ This decorator is useful for handling transient errors, such as networking issues.
405
+ If your task contains operations that can't be retried safely, e.g. database updates,
406
+ it is advisable to annotate it with `@retry(times=0)`.
407
+
408
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
409
+ decorator will execute a no-op task after all retries have been exhausted,
410
+ ensuring that the flow execution can continue.
411
+
343
412
 
344
413
  Parameters
345
414
  ----------
346
- cpu : int, default 1
347
- Number of CPUs required for this step. If `@resources` is
348
- also present, the maximum value from all decorators is used.
349
- memory : int, default 4096
350
- Memory size (in MB) required for this step. If
351
- `@resources` is also present, the maximum value from all decorators is
352
- used.
353
- disk : int, default 10240
354
- Disk size (in MB) required for this step. If
355
- `@resources` is also present, the maximum value from all decorators is
356
- used.
357
- image : str, optional, default None
358
- Docker image to use when launching on Kubernetes. If not specified, and
359
- METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
360
- not, a default Docker image mapping to the current version of Python is used.
361
- image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
362
- If given, the imagePullPolicy to be applied to the Docker image of the step.
363
- image_pull_secrets: List[str], default []
364
- The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
365
- Kubernetes image pull secrets to use when pulling container images
366
- in Kubernetes.
367
- service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
368
- Kubernetes service account to use when launching pod in Kubernetes.
369
- secrets : List[str], optional, default None
370
- Kubernetes secrets to use when launching pod in Kubernetes. These
371
- secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
372
- in Metaflow configuration.
373
- node_selector: Union[Dict[str,str], str], optional, default None
374
- Kubernetes node selector(s) to apply to the pod running the task.
375
- Can be passed in as a comma separated string of values e.g.
376
- 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
377
- {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
378
- namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
379
- Kubernetes namespace to use when launching pod in Kubernetes.
380
- gpu : int, optional, default None
381
- Number of GPUs required for this step. A value of zero implies that
382
- the scheduled node should not have GPUs.
383
- gpu_vendor : str, default KUBERNETES_GPU_VENDOR
384
- The vendor of the GPUs to be used for this step.
385
- tolerations : List[str], default []
386
- The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
387
- Kubernetes tolerations to use when launching pod in Kubernetes.
388
- labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
389
- Kubernetes labels to use when launching pod in Kubernetes.
390
- annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
391
- Kubernetes annotations to use when launching pod in Kubernetes.
392
- use_tmpfs : bool, default False
393
- This enables an explicit tmpfs mount for this step.
394
- tmpfs_tempdir : bool, default True
395
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
396
- tmpfs_size : int, optional, default: None
397
- The value for the size (in MiB) of the tmpfs mount for this step.
398
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
399
- memory allocated for this step.
400
- tmpfs_path : str, optional, default /metaflow_temp
401
- Path to tmpfs mount for this step.
402
- persistent_volume_claims : Dict[str, str], optional, default None
403
- A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
404
- volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
405
- shared_memory: int, optional
406
- Shared memory size (in MiB) required for this step
407
- port: int, optional
408
- Port number to specify in the Kubernetes job object
409
- compute_pool : str, optional, default None
410
- Compute pool to be used for for this step.
411
- If not specified, any accessible compute pool within the perimeter is used.
412
- hostname_resolution_timeout: int, default 10 * 60
413
- Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
414
- Only applicable when @parallel is used.
415
- qos: str, default: Burstable
416
- Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
417
-
418
- security_context: Dict[str, Any], optional, default None
419
- Container security context. Applies to the task container. Allows the following keys:
420
- - privileged: bool, optional, default None
421
- - allow_privilege_escalation: bool, optional, default None
422
- - run_as_user: int, optional, default None
423
- - run_as_group: int, optional, default None
424
- - run_as_non_root: bool, optional, default None
415
+ times : int, default 3
416
+ Number of times to retry this task.
417
+ minutes_between_retries : int, default 2
418
+ Number of minutes between retries.
425
419
  """
426
420
  ...
427
421
 
428
422
  @typing.overload
429
- 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]]]:
423
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
424
+ ...
425
+
426
+ @typing.overload
427
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
428
+ ...
429
+
430
+ def retry(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, times: int = 3, minutes_between_retries: int = 2):
430
431
  """
431
- Specifies a timeout for your step.
432
+ Specifies the number of times the task corresponding
433
+ to a step needs to be retried.
432
434
 
433
- This decorator is useful if this step may hang indefinitely.
435
+ This decorator is useful for handling transient errors, such as networking issues.
436
+ If your task contains operations that can't be retried safely, e.g. database updates,
437
+ it is advisable to annotate it with `@retry(times=0)`.
434
438
 
435
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
436
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
437
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
439
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
440
+ decorator will execute a no-op task after all retries have been exhausted,
441
+ ensuring that the flow execution can continue.
438
442
 
439
- Note that all the values specified in parameters are added together so if you specify
440
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
443
+
444
+ Parameters
445
+ ----------
446
+ times : int, default 3
447
+ Number of times to retry this task.
448
+ minutes_between_retries : int, default 2
449
+ Number of minutes between retries.
450
+ """
451
+ ...
452
+
453
+ @typing.overload
454
+ 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]]]:
455
+ """
456
+ Specifies the Conda environment for the step.
457
+
458
+ Information in this decorator will augment any
459
+ attributes set in the `@conda_base` flow-level decorator. Hence,
460
+ you can use `@conda_base` to set packages required by all
461
+ steps and use `@conda` to specify step-specific overrides.
441
462
 
442
463
 
443
464
  Parameters
444
465
  ----------
445
- seconds : int, default 0
446
- Number of seconds to wait prior to timing out.
447
- minutes : int, default 0
448
- Number of minutes to wait prior to timing out.
449
- hours : int, default 0
450
- Number of hours to wait prior to timing out.
466
+ packages : Dict[str, str], default {}
467
+ Packages to use for this step. The key is the name of the package
468
+ and the value is the version to use.
469
+ libraries : Dict[str, str], default {}
470
+ Supported for backward compatibility. When used with packages, packages will take precedence.
471
+ python : str, optional, default None
472
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
473
+ that the version used will correspond to the version of the Python interpreter used to start the run.
474
+ disabled : bool, default False
475
+ If set to True, disables @conda.
451
476
  """
452
477
  ...
453
478
 
454
479
  @typing.overload
455
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
480
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
456
481
  ...
457
482
 
458
483
  @typing.overload
459
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
484
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
460
485
  ...
461
486
 
462
- 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):
487
+ 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):
463
488
  """
464
- Specifies a timeout for your step.
465
-
466
- This decorator is useful if this step may hang indefinitely.
467
-
468
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
469
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
470
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
489
+ Specifies the Conda environment for the step.
471
490
 
472
- Note that all the values specified in parameters are added together so if you specify
473
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
491
+ Information in this decorator will augment any
492
+ attributes set in the `@conda_base` flow-level decorator. Hence,
493
+ you can use `@conda_base` to set packages required by all
494
+ steps and use `@conda` to specify step-specific overrides.
474
495
 
475
496
 
476
497
  Parameters
477
498
  ----------
478
- seconds : int, default 0
479
- Number of seconds to wait prior to timing out.
480
- minutes : int, default 0
481
- Number of minutes to wait prior to timing out.
482
- hours : int, default 0
483
- Number of hours to wait prior to timing out.
499
+ packages : Dict[str, str], default {}
500
+ Packages to use for this step. The key is the name of the package
501
+ and the value is the version to use.
502
+ libraries : Dict[str, str], default {}
503
+ Supported for backward compatibility. When used with packages, packages will take precedence.
504
+ python : str, optional, default None
505
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
506
+ that the version used will correspond to the version of the Python interpreter used to start the run.
507
+ disabled : bool, default False
508
+ If set to True, disables @conda.
484
509
  """
485
510
  ...
486
511
 
@@ -563,10 +588,9 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
563
588
  """
564
589
  ...
565
590
 
566
- @typing.overload
567
- 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, 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]]]:
591
+ 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[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]]]:
568
592
  """
569
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
593
+ Specifies that this step should execute on Kubernetes.
570
594
 
571
595
 
572
596
  Parameters
@@ -574,165 +598,87 @@ def batch(*, cpu: int = 1, gpu: int = 0, memory: int = 4096, image: typing.Optio
574
598
  cpu : int, default 1
575
599
  Number of CPUs required for this step. If `@resources` is
576
600
  also present, the maximum value from all decorators is used.
577
- gpu : int, default 0
578
- Number of GPUs required for this step. If `@resources` is
579
- also present, the maximum value from all decorators is used.
580
601
  memory : int, default 4096
581
602
  Memory size (in MB) required for this step. If
582
603
  `@resources` is also present, the maximum value from all decorators is
583
604
  used.
584
- image : str, optional, default None
585
- Docker image to use when launching on AWS Batch. If not specified, and
586
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
587
- not, a default Docker image mapping to the current version of Python is used.
588
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
589
- AWS Batch Job Queue to submit the job to.
590
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
591
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
592
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
593
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
594
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
595
- shared_memory : int, optional, default None
596
- The value for the size (in MiB) of the /dev/shm volume for this step.
597
- This parameter maps to the `--shm-size` option in Docker.
598
- max_swap : int, optional, default None
599
- The total amount of swap memory (in MiB) a container can use for this
600
- step. This parameter is translated to the `--memory-swap` option in
601
- Docker where the value is the sum of the container memory plus the
602
- `max_swap` value.
603
- swappiness : int, optional, default None
604
- This allows you to tune memory swappiness behavior for this step.
605
- A swappiness value of 0 causes swapping not to happen unless absolutely
606
- necessary. A swappiness value of 100 causes pages to be swapped very
607
- aggressively. Accepted values are whole numbers between 0 and 100.
608
- use_tmpfs : bool, default False
609
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
610
- not available on Fargate compute environments
611
- tmpfs_tempdir : bool, default True
612
- sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
613
- tmpfs_size : int, optional, default None
614
- The value for the size (in MiB) of the tmpfs mount for this step.
615
- This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
616
- memory allocated for this step.
617
- tmpfs_path : str, optional, default None
618
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
619
- inferentia : int, default 0
620
- Number of Inferentia chips required for this step.
621
- trainium : int, default None
622
- Alias for inferentia. Use only one of the two.
623
- efa : int, default 0
624
- Number of elastic fabric adapter network devices to attach to container
625
- ephemeral_storage : int, default None
626
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
627
- This is only relevant for Fargate compute environments
628
- log_driver: str, optional, default None
629
- The log driver to use for the Amazon ECS container.
630
- log_options: List[str], optional, default None
631
- List of strings containing options for the chosen log driver. The configurable values
632
- depend on the `log driver` chosen. Validation of these options is not supported yet.
633
- Example: [`awslogs-group:aws/batch/job`]
634
- """
635
- ...
636
-
637
- @typing.overload
638
- def batch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
639
- ...
640
-
641
- @typing.overload
642
- def batch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
643
- ...
644
-
645
- 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, 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):
646
- """
647
- Specifies that this step should execute on [AWS Batch](https://aws.amazon.com/batch/).
648
-
649
-
650
- Parameters
651
- ----------
652
- cpu : int, default 1
653
- Number of CPUs required for this step. If `@resources` is
654
- also present, the maximum value from all decorators is used.
655
- gpu : int, default 0
656
- Number of GPUs required for this step. If `@resources` is
657
- also present, the maximum value from all decorators is used.
658
- memory : int, default 4096
659
- Memory size (in MB) required for this step. If
605
+ disk : int, default 10240
606
+ Disk size (in MB) required for this step. If
660
607
  `@resources` is also present, the maximum value from all decorators is
661
608
  used.
662
609
  image : str, optional, default None
663
- Docker image to use when launching on AWS Batch. If not specified, and
664
- METAFLOW_BATCH_CONTAINER_IMAGE is specified, that image is used. If
610
+ Docker image to use when launching on Kubernetes. If not specified, and
611
+ METAFLOW_KUBERNETES_CONTAINER_IMAGE is specified, that image is used. If
665
612
  not, a default Docker image mapping to the current version of Python is used.
666
- queue : str, default METAFLOW_BATCH_JOB_QUEUE
667
- AWS Batch Job Queue to submit the job to.
668
- iam_role : str, default METAFLOW_ECS_S3_ACCESS_IAM_ROLE
669
- AWS IAM role that AWS Batch container uses to access AWS cloud resources.
670
- execution_role : str, default METAFLOW_ECS_FARGATE_EXECUTION_ROLE
671
- AWS IAM role that AWS Batch can use [to trigger AWS Fargate tasks]
672
- (https://docs.aws.amazon.com/batch/latest/userguide/execution-IAM-role.html).
673
- shared_memory : int, optional, default None
674
- The value for the size (in MiB) of the /dev/shm volume for this step.
675
- This parameter maps to the `--shm-size` option in Docker.
676
- max_swap : int, optional, default None
677
- The total amount of swap memory (in MiB) a container can use for this
678
- step. This parameter is translated to the `--memory-swap` option in
679
- Docker where the value is the sum of the container memory plus the
680
- `max_swap` value.
681
- swappiness : int, optional, default None
682
- This allows you to tune memory swappiness behavior for this step.
683
- A swappiness value of 0 causes swapping not to happen unless absolutely
684
- necessary. A swappiness value of 100 causes pages to be swapped very
685
- aggressively. Accepted values are whole numbers between 0 and 100.
613
+ image_pull_policy: str, default KUBERNETES_IMAGE_PULL_POLICY
614
+ If given, the imagePullPolicy to be applied to the Docker image of the step.
615
+ image_pull_secrets: List[str], default []
616
+ The default is extracted from METAFLOW_KUBERNETES_IMAGE_PULL_SECRETS.
617
+ Kubernetes image pull secrets to use when pulling container images
618
+ in Kubernetes.
619
+ service_account : str, default METAFLOW_KUBERNETES_SERVICE_ACCOUNT
620
+ Kubernetes service account to use when launching pod in Kubernetes.
621
+ secrets : List[str], optional, default None
622
+ Kubernetes secrets to use when launching pod in Kubernetes. These
623
+ secrets are in addition to the ones defined in `METAFLOW_KUBERNETES_SECRETS`
624
+ in Metaflow configuration.
625
+ node_selector: Union[Dict[str,str], str], optional, default None
626
+ Kubernetes node selector(s) to apply to the pod running the task.
627
+ Can be passed in as a comma separated string of values e.g.
628
+ 'kubernetes.io/os=linux,kubernetes.io/arch=amd64' or as a dictionary
629
+ {'kubernetes.io/os': 'linux', 'kubernetes.io/arch': 'amd64'}
630
+ namespace : str, default METAFLOW_KUBERNETES_NAMESPACE
631
+ Kubernetes namespace to use when launching pod in Kubernetes.
632
+ gpu : int, optional, default None
633
+ Number of GPUs required for this step. A value of zero implies that
634
+ the scheduled node should not have GPUs.
635
+ gpu_vendor : str, default KUBERNETES_GPU_VENDOR
636
+ The vendor of the GPUs to be used for this step.
637
+ tolerations : List[str], default []
638
+ The default is extracted from METAFLOW_KUBERNETES_TOLERATIONS.
639
+ Kubernetes tolerations to use when launching pod in Kubernetes.
640
+ labels: Dict[str, str], default: METAFLOW_KUBERNETES_LABELS
641
+ Kubernetes labels to use when launching pod in Kubernetes.
642
+ annotations: Dict[str, str], default: METAFLOW_KUBERNETES_ANNOTATIONS
643
+ Kubernetes annotations to use when launching pod in Kubernetes.
686
644
  use_tmpfs : bool, default False
687
- This enables an explicit tmpfs mount for this step. Note that tmpfs is
688
- not available on Fargate compute environments
645
+ This enables an explicit tmpfs mount for this step.
689
646
  tmpfs_tempdir : bool, default True
690
647
  sets METAFLOW_TEMPDIR to tmpfs_path if set for this step.
691
- tmpfs_size : int, optional, default None
648
+ tmpfs_size : int, optional, default: None
692
649
  The value for the size (in MiB) of the tmpfs mount for this step.
693
650
  This parameter maps to the `--tmpfs` option in Docker. Defaults to 50% of the
694
651
  memory allocated for this step.
695
- tmpfs_path : str, optional, default None
696
- Path to tmpfs mount for this step. Defaults to /metaflow_temp.
697
- inferentia : int, default 0
698
- Number of Inferentia chips required for this step.
699
- trainium : int, default None
700
- Alias for inferentia. Use only one of the two.
701
- efa : int, default 0
702
- Number of elastic fabric adapter network devices to attach to container
703
- ephemeral_storage : int, default None
704
- The total amount, in GiB, of ephemeral storage to set for the task, 21-200GiB.
705
- This is only relevant for Fargate compute environments
706
- log_driver: str, optional, default None
707
- The log driver to use for the Amazon ECS container.
708
- log_options: List[str], optional, default None
709
- List of strings containing options for the chosen log driver. The configurable values
710
- depend on the `log driver` chosen. Validation of these options is not supported yet.
711
- Example: [`awslogs-group:aws/batch/job`]
712
- """
713
- ...
714
-
715
- @typing.overload
716
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
717
- """
718
- Decorator prototype for all step decorators. This function gets specialized
719
- and imported for all decorators types by _import_plugin_decorators().
720
- """
721
- ...
722
-
723
- @typing.overload
724
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
725
- ...
726
-
727
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
728
- """
729
- Decorator prototype for all step decorators. This function gets specialized
730
- and imported for all decorators types by _import_plugin_decorators().
652
+ tmpfs_path : str, optional, default /metaflow_temp
653
+ Path to tmpfs mount for this step.
654
+ persistent_volume_claims : Dict[str, str], optional, default None
655
+ A map (dictionary) of persistent volumes to be mounted to the pod for this step. The map is from persistent
656
+ volumes to the path to which the volume is to be mounted, e.g., `{'pvc-name': '/path/to/mount/on'}`.
657
+ shared_memory: int, optional
658
+ Shared memory size (in MiB) required for this step
659
+ port: int, optional
660
+ Port number to specify in the Kubernetes job object
661
+ compute_pool : str, optional, default None
662
+ Compute pool to be used for for this step.
663
+ If not specified, any accessible compute pool within the perimeter is used.
664
+ hostname_resolution_timeout: int, default 10 * 60
665
+ Timeout in seconds for the workers tasks in the gang scheduled cluster to resolve the hostname of control task.
666
+ Only applicable when @parallel is used.
667
+ qos: str, default: Burstable
668
+ Quality of Service class to assign to the pod. Supported values are: Guaranteed, Burstable, BestEffort
669
+
670
+ security_context: Dict[str, Any], optional, default None
671
+ Container security context. Applies to the task container. Allows the following keys:
672
+ - privileged: bool, optional, default None
673
+ - allow_privilege_escalation: bool, optional, default None
674
+ - run_as_user: int, optional, default None
675
+ - run_as_group: int, optional, default None
676
+ - run_as_non_root: bool, optional, default None
731
677
  """
732
678
  ...
733
679
 
734
680
  @typing.overload
735
- def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
681
+ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: 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]]]:
736
682
  """
737
683
  Specifies secrets to be retrieved and injected as environment variables prior to
738
684
  the execution of a step.
@@ -742,6 +688,8 @@ def secrets(*, sources: typing.List[typing.Union[str, typing.Dict[str, typing.An
742
688
  ----------
743
689
  sources : List[Union[str, Dict[str, Any]]], default: []
744
690
  List of secret specs, defining how the secrets are to be retrieved
691
+ role : str, optional, default: None
692
+ Role to use for fetching secrets
745
693
  """
746
694
  ...
747
695
 
@@ -753,7 +701,7 @@ def secrets(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Cal
753
701
  def secrets(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
754
702
  ...
755
703
 
756
- def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = []):
704
+ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, sources: typing.List[typing.Union[str, typing.Dict[str, typing.Any]]] = [], role: typing.Optional[str] = None):
757
705
  """
758
706
  Specifies secrets to be retrieved and injected as environment variables prior to
759
707
  the execution of a step.
@@ -763,165 +711,170 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
763
711
  ----------
764
712
  sources : List[Union[str, Dict[str, Any]]], default: []
765
713
  List of secret specs, defining how the secrets are to be retrieved
714
+ role : str, optional, default: None
715
+ Role to use for fetching secrets
766
716
  """
767
717
  ...
768
718
 
769
719
  @typing.overload
770
- 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]]]:
720
+ 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]]]:
771
721
  """
772
- Specifies the Conda environment for the step.
773
-
774
- Information in this decorator will augment any
775
- attributes set in the `@conda_base` flow-level decorator. Hence,
776
- you can use `@conda_base` to set packages required by all
777
- steps and use `@conda` to specify step-specific overrides.
722
+ Specifies environment variables to be set prior to the execution of a step.
778
723
 
779
724
 
780
725
  Parameters
781
726
  ----------
782
- packages : Dict[str, str], default {}
783
- Packages to use for this step. The key is the name of the package
784
- and the value is the version to use.
785
- libraries : Dict[str, str], default {}
786
- Supported for backward compatibility. When used with packages, packages will take precedence.
787
- python : str, optional, default None
788
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
789
- that the version used will correspond to the version of the Python interpreter used to start the run.
790
- disabled : bool, default False
791
- If set to True, disables @conda.
727
+ vars : Dict[str, str], default {}
728
+ Dictionary of environment variables to set.
792
729
  """
793
730
  ...
794
731
 
795
732
  @typing.overload
796
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
733
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
797
734
  ...
798
735
 
799
736
  @typing.overload
800
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
737
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
801
738
  ...
802
739
 
803
- 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):
740
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
804
741
  """
805
- Specifies the Conda environment for the step.
806
-
807
- Information in this decorator will augment any
808
- attributes set in the `@conda_base` flow-level decorator. Hence,
809
- you can use `@conda_base` to set packages required by all
810
- steps and use `@conda` to specify step-specific overrides.
742
+ Specifies environment variables to be set prior to the execution of a step.
811
743
 
812
744
 
813
745
  Parameters
814
746
  ----------
815
- packages : Dict[str, str], default {}
816
- Packages to use for this step. The key is the name of the package
817
- and the value is the version to use.
818
- libraries : Dict[str, str], default {}
819
- Supported for backward compatibility. When used with packages, packages will take precedence.
820
- python : str, optional, default None
821
- Version of Python to use, e.g. '3.7.4'. A default value of None implies
822
- that the version used will correspond to the version of the Python interpreter used to start the run.
823
- disabled : bool, default False
824
- If set to True, disables @conda.
747
+ vars : Dict[str, str], default {}
748
+ Dictionary of environment variables to set.
825
749
  """
826
750
  ...
827
751
 
828
752
  @typing.overload
829
- 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]]]:
753
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
830
754
  """
831
- Creates a human-readable report, a Metaflow Card, after this step completes.
755
+ Decorator prototype for all step decorators. This function gets specialized
756
+ and imported for all decorators types by _import_plugin_decorators().
757
+ """
758
+ ...
759
+
760
+ @typing.overload
761
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
762
+ ...
763
+
764
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
765
+ """
766
+ Decorator prototype for all step decorators. This function gets specialized
767
+ and imported for all decorators types by _import_plugin_decorators().
768
+ """
769
+ ...
770
+
771
+ @typing.overload
772
+ 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]]]:
773
+ """
774
+ Specifies a timeout for your step.
832
775
 
833
- Note that you may add multiple `@card` decorators in a step with different parameters.
776
+ This decorator is useful if this step may hang indefinitely.
777
+
778
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
779
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
780
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
781
+
782
+ Note that all the values specified in parameters are added together so if you specify
783
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
834
784
 
835
785
 
836
786
  Parameters
837
787
  ----------
838
- type : str, default 'default'
839
- Card type.
840
- id : str, optional, default None
841
- If multiple cards are present, use this id to identify this card.
842
- options : Dict[str, Any], default {}
843
- Options passed to the card. The contents depend on the card type.
844
- timeout : int, default 45
845
- Interrupt reporting if it takes more than this many seconds.
788
+ seconds : int, default 0
789
+ Number of seconds to wait prior to timing out.
790
+ minutes : int, default 0
791
+ Number of minutes to wait prior to timing out.
792
+ hours : int, default 0
793
+ Number of hours to wait prior to timing out.
846
794
  """
847
795
  ...
848
796
 
849
797
  @typing.overload
850
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
798
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
851
799
  ...
852
800
 
853
801
  @typing.overload
854
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
802
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
855
803
  ...
856
804
 
857
- 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):
805
+ 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):
858
806
  """
859
- Creates a human-readable report, a Metaflow Card, after this step completes.
807
+ Specifies a timeout for your step.
808
+
809
+ This decorator is useful if this step may hang indefinitely.
860
810
 
861
- Note that you may add multiple `@card` decorators in a step with different parameters.
811
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
812
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
813
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
814
+
815
+ Note that all the values specified in parameters are added together so if you specify
816
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
862
817
 
863
818
 
864
819
  Parameters
865
820
  ----------
866
- type : str, default 'default'
867
- Card type.
868
- id : str, optional, default None
869
- If multiple cards are present, use this id to identify this card.
870
- options : Dict[str, Any], default {}
871
- Options passed to the card. The contents depend on the card type.
872
- timeout : int, default 45
873
- Interrupt reporting if it takes more than this many seconds.
821
+ seconds : int, default 0
822
+ Number of seconds to wait prior to timing out.
823
+ minutes : int, default 0
824
+ Number of minutes to wait prior to timing out.
825
+ hours : int, default 0
826
+ Number of hours to wait prior to timing out.
874
827
  """
875
828
  ...
876
829
 
877
830
  @typing.overload
878
- def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
831
+ 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]]]:
879
832
  """
880
- Specifies the Conda environment for all steps of the flow.
833
+ Specifies the PyPI packages for the step.
881
834
 
882
- Use `@conda_base` to set common libraries required by all
883
- steps and use `@conda` to specify step-specific additions.
835
+ Information in this decorator will augment any
836
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
837
+ you can use `@pypi_base` to set packages required by all
838
+ steps and use `@pypi` to specify step-specific overrides.
884
839
 
885
840
 
886
841
  Parameters
887
842
  ----------
888
- packages : Dict[str, str], default {}
889
- Packages to use for this flow. The key is the name of the package
843
+ packages : Dict[str, str], default: {}
844
+ Packages to use for this step. The key is the name of the package
890
845
  and the value is the version to use.
891
- libraries : Dict[str, str], default {}
892
- Supported for backward compatibility. When used with packages, packages will take precedence.
893
- python : str, optional, default None
846
+ python : str, optional, default: None
894
847
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
895
848
  that the version used will correspond to the version of the Python interpreter used to start the run.
896
- disabled : bool, default False
897
- If set to True, disables Conda.
898
849
  """
899
850
  ...
900
851
 
901
852
  @typing.overload
902
- def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
853
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
903
854
  ...
904
855
 
905
- def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
856
+ @typing.overload
857
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
858
+ ...
859
+
860
+ 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):
906
861
  """
907
- Specifies the Conda environment for all steps of the flow.
862
+ Specifies the PyPI packages for the step.
908
863
 
909
- Use `@conda_base` to set common libraries required by all
910
- steps and use `@conda` to specify step-specific additions.
864
+ Information in this decorator will augment any
865
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
866
+ you can use `@pypi_base` to set packages required by all
867
+ steps and use `@pypi` to specify step-specific overrides.
911
868
 
912
869
 
913
870
  Parameters
914
871
  ----------
915
- packages : Dict[str, str], default {}
916
- Packages to use for this flow. The key is the name of the package
872
+ packages : Dict[str, str], default: {}
873
+ Packages to use for this step. The key is the name of the package
917
874
  and the value is the version to use.
918
- libraries : Dict[str, str], default {}
919
- Supported for backward compatibility. When used with packages, packages will take precedence.
920
- python : str, optional, default None
875
+ python : str, optional, default: None
921
876
  Version of Python to use, e.g. '3.7.4'. A default value of None implies
922
877
  that the version used will correspond to the version of the Python interpreter used to start the run.
923
- disabled : bool, default False
924
- If set to True, disables Conda.
925
878
  """
926
879
  ...
927
880
 
@@ -1119,38 +1072,46 @@ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *
1119
1072
  """
1120
1073
  ...
1121
1074
 
1122
- def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1075
+ 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]]:
1123
1076
  """
1124
- Specifies what flows belong to the same project.
1125
-
1126
- A project-specific namespace is created for all flows that
1127
- use the same `@project(name)`.
1077
+ 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)
1078
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1079
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1080
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
1081
+ starts only after all sensors finish.
1128
1082
 
1129
1083
 
1130
1084
  Parameters
1131
1085
  ----------
1086
+ timeout : int
1087
+ Time, in seconds before the task times out and fails. (Default: 3600)
1088
+ poke_interval : int
1089
+ Time in seconds that the job should wait in between each try. (Default: 60)
1090
+ mode : str
1091
+ How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1092
+ exponential_backoff : bool
1093
+ allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1094
+ pool : str
1095
+ the slot pool this task should run in,
1096
+ slot pools are a way to limit concurrency for certain tasks. (Default:None)
1097
+ soft_fail : bool
1098
+ Set to true to mark the task as SKIPPED on failure. (Default: False)
1132
1099
  name : str
1133
- Project name. Make sure that the name is unique amongst all
1134
- projects that use the same production scheduler. The name may
1135
- contain only lowercase alphanumeric characters and underscores.
1136
-
1137
- branch : Optional[str], default None
1138
- The branch to use. If not specified, the branch is set to
1139
- `user.<username>` unless `production` is set to `True`. This can
1140
- also be set on the command line using `--branch` as a top-level option.
1141
- It is an error to specify `branch` in the decorator and on the command line.
1142
-
1143
- production : bool, default False
1144
- Whether or not the branch is the production branch. This can also be set on the
1145
- command line using `--production` as a top-level option. It is an error to specify
1146
- `production` in the decorator and on the command line.
1147
- The project branch name will be:
1148
- - if `branch` is specified:
1149
- - if `production` is True: `prod.<branch>`
1150
- - if `production` is False: `test.<branch>`
1151
- - if `branch` is not specified:
1152
- - if `production` is True: `prod`
1153
- - if `production` is False: `user.<username>`
1100
+ Name of the sensor on Airflow
1101
+ description : str
1102
+ Description of sensor in the Airflow UI
1103
+ bucket_key : Union[str, List[str]]
1104
+ The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1105
+ When it's specified as a full s3:// url, please leave `bucket_name` as None
1106
+ bucket_name : str
1107
+ Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1108
+ When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1109
+ wildcard_match : bool
1110
+ whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1111
+ aws_conn_id : str
1112
+ a reference to the s3 connection on Airflow. (Default: None)
1113
+ verify : bool
1114
+ Whether or not to verify SSL certificates for S3 connection. (Default: None)
1154
1115
  """
1155
1116
  ...
1156
1117
 
@@ -1246,6 +1207,92 @@ def schedule(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, hourly:
1246
1207
  """
1247
1208
  ...
1248
1209
 
1210
+ @typing.overload
1211
+ def conda_base(*, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1212
+ """
1213
+ Specifies the Conda environment for all steps of the flow.
1214
+
1215
+ Use `@conda_base` to set common libraries required by all
1216
+ steps and use `@conda` to specify step-specific additions.
1217
+
1218
+
1219
+ Parameters
1220
+ ----------
1221
+ packages : Dict[str, str], default {}
1222
+ Packages to use for this flow. The key is the name of the package
1223
+ and the value is the version to use.
1224
+ libraries : Dict[str, str], default {}
1225
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1226
+ python : str, optional, default None
1227
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1228
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1229
+ disabled : bool, default False
1230
+ If set to True, disables Conda.
1231
+ """
1232
+ ...
1233
+
1234
+ @typing.overload
1235
+ def conda_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
1236
+ ...
1237
+
1238
+ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, libraries: typing.Dict[str, str] = {}, python: typing.Optional[str] = None, disabled: bool = False):
1239
+ """
1240
+ Specifies the Conda environment for all steps of the flow.
1241
+
1242
+ Use `@conda_base` to set common libraries required by all
1243
+ steps and use `@conda` to specify step-specific additions.
1244
+
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
+ libraries : Dict[str, str], default {}
1252
+ Supported for backward compatibility. When used with packages, packages will take precedence.
1253
+ python : str, optional, default None
1254
+ Version of Python to use, e.g. '3.7.4'. A default value of None implies
1255
+ that the version used will correspond to the version of the Python interpreter used to start the run.
1256
+ disabled : bool, default False
1257
+ If set to True, disables Conda.
1258
+ """
1259
+ ...
1260
+
1261
+ def project(*, name: str, branch: typing.Optional[str] = None, production: bool = False) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
1262
+ """
1263
+ Specifies what flows belong to the same project.
1264
+
1265
+ A project-specific namespace is created for all flows that
1266
+ use the same `@project(name)`.
1267
+
1268
+
1269
+ Parameters
1270
+ ----------
1271
+ name : str
1272
+ Project name. Make sure that the name is unique amongst all
1273
+ projects that use the same production scheduler. The name may
1274
+ contain only lowercase alphanumeric characters and underscores.
1275
+
1276
+ branch : Optional[str], default None
1277
+ The branch to use. If not specified, the branch is set to
1278
+ `user.<username>` unless `production` is set to `True`. This can
1279
+ also be set on the command line using `--branch` as a top-level option.
1280
+ It is an error to specify `branch` in the decorator and on the command line.
1281
+
1282
+ production : bool, default False
1283
+ Whether or not the branch is the production branch. This can also be set on the
1284
+ command line using `--production` as a top-level option. It is an error to specify
1285
+ `production` in the decorator and on the command line.
1286
+ The project branch name will be:
1287
+ - if `branch` is specified:
1288
+ - if `production` is True: `prod.<branch>`
1289
+ - if `production` is False: `test.<branch>`
1290
+ - if `branch` is not specified:
1291
+ - if `production` is True: `prod`
1292
+ - if `production` is False: `user.<username>`
1293
+ """
1294
+ ...
1295
+
1249
1296
  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]]:
1250
1297
  """
1251
1298
  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.
@@ -1289,46 +1336,3 @@ def airflow_external_task_sensor(*, timeout: int, poke_interval: int, mode: str,
1289
1336
  """
1290
1337
  ...
1291
1338
 
1292
- 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]]:
1293
- """
1294
- 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)
1295
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
1296
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
1297
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
1298
- starts only after all sensors finish.
1299
-
1300
-
1301
- Parameters
1302
- ----------
1303
- timeout : int
1304
- Time, in seconds before the task times out and fails. (Default: 3600)
1305
- poke_interval : int
1306
- Time in seconds that the job should wait in between each try. (Default: 60)
1307
- mode : str
1308
- How the sensor operates. Options are: { poke | reschedule }. (Default: "poke")
1309
- exponential_backoff : bool
1310
- allow progressive longer waits between pokes by using exponential backoff algorithm. (Default: True)
1311
- pool : str
1312
- the slot pool this task should run in,
1313
- slot pools are a way to limit concurrency for certain tasks. (Default:None)
1314
- soft_fail : bool
1315
- Set to true to mark the task as SKIPPED on failure. (Default: False)
1316
- name : str
1317
- Name of the sensor on Airflow
1318
- description : str
1319
- Description of sensor in the Airflow UI
1320
- bucket_key : Union[str, List[str]]
1321
- The key(s) being waited on. Supports full s3:// style url or relative path from root level.
1322
- When it's specified as a full s3:// url, please leave `bucket_name` as None
1323
- bucket_name : str
1324
- Name of the S3 bucket. Only needed when bucket_key is not provided as a full s3:// url.
1325
- When specified, all the keys passed to bucket_key refers to this bucket. (Default:None)
1326
- wildcard_match : bool
1327
- whether the bucket_key should be interpreted as a Unix wildcard pattern. (Default: False)
1328
- aws_conn_id : str
1329
- a reference to the s3 connection on Airflow. (Default: None)
1330
- verify : bool
1331
- Whether or not to verify SSL certificates for S3 connection. (Default: None)
1332
- """
1333
- ...
1334
-