ob-metaflow-stubs 6.0.3.128__py2.py3-none-any.whl → 6.0.3.129__py2.py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (191) hide show
  1. metaflow-stubs/__init__.pyi +320 -319
  2. metaflow-stubs/cards.pyi +2 -2
  3. metaflow-stubs/cli.pyi +2 -2
  4. metaflow-stubs/cli_components/__init__.pyi +2 -2
  5. metaflow-stubs/cli_components/utils.pyi +2 -2
  6. metaflow-stubs/client/__init__.pyi +2 -2
  7. metaflow-stubs/client/core.pyi +6 -6
  8. metaflow-stubs/client/filecache.pyi +2 -2
  9. metaflow-stubs/events.pyi +2 -2
  10. metaflow-stubs/exception.pyi +3 -3
  11. metaflow-stubs/flowspec.pyi +8 -8
  12. metaflow-stubs/generated_for.txt +1 -1
  13. metaflow-stubs/includefile.pyi +5 -5
  14. metaflow-stubs/info_file.pyi +2 -2
  15. metaflow-stubs/metadata_provider/__init__.pyi +2 -2
  16. metaflow-stubs/metadata_provider/heartbeat.pyi +2 -2
  17. metaflow-stubs/metadata_provider/metadata.pyi +2 -2
  18. metaflow-stubs/metadata_provider/util.pyi +2 -2
  19. metaflow-stubs/metaflow_config.pyi +4 -2
  20. metaflow-stubs/metaflow_current.pyi +136 -136
  21. metaflow-stubs/mf_extensions/__init__.pyi +2 -2
  22. metaflow-stubs/mf_extensions/obcheckpoint/__init__.pyi +2 -2
  23. metaflow-stubs/mf_extensions/obcheckpoint/plugins/__init__.pyi +2 -2
  24. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/__init__.pyi +2 -2
  25. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/__init__.pyi +2 -2
  26. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/async_cards.pyi +3 -3
  27. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/deco_injection_mixin.pyi +2 -2
  28. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/card_utils/extra_components.pyi +4 -4
  29. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/__init__.pyi +2 -2
  30. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/__init__.pyi +2 -2
  31. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/checkpoint_lister.pyi +4 -4
  32. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/cards/lineage_card.pyi +2 -2
  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 +2 -2
  35. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/core.pyi +4 -4
  36. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/decorator.pyi +5 -5
  37. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/exceptions.pyi +2 -2
  38. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/final_api.pyi +3 -3
  39. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/checkpoints/lineage.pyi +2 -2
  40. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/__init__.pyi +2 -2
  41. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/core.pyi +3 -3
  42. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastore/exceptions.pyi +2 -2
  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 +2 -2
  45. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/datastructures.pyi +3 -3
  46. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/exceptions.pyi +2 -2
  47. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/hf_hub/__init__.pyi +2 -2
  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 +2 -2
  50. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/core.pyi +3 -3
  51. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/exceptions.pyi +2 -2
  52. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/modeling_utils/model_storage.pyi +5 -5
  53. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/__init__.pyi +2 -2
  54. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/flowspec_utils.pyi +2 -2
  55. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/general.pyi +2 -2
  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 +2 -2
  58. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/base.pyi +2 -2
  59. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/serialization_handler/tar.pyi +3 -3
  60. metaflow-stubs/mf_extensions/obcheckpoint/plugins/machine_learning_utilities/utils/tar_utils.pyi +3 -3
  61. metaflow-stubs/mf_extensions/outerbounds/__init__.pyi +2 -2
  62. metaflow-stubs/mf_extensions/outerbounds/plugins/__init__.pyi +2 -2
  63. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/__init__.pyi +2 -2
  64. metaflow-stubs/mf_extensions/outerbounds/plugins/snowflake/snowflake.pyi +2 -2
  65. metaflow-stubs/mf_extensions/outerbounds/profilers/__init__.pyi +2 -2
  66. metaflow-stubs/mf_extensions/outerbounds/profilers/gpu.pyi +2 -2
  67. metaflow-stubs/mf_extensions/outerbounds/remote_config.pyi +2 -2
  68. metaflow-stubs/mf_extensions/outerbounds/toplevel/__init__.pyi +2 -2
  69. metaflow-stubs/mf_extensions/outerbounds/toplevel/global_aliases_for_metaflow_package.pyi +2 -2
  70. metaflow-stubs/multicore_utils.pyi +2 -2
  71. metaflow-stubs/parameters.pyi +5 -5
  72. metaflow-stubs/plugins/__init__.pyi +13 -13
  73. metaflow-stubs/plugins/airflow/__init__.pyi +2 -2
  74. metaflow-stubs/plugins/airflow/airflow_utils.pyi +2 -2
  75. metaflow-stubs/plugins/airflow/exception.pyi +2 -2
  76. metaflow-stubs/plugins/airflow/sensors/__init__.pyi +2 -2
  77. metaflow-stubs/plugins/airflow/sensors/base_sensor.pyi +2 -2
  78. metaflow-stubs/plugins/airflow/sensors/external_task_sensor.pyi +2 -2
  79. metaflow-stubs/plugins/airflow/sensors/s3_sensor.pyi +2 -2
  80. metaflow-stubs/plugins/argo/__init__.pyi +2 -2
  81. metaflow-stubs/plugins/argo/argo_client.pyi +2 -2
  82. metaflow-stubs/plugins/argo/argo_events.pyi +2 -2
  83. metaflow-stubs/plugins/argo/argo_workflows.pyi +2 -2
  84. metaflow-stubs/plugins/argo/argo_workflows_decorator.pyi +3 -3
  85. metaflow-stubs/plugins/argo/argo_workflows_deployer.pyi +4 -4
  86. metaflow-stubs/plugins/argo/argo_workflows_deployer_objects.pyi +4 -4
  87. metaflow-stubs/plugins/aws/__init__.pyi +2 -2
  88. metaflow-stubs/plugins/aws/aws_client.pyi +2 -2
  89. metaflow-stubs/plugins/aws/aws_utils.pyi +2 -2
  90. metaflow-stubs/plugins/aws/batch/__init__.pyi +2 -2
  91. metaflow-stubs/plugins/aws/batch/batch.pyi +2 -2
  92. metaflow-stubs/plugins/aws/batch/batch_client.pyi +2 -2
  93. metaflow-stubs/plugins/aws/batch/batch_decorator.pyi +2 -2
  94. metaflow-stubs/plugins/aws/secrets_manager/__init__.pyi +2 -2
  95. metaflow-stubs/plugins/aws/secrets_manager/aws_secrets_manager_secrets_provider.pyi +15 -13
  96. metaflow-stubs/plugins/aws/step_functions/__init__.pyi +2 -2
  97. metaflow-stubs/plugins/aws/step_functions/event_bridge_client.pyi +2 -2
  98. metaflow-stubs/plugins/aws/step_functions/schedule_decorator.pyi +2 -2
  99. metaflow-stubs/plugins/aws/step_functions/step_functions.pyi +2 -2
  100. metaflow-stubs/plugins/aws/step_functions/step_functions_client.pyi +2 -2
  101. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer.pyi +4 -4
  102. metaflow-stubs/plugins/aws/step_functions/step_functions_deployer_objects.pyi +3 -3
  103. metaflow-stubs/plugins/azure/__init__.pyi +2 -2
  104. metaflow-stubs/plugins/azure/azure_credential.pyi +2 -2
  105. metaflow-stubs/plugins/azure/azure_exceptions.pyi +2 -2
  106. metaflow-stubs/plugins/azure/azure_secret_manager_secrets_provider.pyi +3 -3
  107. metaflow-stubs/plugins/azure/azure_utils.pyi +2 -2
  108. metaflow-stubs/plugins/azure/blob_service_client_factory.pyi +2 -2
  109. metaflow-stubs/plugins/azure/includefile_support.pyi +2 -2
  110. metaflow-stubs/plugins/cards/__init__.pyi +2 -2
  111. metaflow-stubs/plugins/cards/card_client.pyi +2 -2
  112. metaflow-stubs/plugins/cards/card_creator.pyi +2 -2
  113. metaflow-stubs/plugins/cards/card_datastore.pyi +2 -2
  114. metaflow-stubs/plugins/cards/card_decorator.pyi +4 -2
  115. metaflow-stubs/plugins/cards/card_modules/__init__.pyi +2 -2
  116. metaflow-stubs/plugins/cards/card_modules/basic.pyi +3 -3
  117. metaflow-stubs/plugins/cards/card_modules/card.pyi +2 -2
  118. metaflow-stubs/plugins/cards/card_modules/components.pyi +4 -4
  119. metaflow-stubs/plugins/cards/card_modules/convert_to_native_type.pyi +2 -2
  120. metaflow-stubs/plugins/cards/card_modules/renderer_tools.pyi +2 -2
  121. metaflow-stubs/plugins/cards/card_modules/test_cards.pyi +2 -2
  122. metaflow-stubs/plugins/cards/card_resolver.pyi +2 -2
  123. metaflow-stubs/plugins/cards/component_serializer.pyi +2 -2
  124. metaflow-stubs/plugins/cards/exception.pyi +2 -2
  125. metaflow-stubs/plugins/catch_decorator.pyi +3 -3
  126. metaflow-stubs/plugins/datatools/__init__.pyi +2 -2
  127. metaflow-stubs/plugins/datatools/local.pyi +2 -2
  128. metaflow-stubs/plugins/datatools/s3/__init__.pyi +2 -2
  129. metaflow-stubs/plugins/datatools/s3/s3.pyi +4 -4
  130. metaflow-stubs/plugins/datatools/s3/s3tail.pyi +2 -2
  131. metaflow-stubs/plugins/datatools/s3/s3util.pyi +2 -2
  132. metaflow-stubs/plugins/debug_logger.pyi +2 -2
  133. metaflow-stubs/plugins/debug_monitor.pyi +2 -2
  134. metaflow-stubs/plugins/environment_decorator.pyi +2 -2
  135. metaflow-stubs/plugins/events_decorator.pyi +2 -2
  136. metaflow-stubs/plugins/frameworks/__init__.pyi +2 -2
  137. metaflow-stubs/plugins/frameworks/pytorch.pyi +2 -2
  138. metaflow-stubs/plugins/gcp/__init__.pyi +2 -2
  139. metaflow-stubs/plugins/gcp/gcp_secret_manager_secrets_provider.pyi +3 -3
  140. metaflow-stubs/plugins/gcp/gs_exceptions.pyi +2 -2
  141. metaflow-stubs/plugins/gcp/gs_storage_client_factory.pyi +2 -2
  142. metaflow-stubs/plugins/gcp/gs_utils.pyi +2 -2
  143. metaflow-stubs/plugins/gcp/includefile_support.pyi +2 -2
  144. metaflow-stubs/plugins/kubernetes/__init__.pyi +2 -2
  145. metaflow-stubs/plugins/kubernetes/kube_utils.pyi +2 -2
  146. metaflow-stubs/plugins/kubernetes/kubernetes.pyi +3 -3
  147. metaflow-stubs/plugins/kubernetes/kubernetes_client.pyi +2 -2
  148. metaflow-stubs/plugins/kubernetes/kubernetes_decorator.pyi +2 -2
  149. metaflow-stubs/plugins/kubernetes/kubernetes_jobsets.pyi +2 -2
  150. metaflow-stubs/plugins/parallel_decorator.pyi +2 -2
  151. metaflow-stubs/plugins/perimeters.pyi +2 -2
  152. metaflow-stubs/plugins/project_decorator.pyi +2 -2
  153. metaflow-stubs/plugins/pypi/__init__.pyi +2 -2
  154. metaflow-stubs/plugins/pypi/conda_decorator.pyi +2 -2
  155. metaflow-stubs/plugins/pypi/conda_environment.pyi +5 -5
  156. metaflow-stubs/plugins/pypi/pypi_decorator.pyi +2 -2
  157. metaflow-stubs/plugins/pypi/pypi_environment.pyi +2 -2
  158. metaflow-stubs/plugins/pypi/utils.pyi +2 -2
  159. metaflow-stubs/plugins/resources_decorator.pyi +2 -2
  160. metaflow-stubs/plugins/retry_decorator.pyi +2 -2
  161. metaflow-stubs/plugins/secrets/__init__.pyi +3 -3
  162. metaflow-stubs/plugins/secrets/inline_secrets_provider.pyi +3 -3
  163. metaflow-stubs/plugins/secrets/secrets_decorator.pyi +2 -2
  164. metaflow-stubs/plugins/snowflake/__init__.pyi +2 -2
  165. metaflow-stubs/plugins/storage_executor.pyi +2 -2
  166. metaflow-stubs/plugins/test_unbounded_foreach_decorator.pyi +2 -2
  167. metaflow-stubs/plugins/timeout_decorator.pyi +3 -3
  168. metaflow-stubs/profilers/__init__.pyi +2 -2
  169. metaflow-stubs/pylint_wrapper.pyi +2 -2
  170. metaflow-stubs/runner/__init__.pyi +2 -2
  171. metaflow-stubs/runner/deployer.pyi +7 -7
  172. metaflow-stubs/runner/deployer_impl.pyi +2 -2
  173. metaflow-stubs/runner/metaflow_runner.pyi +2 -2
  174. metaflow-stubs/runner/nbdeploy.pyi +2 -2
  175. metaflow-stubs/runner/nbrun.pyi +2 -2
  176. metaflow-stubs/runner/subprocess_manager.pyi +2 -2
  177. metaflow-stubs/runner/utils.pyi +3 -3
  178. metaflow-stubs/system/__init__.pyi +2 -2
  179. metaflow-stubs/system/system_logger.pyi +3 -3
  180. metaflow-stubs/system/system_monitor.pyi +2 -2
  181. metaflow-stubs/tagging_util.pyi +2 -2
  182. metaflow-stubs/tuple_util.pyi +2 -2
  183. metaflow-stubs/user_configs/__init__.pyi +2 -2
  184. metaflow-stubs/user_configs/config_decorators.pyi +6 -6
  185. metaflow-stubs/user_configs/config_options.pyi +10 -5
  186. metaflow-stubs/user_configs/config_parameters.pyi +10 -10
  187. {ob_metaflow_stubs-6.0.3.128.dist-info → ob_metaflow_stubs-6.0.3.129.dist-info}/METADATA +1 -1
  188. ob_metaflow_stubs-6.0.3.129.dist-info/RECORD +191 -0
  189. ob_metaflow_stubs-6.0.3.128.dist-info/RECORD +0 -191
  190. {ob_metaflow_stubs-6.0.3.128.dist-info → ob_metaflow_stubs-6.0.3.129.dist-info}/WHEEL +0 -0
  191. {ob_metaflow_stubs-6.0.3.128.dist-info → ob_metaflow_stubs-6.0.3.129.dist-info}/top_level.txt +0 -0
@@ -1,15 +1,15 @@
1
1
  ######################################################################################################
2
2
  # Auto-generated Metaflow stub file #
3
- # MF version: 2.12.39.1+obcheckpoint(0.1.6);ob(v1) #
4
- # Generated on 2024-12-18T07:29:34.132398 #
3
+ # MF version: 2.13.0.1+obcheckpoint(0.1.6);ob(v1) #
4
+ # Generated on 2024-12-20T17:11:09.779274 #
5
5
  ######################################################################################################
6
6
 
7
7
  from __future__ import annotations
8
8
 
9
9
  import typing
10
10
  if typing.TYPE_CHECKING:
11
- import typing
12
11
  import datetime
12
+ import typing
13
13
  FlowSpecDerived = typing.TypeVar("FlowSpecDerived", bound="FlowSpec", contravariant=False, covariant=False)
14
14
  StepFlag = typing.NewType("StepFlag", bool)
15
15
 
@@ -31,11 +31,12 @@ from .parameters import Parameter as Parameter
31
31
  from .parameters import JSONTypeClass as JSONTypeClass
32
32
  from .parameters import JSONType as JSONType
33
33
  from .user_configs.config_parameters import Config as Config
34
+ from .user_configs.config_parameters import ConfigValue as ConfigValue
34
35
  from .user_configs.config_parameters import config_expr as config_expr
35
36
  from .user_configs.config_decorators import CustomFlowDecorator as CustomFlowDecorator
36
37
  from .user_configs.config_decorators import CustomStepDecorator as CustomStepDecorator
37
- from . import events as events
38
38
  from . import cards as cards
39
+ from . import events as events
39
40
  from . import tuple_util as tuple_util
40
41
  from . import runner as runner
41
42
  from . import plugins as plugins
@@ -145,40 +146,176 @@ def step(f: typing.Union[typing.Callable[[FlowSpecDerived], None], typing.Callab
145
146
  """
146
147
  ...
147
148
 
148
- 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] = [], 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]]]:
149
+ @typing.overload
150
+ def environment(*, vars: typing.Dict[str, str] = {}) -> typing.Callable[[typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]], typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]]]:
149
151
  """
150
- Specifies that this step should execute on Kubernetes.
152
+ Specifies environment variables to be set prior to the execution of a step.
151
153
  """
152
154
  ...
153
155
 
154
156
  @typing.overload
155
- 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]]]:
157
+ def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
158
+ ...
159
+
160
+ @typing.overload
161
+ def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
162
+ ...
163
+
164
+ def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
156
165
  """
157
- Specifies the Conda environment for the step.
166
+ Specifies environment variables to be set prior to the execution of a step.
167
+ """
168
+ ...
169
+
170
+ @typing.overload
171
+ 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]]]:
172
+ """
173
+ Specifies the number of times the task corresponding
174
+ to a step needs to be retried.
158
175
 
159
- Information in this decorator will augment any
160
- attributes set in the `@conda_base` flow-level decorator. Hence,
161
- you can use `@conda_base` to set packages required by all
162
- steps and use `@conda` to specify step-specific overrides.
176
+ This decorator is useful for handling transient errors, such as networking issues.
177
+ If your task contains operations that can't be retried safely, e.g. database updates,
178
+ it is advisable to annotate it with `@retry(times=0)`.
179
+
180
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
181
+ decorator will execute a no-op task after all retries have been exhausted,
182
+ ensuring that the flow execution can continue.
163
183
  """
164
184
  ...
165
185
 
166
186
  @typing.overload
167
- def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
187
+ def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
168
188
  ...
169
189
 
170
190
  @typing.overload
171
- def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
191
+ def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
172
192
  ...
173
193
 
174
- 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):
194
+ 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):
175
195
  """
176
- Specifies the Conda environment for the step.
196
+ Specifies the number of times the task corresponding
197
+ to a step needs to be retried.
177
198
 
178
- Information in this decorator will augment any
179
- attributes set in the `@conda_base` flow-level decorator. Hence,
180
- you can use `@conda_base` to set packages required by all
181
- steps and use `@conda` to specify step-specific overrides.
199
+ This decorator is useful for handling transient errors, such as networking issues.
200
+ If your task contains operations that can't be retried safely, e.g. database updates,
201
+ it is advisable to annotate it with `@retry(times=0)`.
202
+
203
+ This can be used in conjunction with the `@catch` decorator. The `@catch`
204
+ decorator will execute a no-op task after all retries have been exhausted,
205
+ ensuring that the flow execution can continue.
206
+ """
207
+ ...
208
+
209
+ @typing.overload
210
+ 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]]]:
211
+ """
212
+ Creates a human-readable report, a Metaflow Card, after this step completes.
213
+
214
+ Note that you may add multiple `@card` decorators in a step with different parameters.
215
+ """
216
+ ...
217
+
218
+ @typing.overload
219
+ def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
220
+ ...
221
+
222
+ @typing.overload
223
+ def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
224
+ ...
225
+
226
+ 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):
227
+ """
228
+ Creates a human-readable report, a Metaflow Card, after this step completes.
229
+
230
+ Note that you may add multiple `@card` decorators in a step with different parameters.
231
+ """
232
+ ...
233
+
234
+ 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]]]:
235
+ """
236
+ Decorator that helps cache, version and store models/datasets from huggingface hub.
237
+ """
238
+ ...
239
+
240
+ def nvidia(*, gpu: int, gpu_type: 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]]]:
241
+ """
242
+ Specifies that this step should execute on DGX cloud.
243
+ """
244
+ ...
245
+
246
+ @typing.overload
247
+ 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]]]:
248
+ """
249
+ Specifies that the step will success under all circumstances.
250
+
251
+ The decorator will create an optional artifact, specified by `var`, which
252
+ contains the exception raised. You can use it to detect the presence
253
+ of errors, indicating that all happy-path artifacts produced by the step
254
+ are missing.
255
+ """
256
+ ...
257
+
258
+ @typing.overload
259
+ def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
260
+ ...
261
+
262
+ @typing.overload
263
+ def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
264
+ ...
265
+
266
+ 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):
267
+ """
268
+ Specifies that the step will success under all circumstances.
269
+
270
+ The decorator will create an optional artifact, specified by `var`, which
271
+ contains the exception raised. You can use it to detect the presence
272
+ of errors, indicating that all happy-path artifacts produced by the step
273
+ are missing.
274
+ """
275
+ ...
276
+
277
+ @typing.overload
278
+ 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]]]:
279
+ """
280
+ Specifies a timeout for your step.
281
+
282
+ This decorator is useful if this step may hang indefinitely.
283
+
284
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
285
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
286
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
287
+
288
+ Note that all the values specified in parameters are added together so if you specify
289
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
290
+ """
291
+ ...
292
+
293
+ @typing.overload
294
+ def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
295
+ ...
296
+
297
+ @typing.overload
298
+ def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
299
+ ...
300
+
301
+ 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):
302
+ """
303
+ Specifies a timeout for your step.
304
+
305
+ This decorator is useful if this step may hang indefinitely.
306
+
307
+ This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
308
+ A timeout is considered to be an exception thrown by the step. It will cause the step to be
309
+ retried if needed and the exception will be caught by the `@catch` decorator, if present.
310
+
311
+ Note that all the values specified in parameters are added together so if you specify
312
+ 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
313
+ """
314
+ ...
315
+
316
+ 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] = [], 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]]]:
317
+ """
318
+ Specifies that this step should execute on Kubernetes.
182
319
  """
183
320
  ...
184
321
 
@@ -232,81 +369,52 @@ def resources(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None]
232
369
  ...
233
370
 
234
371
  @typing.overload
235
- 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]]]:
372
+ 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]]]:
236
373
  """
237
- Specifies the PyPI packages for the step.
374
+ Specifies the Conda environment for the step.
238
375
 
239
376
  Information in this decorator will augment any
240
- attributes set in the `@pyi_base` flow-level decorator. Hence,
241
- you can use `@pypi_base` to set packages required by all
242
- steps and use `@pypi` to specify step-specific overrides.
377
+ attributes set in the `@conda_base` flow-level decorator. Hence,
378
+ you can use `@conda_base` to set packages required by all
379
+ steps and use `@conda` to specify step-specific overrides.
243
380
  """
244
381
  ...
245
382
 
246
383
  @typing.overload
247
- def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
384
+ def conda(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
248
385
  ...
249
386
 
250
387
  @typing.overload
251
- def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
388
+ def conda(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
252
389
  ...
253
390
 
254
- 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):
391
+ 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):
255
392
  """
256
- Specifies the PyPI packages for the step.
393
+ Specifies the Conda environment for the step.
257
394
 
258
395
  Information in this decorator will augment any
259
- attributes set in the `@pyi_base` flow-level decorator. Hence,
260
- you can use `@pypi_base` to set packages required by all
261
- steps and use `@pypi` to specify step-specific overrides.
262
- """
263
- ...
264
-
265
- @typing.overload
266
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
267
- """
268
- Internal decorator to support Fast bakery
269
- """
270
- ...
271
-
272
- @typing.overload
273
- def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
274
- ...
275
-
276
- def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
277
- """
278
- Internal decorator to support Fast bakery
396
+ attributes set in the `@conda_base` flow-level decorator. Hence,
397
+ you can use `@conda_base` to set packages required by all
398
+ steps and use `@conda` to specify step-specific overrides.
279
399
  """
280
400
  ...
281
401
 
282
402
  @typing.overload
283
- 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]]]:
403
+ def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
284
404
  """
285
- Specifies that the step will success under all circumstances.
286
-
287
- The decorator will create an optional artifact, specified by `var`, which
288
- contains the exception raised. You can use it to detect the presence
289
- of errors, indicating that all happy-path artifacts produced by the step
290
- are missing.
405
+ Decorator prototype for all step decorators. This function gets specialized
406
+ and imported for all decorators types by _import_plugin_decorators().
291
407
  """
292
408
  ...
293
409
 
294
410
  @typing.overload
295
- def catch(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
296
- ...
297
-
298
- @typing.overload
299
- def catch(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
411
+ def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
300
412
  ...
301
413
 
302
- 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):
414
+ def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
303
415
  """
304
- Specifies that the step will success under all circumstances.
305
-
306
- The decorator will create an optional artifact, specified by `var`, which
307
- contains the exception raised. You can use it to detect the presence
308
- of errors, indicating that all happy-path artifacts produced by the step
309
- are missing.
416
+ Decorator prototype for all step decorators. This function gets specialized
417
+ and imported for all decorators types by _import_plugin_decorators().
310
418
  """
311
419
  ...
312
420
 
@@ -332,21 +440,23 @@ def model(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], ty
332
440
  ...
333
441
 
334
442
  @typing.overload
335
- def parallel(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
443
+ 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]]]:
336
444
  """
337
- Decorator prototype for all step decorators. This function gets specialized
338
- and imported for all decorators types by _import_plugin_decorators().
445
+ Enables checkpointing for a step.
339
446
  """
340
447
  ...
341
448
 
342
449
  @typing.overload
343
- def parallel(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
450
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
344
451
  ...
345
452
 
346
- def parallel(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
453
+ @typing.overload
454
+ def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
455
+ ...
456
+
457
+ 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):
347
458
  """
348
- Decorator prototype for all step decorators. This function gets specialized
349
- and imported for all decorators types by _import_plugin_decorators().
459
+ Enables checkpointing for a step.
350
460
  """
351
461
  ...
352
462
 
@@ -374,159 +484,181 @@ def secrets(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None],
374
484
  ...
375
485
 
376
486
  @typing.overload
377
- 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]]]:
487
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
378
488
  """
379
- Specifies the number of times the task corresponding
380
- to a step needs to be retried.
381
-
382
- This decorator is useful for handling transient errors, such as networking issues.
383
- If your task contains operations that can't be retried safely, e.g. database updates,
384
- it is advisable to annotate it with `@retry(times=0)`.
385
-
386
- This can be used in conjunction with the `@catch` decorator. The `@catch`
387
- decorator will execute a no-op task after all retries have been exhausted,
388
- ensuring that the flow execution can continue.
489
+ Internal decorator to support Fast bakery
389
490
  """
390
491
  ...
391
492
 
392
493
  @typing.overload
393
- def retry(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
394
- ...
395
-
396
- @typing.overload
397
- def retry(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
494
+ def fast_bakery_internal(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
398
495
  ...
399
496
 
400
- 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):
497
+ def fast_bakery_internal(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None):
401
498
  """
402
- Specifies the number of times the task corresponding
403
- to a step needs to be retried.
404
-
405
- This decorator is useful for handling transient errors, such as networking issues.
406
- If your task contains operations that can't be retried safely, e.g. database updates,
407
- it is advisable to annotate it with `@retry(times=0)`.
408
-
409
- This can be used in conjunction with the `@catch` decorator. The `@catch`
410
- decorator will execute a no-op task after all retries have been exhausted,
411
- ensuring that the flow execution can continue.
499
+ Internal decorator to support Fast bakery
412
500
  """
413
501
  ...
414
502
 
415
503
  @typing.overload
416
- 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]]]:
504
+ 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]]]:
417
505
  """
418
- Enables checkpointing for a step.
506
+ Specifies the PyPI packages for the step.
507
+
508
+ Information in this decorator will augment any
509
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
510
+ you can use `@pypi_base` to set packages required by all
511
+ steps and use `@pypi` to specify step-specific overrides.
419
512
  """
420
513
  ...
421
514
 
422
515
  @typing.overload
423
- def checkpoint(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
516
+ def pypi(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
424
517
  ...
425
518
 
426
519
  @typing.overload
427
- def checkpoint(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
428
- ...
429
-
430
- 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):
431
- """
432
- Enables checkpointing for a step.
433
- """
520
+ def pypi(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
434
521
  ...
435
522
 
436
- def nvidia(*, gpu: int, gpu_type: 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]]]:
523
+ def pypi(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
437
524
  """
438
- Specifies that this step should execute on DGX cloud.
525
+ Specifies the PyPI packages for the step.
526
+
527
+ Information in this decorator will augment any
528
+ attributes set in the `@pyi_base` flow-level decorator. Hence,
529
+ you can use `@pypi_base` to set packages required by all
530
+ steps and use `@pypi` to specify step-specific overrides.
439
531
  """
440
532
  ...
441
533
 
442
- @typing.overload
443
- 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]]]:
534
+ 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]]:
444
535
  """
445
- Specifies environment variables to be set prior to the execution of a step.
536
+ 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.
537
+ 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.
446
538
  """
447
539
  ...
448
540
 
449
- @typing.overload
450
- def environment(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
451
- ...
452
-
453
- @typing.overload
454
- def environment(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
455
- ...
456
-
457
- def environment(f: typing.Union[typing.Callable[[FlowSpecDerived, StepFlag], None], typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None], None] = None, *, vars: typing.Dict[str, str] = {}):
541
+ 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]]:
458
542
  """
459
- Specifies environment variables to be set prior to the execution of a step.
543
+ 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)
544
+ before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
545
+ and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
546
+ added as a flow decorators. Adding more than one decorator will ensure that `start` step
547
+ starts only after all sensors finish.
460
548
  """
461
549
  ...
462
550
 
463
- 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]]]:
551
+ def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
464
552
  """
465
- Decorator that helps cache, version and store models/datasets from huggingface hub.
553
+ Specifies what flows belong to the same project.
554
+
555
+ A project-specific namespace is created for all flows that
556
+ use the same `@project(name)`.
466
557
  """
467
558
  ...
468
559
 
469
- @typing.overload
470
- 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]]]:
560
+ def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
471
561
  """
472
- Specifies a timeout for your step.
562
+ This decorator is used to run NIM containers in Metaflow tasks as sidecars.
473
563
 
474
- This decorator is useful if this step may hang indefinitely.
564
+ User code call
565
+ -----------
566
+ @nim(
567
+ models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
568
+ backend='managed'
569
+ )
475
570
 
476
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
477
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
478
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
571
+ Valid backend options
572
+ ---------------------
573
+ - 'managed': Outerbounds selects a compute provider based on the model.
574
+ - 🚧 'dataplane': Run in your account.
479
575
 
480
- Note that all the values specified in parameters are added together so if you specify
481
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
576
+ Valid model options
577
+ ----------------
578
+ - 'meta/llama3-8b-instruct': 8B parameter model
579
+ - 'meta/llama3-70b-instruct': 70B parameter model
580
+ - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
482
581
  """
483
582
  ...
484
583
 
485
584
  @typing.overload
486
- def timeout(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
487
- ...
488
-
489
- @typing.overload
490
- def timeout(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
491
- ...
492
-
493
- 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):
585
+ def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
494
586
  """
495
- Specifies a timeout for your step.
587
+ Specifies the flow(s) that this flow depends on.
496
588
 
497
- This decorator is useful if this step may hang indefinitely.
589
+ ```
590
+ @trigger_on_finish(flow='FooFlow')
591
+ ```
592
+ or
593
+ ```
594
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
595
+ ```
596
+ This decorator respects the @project decorator and triggers the flow
597
+ when upstream runs within the same namespace complete successfully
498
598
 
499
- This can be used in conjunction with the `@retry` decorator as well as the `@catch` decorator.
500
- A timeout is considered to be an exception thrown by the step. It will cause the step to be
501
- retried if needed and the exception will be caught by the `@catch` decorator, if present.
599
+ Additionally, you can specify project aware upstream flow dependencies
600
+ by specifying the fully qualified project_flow_name.
601
+ ```
602
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
603
+ ```
604
+ or
605
+ ```
606
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
607
+ ```
502
608
 
503
- Note that all the values specified in parameters are added together so if you specify
504
- 60 seconds and 1 hour, the decorator will have an effective timeout of 1 hour and 1 minute.
505
- """
506
- ...
507
-
508
- @typing.overload
509
- 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]]]:
510
- """
511
- Creates a human-readable report, a Metaflow Card, after this step completes.
609
+ You can also specify just the project or project branch (other values will be
610
+ inferred from the current project or project branch):
611
+ ```
612
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
613
+ ```
512
614
 
513
- Note that you may add multiple `@card` decorators in a step with different parameters.
615
+ Note that `branch` is typically one of:
616
+ - `prod`
617
+ - `user.bob`
618
+ - `test.my_experiment`
619
+ - `prod.staging`
514
620
  """
515
621
  ...
516
622
 
517
623
  @typing.overload
518
- def card(f: typing.Callable[[FlowSpecDerived, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, StepFlag], None]:
519
- ...
520
-
521
- @typing.overload
522
- def card(f: typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]) -> typing.Callable[[FlowSpecDerived, typing.Any, StepFlag], None]:
624
+ def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
523
625
  ...
524
626
 
525
- 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
+ def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
526
628
  """
527
- Creates a human-readable report, a Metaflow Card, after this step completes.
629
+ Specifies the flow(s) that this flow depends on.
528
630
 
529
- Note that you may add multiple `@card` decorators in a step with different parameters.
631
+ ```
632
+ @trigger_on_finish(flow='FooFlow')
633
+ ```
634
+ or
635
+ ```
636
+ @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
637
+ ```
638
+ This decorator respects the @project decorator and triggers the flow
639
+ when upstream runs within the same namespace complete successfully
640
+
641
+ Additionally, you can specify project aware upstream flow dependencies
642
+ by specifying the fully qualified project_flow_name.
643
+ ```
644
+ @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
645
+ ```
646
+ or
647
+ ```
648
+ @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
649
+ ```
650
+
651
+ You can also specify just the project or project branch (other values will be
652
+ inferred from the current project or project branch):
653
+ ```
654
+ @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
655
+ ```
656
+
657
+ Note that `branch` is typically one of:
658
+ - `prod`
659
+ - `user.bob`
660
+ - `test.my_experiment`
661
+ - `prod.staging`
530
662
  """
531
663
  ...
532
664
 
@@ -572,13 +704,26 @@ def conda_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packa
572
704
  """
573
705
  ...
574
706
 
575
- 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]]:
707
+ @typing.overload
708
+ def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
576
709
  """
577
- 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)
578
- before the start step of the flow. This decorator only works when a flow is scheduled on Airflow
579
- and is compiled using `airflow create`. More than one `@airflow_s3_key_sensor` can be
580
- added as a flow decorators. Adding more than one decorator will ensure that `start` step
581
- starts only after all sensors finish.
710
+ Specifies the PyPI packages for all steps of the flow.
711
+
712
+ Use `@pypi_base` to set common packages required by all
713
+ steps and use `@pypi` to specify step-specific overrides.
714
+ """
715
+ ...
716
+
717
+ @typing.overload
718
+ def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
719
+ ...
720
+
721
+ def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
722
+ """
723
+ Specifies the PyPI packages for all steps of the flow.
724
+
725
+ Use `@pypi_base` to set common packages required by all
726
+ steps and use `@pypi` to specify step-specific overrides.
582
727
  """
583
728
  ...
584
729
 
@@ -655,149 +800,5 @@ def trigger(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, event: t
655
800
  """
656
801
  ...
657
802
 
658
- @typing.overload
659
- def pypi_base(*, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
660
- """
661
- Specifies the PyPI packages for all steps of the flow.
662
-
663
- Use `@pypi_base` to set common packages required by all
664
- steps and use `@pypi` to specify step-specific overrides.
665
- """
666
- ...
667
-
668
- @typing.overload
669
- def pypi_base(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
670
- ...
671
-
672
- def pypi_base(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, packages: typing.Dict[str, str] = {}, python: typing.Optional[str] = None):
673
- """
674
- Specifies the PyPI packages for all steps of the flow.
675
-
676
- Use `@pypi_base` to set common packages required by all
677
- steps and use `@pypi` to specify step-specific overrides.
678
- """
679
- ...
680
-
681
- 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]]:
682
- """
683
- 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.
684
- 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.
685
- """
686
- ...
687
-
688
- def project(*, name: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
689
- """
690
- Specifies what flows belong to the same project.
691
-
692
- A project-specific namespace is created for all flows that
693
- use the same `@project(name)`.
694
- """
695
- ...
696
-
697
- def nim(*, models: "list[NIM]", backend: str) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
698
- """
699
- This decorator is used to run NIM containers in Metaflow tasks as sidecars.
700
-
701
- User code call
702
- -----------
703
- @nim(
704
- models=['meta/llama3-8b-instruct', 'meta/llama3-70b-instruct'],
705
- backend='managed'
706
- )
707
-
708
- Valid backend options
709
- ---------------------
710
- - 'managed': Outerbounds selects a compute provider based on the model.
711
- - 🚧 'dataplane': Run in your account.
712
-
713
- Valid model options
714
- ----------------
715
- - 'meta/llama3-8b-instruct': 8B parameter model
716
- - 'meta/llama3-70b-instruct': 70B parameter model
717
- - Upon request, any model here: https://nvcf.ngc.nvidia.com/functions?filter=nvidia-functions
718
- """
719
- ...
720
-
721
- @typing.overload
722
- def trigger_on_finish(*, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}) -> typing.Callable[[typing.Type[FlowSpecDerived]], typing.Type[FlowSpecDerived]]:
723
- """
724
- Specifies the flow(s) that this flow depends on.
725
-
726
- ```
727
- @trigger_on_finish(flow='FooFlow')
728
- ```
729
- or
730
- ```
731
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
732
- ```
733
- This decorator respects the @project decorator and triggers the flow
734
- when upstream runs within the same namespace complete successfully
735
-
736
- Additionally, you can specify project aware upstream flow dependencies
737
- by specifying the fully qualified project_flow_name.
738
- ```
739
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
740
- ```
741
- or
742
- ```
743
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
744
- ```
745
-
746
- You can also specify just the project or project branch (other values will be
747
- inferred from the current project or project branch):
748
- ```
749
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
750
- ```
751
-
752
- Note that `branch` is typically one of:
753
- - `prod`
754
- - `user.bob`
755
- - `test.my_experiment`
756
- - `prod.staging`
757
- """
758
- ...
759
-
760
- @typing.overload
761
- def trigger_on_finish(f: typing.Type[FlowSpecDerived]) -> typing.Type[FlowSpecDerived]:
762
- ...
763
-
764
- def trigger_on_finish(f: typing.Optional[typing.Type[FlowSpecDerived]] = None, *, flow: typing.Union[typing.Dict[str, str], str, None] = None, flows: typing.List[typing.Union[str, typing.Dict[str, str]]] = [], options: typing.Dict[str, typing.Any] = {}):
765
- """
766
- Specifies the flow(s) that this flow depends on.
767
-
768
- ```
769
- @trigger_on_finish(flow='FooFlow')
770
- ```
771
- or
772
- ```
773
- @trigger_on_finish(flows=['FooFlow', 'BarFlow'])
774
- ```
775
- This decorator respects the @project decorator and triggers the flow
776
- when upstream runs within the same namespace complete successfully
777
-
778
- Additionally, you can specify project aware upstream flow dependencies
779
- by specifying the fully qualified project_flow_name.
780
- ```
781
- @trigger_on_finish(flow='my_project.branch.my_branch.FooFlow')
782
- ```
783
- or
784
- ```
785
- @trigger_on_finish(flows=['my_project.branch.my_branch.FooFlow', 'BarFlow'])
786
- ```
787
-
788
- You can also specify just the project or project branch (other values will be
789
- inferred from the current project or project branch):
790
- ```
791
- @trigger_on_finish(flow={"name": "FooFlow", "project": "my_project", "project_branch": "branch"})
792
- ```
793
-
794
- Note that `branch` is typically one of:
795
- - `prod`
796
- - `user.bob`
797
- - `test.my_experiment`
798
- - `prod.staging`
799
- """
800
- ...
801
-
802
803
  pkg_name: str
803
804