ob-metaflow-stubs 6.0.3.155__py2.py3-none-any.whl → 6.0.3.156__py2.py3-none-any.whl

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